Push final-ish state

Author: dan-pandori

src/imitation/policies/replay_buffer_wrapper.py

@@ -137,11 +137,6 @@ def __init__(
         """
         # TODO should we limit by number of batches (as this does)
         #      or number of observations returned?
-        self.sample_count = 0
-        self.entropy_as_reward_samples = entropy_as_reward_samples
-        self.k = k
-        # TODO support n_envs > 1
-        self.entropy_stats = util.RunningMeanAndVar(shape=(1,))
         super().__init__(
             buffer_size,
             observation_space,
@@ -150,14 +145,18 @@ def __init__(
             reward_fn=reward_fn,
             **kwargs,
         )
+        self.sample_count = 0
+        self.k = k
+        # TODO support n_envs > 1
+        self.entropy_stats = util.RunningMeanAndVar(shape=(1,))
+        self.entropy_as_reward_samples = entropy_as_reward_samples
 
-    # TODO this seems to never actually get called?
     def sample(self, *args, **kwargs):
         self.sample_count += 1
         samples = super().sample(*args, **kwargs)
-        print(self.sample_count)
-        print(self.entropy_as_reward_samples)
-        if self.sample_count > 500:
+        # For some reason self.entropy_as_reward_samples seems to get cleared,
+        # and I have no idea why.
+        if self.sample_count > self.entropy_as_reward_samples:
             return samples
         # TODO we really ought to reset the reward network once we are done w/
         #      the entropy based pre-training. We also have no reason to train

tests/policies/test_replay_buffer_wrapper.py

@@ -190,7 +190,7 @@ def reset(self):
 
 def test_entropy_wrapper_class(tmpdir, rng):
     buffer_size = 20
-    entropy_samples = 40
+    entropy_samples = 500
     k = 4
 
     venv = DummyVecEnv([ActionIsObsEnv])
@@ -211,17 +211,17 @@ def test_entropy_wrapper_class(tmpdir, rng):
 
     rl_algo.learn(total_timesteps=buffer_size)
     initial_entropy = util.compute_state_entropy(
-        th.Tensor(rl_algo.replay_buffer.observations),
-        th.Tensor(rl_algo.replay_buffer.observations),
+        th.Tensor(rl_algo.replay_buffer.replay_buffer.observations),
+        th.Tensor(rl_algo.replay_buffer.replay_buffer.observations),
         k=k,
     )
 
     rl_algo.learn(total_timesteps=entropy_samples - buffer_size)
     # Expect that the entropy of our replay buffer is now higher,
     # since we trained with that as the reward.
     trained_entropy = util.compute_state_entropy(
-        th.Tensor(rl_algo.replay_buffer.observations),
-        th.Tensor(rl_algo.replay_buffer.observations),
+        th.Tensor(rl_algo.replay_buffer.replay_buffer.observations),
+        th.Tensor(rl_algo.replay_buffer.replay_buffer.observations),
         k=k,
     )
     assert trained_entropy.mean() > initial_entropy.mean()