#641 code review: refactor PebbleStateEntropyReward so that inner RewardNets can be injected from the outside

Author: Jan Michelfeit

src/imitation/algorithms/pebble/entropy_reward.py

@@ -13,47 +13,46 @@
     ReplayBufferView,
 )
 from imitation.rewards.reward_function import RewardFn
-from imitation.rewards.reward_nets import NormalizedRewardNet, RewardNet
+from imitation.rewards.reward_nets import RewardNet
 from imitation.util import util
-from imitation.util.networks import RunningNorm
-
-
-class PebbleRewardPhase(enum.Enum):
-    """States representing different behaviors for PebbleStateEntropyReward."""
-
-    UNSUPERVISED_EXPLORATION = enum.auto()  # Entropy based reward
-    POLICY_AND_REWARD_LEARNING = enum.auto()  # Learned reward
 
 
 class InsufficientObservations(RuntimeError):
     pass
 
 
-class EntropyRewardNet(RewardNet):
+class EntropyRewardNet(RewardNet, ReplayBufferAwareRewardFn):
     def __init__(
         self,
         nearest_neighbor_k: int,
-        replay_buffer_view: ReplayBufferView,
         observation_space: gym.Space,
         action_space: gym.Space,
         normalize_images: bool = True,
+        replay_buffer_view: Optional[ReplayBufferView] = None,
     ):
         """Initialize the RewardNet.
 
         Args:
+            nearest_neighbor_k: Parameter for entropy computation (see
+                compute_state_entropy())
             observation_space: the observation space of the environment
             action_space: the action space of the environment
             normalize_images: whether to automatically normalize
                 image observations to [0, 1] (from 0 to 255). Defaults to True.
+            replay_buffer_view: Replay buffer view with observations to compare
+                against when computing entropy. If None is given, the buffer needs to
+                be set with on_replay_buffer_initialized() before EntropyRewardNet can
+                be used
         """
         super().__init__(observation_space, action_space, normalize_images)
         self.nearest_neighbor_k = nearest_neighbor_k
         self._replay_buffer_view = replay_buffer_view
 
-    def set_replay_buffer(self, replay_buffer: ReplayBufferRewardWrapper):
-        """This method needs to be called after unpickling.
+    def on_replay_buffer_initialized(self, replay_buffer: ReplayBufferRewardWrapper):
+        """Sets replay buffer.
 
-        See also __getstate__() / __setstate__()
+        This method needs to be called, e.g., after unpickling.
+        See also __getstate__() / __setstate__().
         """
         assert self.observation_space == replay_buffer.observation_space
         assert self.action_space == replay_buffer.action_space
@@ -111,6 +110,13 @@ def __setstate__(self, state):
         self._replay_buffer_view = None
 
 
+class PebbleRewardPhase(enum.Enum):
+    """States representing different behaviors for PebbleStateEntropyReward."""
+
+    UNSUPERVISED_EXPLORATION = enum.auto()  # Entropy based reward
+    POLICY_AND_REWARD_LEARNING = enum.auto()  # Learned reward
+
+
 class PebbleStateEntropyReward(ReplayBufferAwareRewardFn):
     """Reward function for implementation of the PEBBLE learning algorithm.
 
@@ -126,48 +132,30 @@ class PebbleStateEntropyReward(ReplayBufferAwareRewardFn):
     reward is returned.
 
     The second phase requires that a buffer with observations to compare against is
-    supplied with set_replay_buffer() or on_replay_buffer_initialized().
-    To transition to the last phase, unsupervised_exploration_finish() needs
-    to be called.
+    supplied with on_replay_buffer_initialized(). To transition to the last phase,
+    unsupervised_exploration_finish() needs to be called.
     """
 
     def __init__(
         self,
+        entropy_reward_fn: RewardFn,
         learned_reward_fn: RewardFn,
-        nearest_neighbor_k: int = 5,
     ):
         """Builds this class.
 
         Args:
+            entropy_reward_fn: The entropy-based reward function used during
+                unsupervised exploration
             learned_reward_fn: The learned reward function used after unsupervised
                 exploration is finished
-            nearest_neighbor_k: Parameter for entropy computation (see
-                compute_state_entropy())
         """
+        self.entropy_reward_fn = entropy_reward_fn
         self.learned_reward_fn = learned_reward_fn
-        self.nearest_neighbor_k = nearest_neighbor_k
-
         self.state = PebbleRewardPhase.UNSUPERVISED_EXPLORATION
 
-        # These two need to be set with set_replay_buffer():
-        self._entropy_reward_net: Optional[EntropyRewardNet] = None
-        self._normalized_entropy_reward_net: Optional[RewardNet] = None
-
     def on_replay_buffer_initialized(self, replay_buffer: ReplayBufferRewardWrapper):
-        if self._normalized_entropy_reward_net is None:
-            self._entropy_reward_net = EntropyRewardNet(
-                nearest_neighbor_k=self.nearest_neighbor_k,
-                replay_buffer_view=replay_buffer.buffer_view,
-                observation_space=replay_buffer.observation_space,
-                action_space=replay_buffer.action_space,
-                normalize_images=False,
-            )
-            self._normalized_entropy_reward_net = NormalizedRewardNet(
-                self._entropy_reward_net, RunningNorm
-            )
-        else:
-            assert self._entropy_reward_net is not None
-            self._entropy_reward_net.set_replay_buffer(replay_buffer)
+        if isinstance(self.entropy_reward_fn, ReplayBufferAwareRewardFn):
+            self.entropy_reward_fn.on_replay_buffer_initialized(replay_buffer)
 
     def unsupervised_exploration_finish(self):
         assert self.state == PebbleRewardPhase.UNSUPERVISED_EXPLORATION
@@ -181,20 +169,11 @@ def __call__(
         done: np.ndarray,
     ) -> np.ndarray:
         if self.state == PebbleRewardPhase.UNSUPERVISED_EXPLORATION:
-            return self._entropy_reward(state, action, next_state, done)
+            try:
+                return self.entropy_reward_fn(state, action, next_state, done)
+            except InsufficientObservations:
+                # not enough observations to compare to, fall back to the learned function;
+                # (falling back to a constant may also be ok)
+                return self.learned_reward_fn(state, action, next_state, done)
         else:
             return self.learned_reward_fn(state, action, next_state, done)
-
-    def _entropy_reward(self, state, action, next_state, done):
-        if self._normalized_entropy_reward_net is None:
-            raise ValueError(
-                "Replay buffer must be supplied before entropy reward can be used",
-            )
-        try:
-            return self._normalized_entropy_reward_net.predict_processed(
-                state, action, next_state, done, update_stats=True
-            )
-        except InsufficientObservations:
-            # not enough observations to compare to, fall back to the learned function;
-            # (falling back to a constant may also be ok)
-            return self.learned_reward_fn(state, action, next_state, done)

src/imitation/scripts/train_preference_comparisons.py

@@ -13,17 +13,26 @@
 from stable_baselines3.common import base_class, type_aliases, vec_env
 
 from imitation.algorithms import preference_comparisons
-from imitation.algorithms.pebble.entropy_reward import PebbleStateEntropyReward
+from imitation.algorithms.pebble.entropy_reward import (
+    EntropyRewardNet,
+    PebbleStateEntropyReward,
+)
 from imitation.data import types
 from imitation.policies import serialize
+from imitation.policies.replay_buffer_wrapper import (
+    ReplayBufferAwareRewardFn,
+    ReplayBufferRewardWrapper,
+)
 from imitation.rewards import reward_function, reward_nets
+from imitation.rewards.reward_nets import NormalizedRewardNet
 from imitation.scripts.common import common, reward
 from imitation.scripts.common import rl as rl_common
 from imitation.scripts.common import train
 from imitation.scripts.config.train_preference_comparisons import (
     train_preference_comparisons_ex,
 )
 from imitation.util import logger as imit_logger
+from imitation.util.networks import RunningNorm
 
 
 def save_model(
@@ -71,14 +80,47 @@ def make_reward_function(
         reward_net.predict_processed,
         update_stats=False,
     )
+    observation_space = reward_net.observation_space
+    action_space = reward_net.action_space
     if pebble_enabled:
-        relabel_reward_fn = PebbleStateEntropyReward(
-            relabel_reward_fn,  # type: ignore[assignment]
+        relabel_reward_fn = create_pebble_reward_fn(
+            relabel_reward_fn,
             pebble_nearest_neighbor_k,
+            action_space,
+            observation_space,
         )
     return relabel_reward_fn
 
 
+def create_pebble_reward_fn(
+    relabel_reward_fn, pebble_nearest_neighbor_k, action_space, observation_space
+):
+    entropy_reward_net = EntropyRewardNet(
+        nearest_neighbor_k=pebble_nearest_neighbor_k,
+        observation_space=observation_space,
+        action_space=action_space,
+        normalize_images=False,
+    )
+    normalized_entropy_reward_net = NormalizedRewardNet(entropy_reward_net, RunningNorm)
+
+    class EntropyRewardFn(ReplayBufferAwareRewardFn):
+        """Adapter for entropy reward adding on_replay_buffer_initialized() hook."""
+
+        def __call__(self, *args, **kwargs) -> np.ndarray:
+            kwargs["update_stats"] = True
+            return normalized_entropy_reward_net.predict_processed(*args, **kwargs)
+
+        def on_replay_buffer_initialized(
+            self, replay_buffer: ReplayBufferRewardWrapper
+        ):
+            entropy_reward_net.on_replay_buffer_initialized(replay_buffer)
+
+    return PebbleStateEntropyReward(
+        EntropyRewardFn(),
+        relabel_reward_fn,  # type: ignore[assignment]
+    )
+
+
 @train_preference_comparisons_ex.capture
 def make_agent_trajectory_generator(
     venv: vec_env.VecEnv,

src/imitation/util/util.py

@@ -395,4 +395,3 @@ def compute_state_entropy(
         all_dists = th.cat(dists, dim=1)
         knn_dists = th.kthvalue(all_dists, k=k + 1, dim=1).values
         return knn_dists
-