#641 increase coverage

Author: Jan Michelfeit

src/imitation/policies/replay_buffer_wrapper.py

@@ -143,5 +143,12 @@ class ReplayBufferAwareRewardFn(RewardFn, abc.ABC):
     def on_replay_buffer_initialized(
         self,
         replay_buffer: ReplayBufferRewardWrapper,
-    ):
-        pass
+    ) -> None:
+        """Hook method to be called when ReplayBuffer is initialized.
+
+        Needed to propagate the ReplayBuffer to a reward function because the buffer
+        is created indirectly in ReplayBufferRewardWrapper.
+
+        Args:
+            replay_buffer: the created ReplayBuffer
+        """  # noqa: DAR202

tests/algorithms/pebble/test_entropy_reward.py

@@ -155,10 +155,6 @@ def test_entropy_reward_net_can_pickle(rng):
     np.testing.assert_allclose(actual_result, expected_result)
 
 
-def reward_fn_stub(state, action, next_state, done):
-    return state
-
-
 def replay_buffer_mock(all_observations: np.ndarray, obs_space: Space = SPACE) -> Mock:
     buffer_view = ReplayBufferView(all_observations, lambda: slice(None))
     mock = Mock()

tests/algorithms/test_preference_comparisons.py

@@ -18,11 +18,16 @@
 
 import imitation.testing.reward_nets as testing_reward_nets
 from imitation.algorithms import preference_comparisons
+from imitation.algorithms.preference_comparisons import (
+    PebbleAgentTrainer,
+    TrajectoryGenerator,
+)
 from imitation.data import types
 from imitation.data.types import TrajectoryWithRew
 from imitation.policies.replay_buffer_wrapper import ReplayBufferView
 from imitation.regularization import regularizers, updaters
 from imitation.rewards import reward_nets
+from imitation.rewards.reward_function import RewardFn
 from imitation.scripts.train_preference_comparisons import create_pebble_reward_fn
 from imitation.util import networks, util
 
@@ -1120,3 +1125,26 @@ def test_that_trainer_improves(
     )
 
     assert np.mean(trained_agent_rewards) > np.mean(novice_agent_rewards)
+
+
+def test_trajectory_generator_raises_on_pretrain_if_not_implemented():
+    class TrajectoryGeneratorTestImpl(TrajectoryGenerator):
+        def sample(self, steps: int) -> Sequence[TrajectoryWithRew]:
+            return []
+
+    generator = TrajectoryGeneratorTestImpl()
+    assert generator.has_pretraining is False
+    with pytest.raises(ValueError, match="should not consume any timesteps"):
+        generator.unsupervised_pretrain(1)
+
+
+def test_pebble_agent_trainer_expects_pebble_reward(agent, venv, rng):
+    reward_fn: RewardFn = lambda state, action, next, done: state
+
+    with pytest.raises(ValueError, match="PebbleStateEntropyReward"):
+        PebbleAgentTrainer(
+            algorithm=agent,
+            reward_fn=reward_fn,  # type:ignore[arg-type]
+            venv=venv,
+            rng=rng,
+        )

tests/policies/test_replay_buffer_wrapper.py

@@ -4,7 +4,6 @@
 from typing import Type
 from unittest.mock import Mock
 
-import gym
 import numpy as np
 import pytest
 import stable_baselines3 as sb3
@@ -122,26 +121,6 @@ def test_wrapper_class(tmpdir, rng):
         replay_buffer_wrapper._get_samples()
 
 
-class ActionIsObsEnv(gym.Env):
-    """Simple environment where the obs is the action."""
-
-    def __init__(self):
-        """Initialize environment."""
-        super().__init__()
-        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
-        reward = 0
-        done = False
-        info = {}
-        return obs, reward, done, info
-
-    def reset(self):
-        return np.array([0])
-
-
 def test_replay_buffer_view_provides_buffered_observations():
     space = spaces.Box(np.array([0]), np.array([5]))
     n_envs = 2

tests/scripts/test_scripts.py

@@ -254,6 +254,20 @@ def test_train_preference_comparisons_reward_named_config(tmpdir, named_configs)
     assert isinstance(run.result, dict)
 
 
+def test_train_preference_comparisons_pebble_config(tmpdir):
+    config_updates = dict(common=dict(log_root=tmpdir))
+    run = train_preference_comparisons.train_preference_comparisons_ex.run(
+        # make sure rl.sac named_config is called after rl.fast to overwrite
+        # rl_kwargs.batch_size to None
+        named_configs=ALGO_FAST_CONFIGS["preference_comparison"]
+        + ["pebble", "mountain_car_continuous"],
+        config_updates=config_updates,
+    )
+    assert run.config["rl"]["rl_cls"] is stable_baselines3.SAC
+    assert run.status == "COMPLETED"
+    assert isinstance(run.result, dict)
+
+
 def test_train_dagger_main(tmpdir):
     with pytest.warns(None) as record:
         run = train_imitation.train_imitation_ex.run(

tests/scripts/test_train_preference_comparisons.py

@@ -52,6 +52,9 @@ def test_creates_normalized_entropy_pebble_reward():
             atol=0.05,
         )
 
+        # Just to make coverage happy:
+        reward_fn_stub(state, PLACEHOLDER, PLACEHOLDER, PLACEHOLDER)
+
 
 def reward_fn_stub(state, action, next_state, done):
     return state