Bug in assert statement in RewardNet that prevents using dict observations even if properly linearized by a feature extractor.

[WilliamDormer]

Bug description

In predict_th, there's an assert statement that says the following:
assert rew_th.shape == state.shape[:1]
This will fail, even if you've modified state_th using self.preprocess to be a valid tensor.

Steps to reproduce

Attempt to pre-process a dictionary style state with the preprocess function of the reward network. Even if you return a valid state_th, it checks against the original state, which is incorrect.

Instead I believe it should be:
assert rew_th.shape == state_th.shape[:1]

I believe this was a typo in the implementation.

Environment

• Linux 24.04
• Python version: 3.10.17

[jarvasquezromo]

Context:
I am encountering an AssertionError when using SimpleDAggerTrainer with gymnasium==0.29.1 and imitation==1.0.1. The error triggers specifically during the .train() phase when the DAgger algorithm attempts to reset the vectorized environment. I'm trying to train my own custom environment and even if I wrapper the env with special function to return an array it's failing because the train cant use the dict type observation from gym.

The Traceback:

  File ".../imitation/algorithms/dagger.py", line 669, in train
    trajectories = rollout.generate_trajectories(
  File ".../imitation/data/rollout.py", line 418, in generate_trajectories
    obs = venv.reset()
  File ".../imitation/algorithms/dagger.py", line 224, in reset
    assert isinstance(obs, np.ndarray)
AssertionError

What is happening:

1. API Mismatch: Since Gymnasium v0.26, env.reset() returns a tuple of (observation, info).
2. Internal Wrapper Logic: The SimpleDAggerTrainer uses a specific wrapper logic for DAgger that intercepts the reset() call.
3. The Failure: At imitation/algorithms/dagger.py:224, the code expects venv.reset() to return only a numpy.ndarray (the observation). Because it is receiving the new Gymnasium tuple (obs, info), the assertion fails.
4. Vectorized Env Behavior: Even when using stable-baselines3's DummyVecEnv, the way imitation interacts with the underlying environment within the DAgger-specific logic is not stripping the info dictionary correctly for this specific algorithm.

Environment:

• gymnasium: 0.29.1
• stable-baselines3: 2.2.1
• imitation: 1.0.1
• Python: 3.10