Commit for diff

Author: dan-pandori

src/imitation/policies/replay_buffer_wrapper.py

@@ -135,6 +135,13 @@ def __init__(
             k: Use the k'th nearest neighbor's distance when computing state entropy.
             **kwargs: keyword arguments for ReplayBuffer.
         """
+        # 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,
@@ -143,18 +150,14 @@ def __init__(
             reward_fn=reward_fn,
             **kwargs,
         )
-        # TODO should we limit by number of batches (as this does)
-        #      or number of observations returned?
-        self.samples = 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,))
 
+    # TODO this seems to never actually get called?
     def sample(self, *args, **kwargs):
-        self.samples += 1
+        self.sample_count += 1
         samples = super().sample(*args, **kwargs)
-        if self.samples > self.entropy_as_reward_samples:
+        print(self.sample_count)
+        print(self.entropy_as_reward_samples)
+        if self.sample_count > 500:
             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
@@ -164,7 +167,12 @@ def sample(self, *args, **kwargs):
             all_obs = self.observations
         else:
             all_obs = self.observations[: self.pos]
-        entropies = util.compute_state_entropy(samples.observations, all_obs, self.k)
+        entropies = util.compute_state_entropy(
+            # TODO support multiple environments
+            samples.observations.unsqueeze(1),
+            all_obs,
+            self.k,
+        )
 
         # Normalize to have mean of 0 and standard deviation of 1
         self.entropy_stats.update(entropies)

tests/policies/test_replay_buffer_wrapper.py

@@ -12,6 +12,7 @@
 from stable_baselines3.common import buffers, off_policy_algorithm, policies
 from stable_baselines3.common.policies import BasePolicy
 from stable_baselines3.common.save_util import load_from_pkl
+from stable_baselines3.common.vec_env import DummyVecEnv
 
 from imitation.policies.replay_buffer_wrapper import (
     ReplayBufferEntropyRewardWrapper,
@@ -123,6 +124,7 @@ def test_wrapper_class(tmpdir, rng):
 def test_entropy_wrapper_class_no_op(tmpdir, rng):
     buffer_size = 15
     total_timesteps = 20
+    entropy_samples = 0
 
     venv = util.make_vec_env("Pendulum-v1", n_envs=1, rng=rng)
     rl_algo = sb3.SAC(
@@ -134,7 +136,7 @@ def test_entropy_wrapper_class_no_op(tmpdir, rng):
         replay_buffer_kwargs=dict(
             replay_buffer_class=buffers.ReplayBuffer,
             reward_fn=zero_reward_fn,
-            entropy_as_reward_samples=0,
+            entropy_as_reward_samples=entropy_samples,
         ),
         buffer_size=buffer_size,
     )
@@ -172,8 +174,8 @@ class ActionIsObsEnv(gym.Env):
     def __init__(self):
         """Initialize environment."""
         super().__init__()
-        self.action_space = spaces.Discrete(50)
-        self.observation_space = spaces.Discrete(50)
+        self.action_space = spaces.Box(np.array([0]), np.array([1]))
+        self.observation_space = spaces.Box(np.array([0]), np.array([1]))
 
     def step(self, action):
         obs = action
@@ -183,17 +185,15 @@ def step(self, action):
         return obs, reward, done, info
 
     def reset(self):
-        return self.action_space.sample()
+        return np.array([0])
 
 
 def test_entropy_wrapper_class(tmpdir, rng):
-    buffer_size = 15
-    entropy_samples = 10
-    total_timesteps = 20
-
-    # TODO expect that our behavior is approximately uniformly distributed
+    buffer_size = 20
+    entropy_samples = 40
+    k = 4
 
-    venv = util.make_vec_env("Pendulum-v1", n_envs=1, rng=rng)
+    venv = DummyVecEnv([ActionIsObsEnv])
     rl_algo = sb3.SAC(
         policy=sb3.sac.policies.SACPolicy,
         policy_kwargs=dict(),
@@ -204,32 +204,24 @@ def test_entropy_wrapper_class(tmpdir, rng):
             replay_buffer_class=buffers.ReplayBuffer,
             reward_fn=zero_reward_fn,
             entropy_as_reward_samples=entropy_samples,
+            k=k,
         ),
         buffer_size=buffer_size,
     )
 
-    rl_algo.learn(total_timesteps=total_timesteps)
-
-    buffer_path = osp.join(tmpdir, "buffer.pkl")
-    rl_algo.save_replay_buffer(buffer_path)
-    replay_buffer_wrapper = load_from_pkl(buffer_path)
-    replay_buffer = replay_buffer_wrapper.replay_buffer
-
-    # replay_buffer_wrapper.sample(...) should return zero-reward transitions
-    assert buffer_size == replay_buffer_wrapper.size() == replay_buffer.size()
-    assert (replay_buffer_wrapper.sample(total_timesteps).rewards == 0.0).all()
-    assert (replay_buffer.sample(total_timesteps).rewards != 0.0).all()  # seed=42
-
-    # replay_buffer_wrapper.pos, replay_buffer_wrapper.full
-    assert replay_buffer_wrapper.pos == total_timesteps - buffer_size
-    assert replay_buffer_wrapper.full
-
-    # reset()
-    replay_buffer_wrapper.reset()
-    assert 0 == replay_buffer_wrapper.size() == replay_buffer.size()
-    assert replay_buffer_wrapper.pos == 0
-    assert not replay_buffer_wrapper.full
+    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),
+        k=k,
+    )
 
-    # to_torch()
-    tensor = replay_buffer_wrapper.to_torch(np.ones(42))
-    assert type(tensor) is th.Tensor
+    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),
+        k=k,
+    )
+    assert trained_entropy.mean() > initial_entropy.mean()