#625 make entropy reward serializable with pickle

Author: Jan Michelfeit

src/imitation/algorithms/pebble/entropy_reward.py

@@ -14,13 +14,14 @@ def __init__(self, nearest_neighbor_k: int, observation_space: spaces.Space):
         self.nearest_neighbor_k = nearest_neighbor_k
         # TODO support n_envs > 1
         self.entropy_stats = RunningNorm(1)
+        self.observation_space = observation_space
         self.obs_shape = get_obs_shape(observation_space)
         self.replay_buffer_view = ReplayBufferView(
             np.empty(0, dtype=observation_space.dtype), lambda: slice(0)
         )
 
-    def set_buffer_view(self, replay_buffer_view: ReplayBufferView):
-        self.replay_buffer_view = replay_buffer_view
+    def set_replay_buffer(self, replay_buffer: ReplayBufferView):
+        self.replay_buffer_view = replay_buffer
 
     def __call__(
         self,
@@ -42,3 +43,14 @@ def __call__(
         )
         normalized_entropies = self.entropy_stats.forward(th.as_tensor(entropies))
         return normalized_entropies.numpy()
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        del state["replay_buffer_view"]
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+        self.replay_buffer_view = ReplayBufferView(
+            np.empty(0, self.observation_space.dtype), lambda: slice(0)
+        )

src/imitation/policies/replay_buffer_wrapper.py

@@ -10,6 +10,7 @@
 from imitation.rewards.reward_function import RewardFn
 from imitation.util import util
 from imitation.util.networks import RunningNorm
+from typing import Callable
 
 
 def _samples_to_reward_fn_input(

tests/algorithms/pebble/test_entropy_reward.py

@@ -1,3 +1,4 @@
+import pickle
 from unittest.mock import patch
 
 import numpy as np
@@ -33,7 +34,9 @@ def test_state_entropy_reward_returns_entropy(rng):
     expected = util.compute_state_entropy(
         observations, all_observations.reshape(-1, *obs_shape), K
     )
-    expected_normalized = reward_fn.entropy_stats.normalize(th.as_tensor(expected)).numpy()
+    expected_normalized = reward_fn.entropy_stats.normalize(
+        th.as_tensor(expected)
+    ).numpy()
     np.testing.assert_allclose(reward, expected_normalized)
 
 
@@ -44,7 +47,7 @@ def test_state_entropy_reward_returns_normalized_values():
 
         reward_fn = StateEntropyReward(K, SPACE)
         all_observations = np.empty((BUFFER_SIZE, VENVS, *get_obs_shape(SPACE)))
-        reward_fn.set_buffer_view(
+        reward_fn.set_replay_buffer(
             ReplayBufferView(all_observations, lambda: slice(None))
         )
 
@@ -68,3 +71,24 @@ def test_state_entropy_reward_returns_normalized_values():
             rtol=0.05,
             atol=0.05,
         )
+
+
+def test_state_entropy_reward_can_pickle():
+    all_observations = np.empty((BUFFER_SIZE, VENVS, *get_obs_shape(SPACE)))
+    replay_buffer = ReplayBufferView(all_observations, lambda: slice(None))
+
+    obs1 = np.random.rand(VENVS, *get_obs_shape(SPACE))
+    reward_fn = StateEntropyReward(K, SPACE)
+    reward_fn.set_replay_buffer(replay_buffer)
+    reward_fn(obs1, PLACEHOLDER, PLACEHOLDER, PLACEHOLDER)
+
+    # Act
+    pickled = pickle.dumps(reward_fn)
+    reward_fn_deserialized = pickle.loads(pickled)
+    reward_fn_deserialized.set_replay_buffer(replay_buffer)
+
+    # Assert
+    obs2 = np.random.rand(VENVS, *get_obs_shape(SPACE))
+    expected_result = reward_fn(obs2, PLACEHOLDER, PLACEHOLDER, PLACEHOLDER)
+    actual_result = reward_fn_deserialized(obs2, PLACEHOLDER, PLACEHOLDER, PLACEHOLDER)
+    np.testing.assert_allclose(actual_result, expected_result)