[BugFix] Remove reset on last step of a rollout (#1936)

Co-authored-by: vmoens <[email protected]>

test/test_env.py

@@ -217,6 +217,35 @@ def test_rollout(env_name, frame_skip, seed=0):
  env.close()
 
 
+@pytest.mark.parametrize("max_steps", [1, 5])
+def test_rollouts_chaining(max_steps, batch_size=(4,), epochs=4):
+ # CountingEnv is done at max_steps + 1, so to emulate it being done at max_steps, we feed max_steps=max_steps - 1
+ env = CountingEnv(max_steps=max_steps - 1, batch_size=batch_size)
+ policy = CountingEnvCountPolicy(
+ action_spec=env.action_spec, action_key=env.action_key
+ )
+
+ input_td = env.reset()
+ for _ in range(epochs):
+ rollout_td = env.rollout(
+ max_steps=max_steps,
+ policy=policy,
+ auto_reset=False,
+ break_when_any_done=False,
+ tensordict=input_td,
+ )
+ assert (env.count == max_steps).all()
+ input_td = step_mdp(
+ rollout_td[..., -1],
+ keep_other=True,
+ exclude_action=False,
+ exclude_reward=True,
+ reward_keys=env.reward_keys,
+ action_keys=env.action_keys,
+ done_keys=env.done_keys,
+ )
+
+
 @pytest.mark.parametrize("device", get_default_devices())
 def test_rollout_predictability(device):
  env = MockSerialEnv(device=device)

torchrl/envs/common.py

@@ -2273,7 +2273,9 @@ def rollout(
  called on the sub-envs that are done. Default is True.
  return_contiguous (bool): if False, a LazyStackedTensorDict will be returned. Default is True.
  tensordict (TensorDict, optional): if auto_reset is False, an initial
- tensordict must be provided.
+ tensordict must be provided. Rollout will check if this tensordict has done flags and reset the
+ environment in those dimensions (if needed). This normally should not occur if ``tensordict`` is the
+ output of a reset, but can occur if ``tensordict`` is the last step of a previous rollout.
 
  Returns:
  TensorDict object containing the resulting trajectory.
@@ -2369,6 +2371,26 @@ def rollout(
  >>> print(rollout.names)
  [None, 'time']
 
+ Rollouts can be used in a loop to emulate data collection.
+ To do so, you need to pass as input the last tensordict coming from the previous rollout after calling
+ :func:`~torchrl.envs.utils.step_mdp` on it.
+
+ Examples:
+ >>> from torchrl.envs import GymEnv, step_mdp
+ >>> env = GymEnv("CartPole-v1")
+ >>> epochs = 10
+ >>> input_td = env.reset()
+ >>> for i in range(epochs):
+ ... rollout_td = env.rollout(
+ ... max_steps=100,
+ ... break_when_any_done=False,
+ ... auto_reset=False,
+ ... tensordict=input_td,
+ ... )
+ ... input_td = step_mdp(
+ ... rollout_td[..., -1],
+ ... )
+
  """
  if auto_cast_to_device:
  try:
@@ -2388,6 +2410,9 @@ def rollout(
  tensordict = self.reset()
  elif tensordict is None:
  raise RuntimeError("tensordict must be provided when auto_reset is False")
+ else:
+ tensordict = self.maybe_reset(tensordict)
+
  if policy is None:
 
  policy = self.rand_action
@@ -2493,7 +2518,10 @@ def _rollout_nonstop(
  tensordict_ = tensordict_.to(env_device, non_blocking=True)
  else:
  tensordict_.clear_device_()
- tensordict, tensordict_ = self.step_and_maybe_reset(tensordict_)
+ if i == max_steps - 1:
+ tensordict = self.step(tensordict_)
+ else:
+ tensordict, tensordict_ = self.step_and_maybe_reset(tensordict_)
  tensordicts.append(tensordict)
  if i == max_steps - 1:
  # we don't truncated as one could potentially continue the run
@@ -2557,14 +2585,28 @@ def step_and_maybe_reset(
  action_keys=self.action_keys,
  done_keys=self.done_keys,
  )
+ tensordict_ = self.maybe_reset(tensordict_)
+ return tensordict, tensordict_
+
+ def maybe_reset(self, tensordict: TensorDictBase) -> TensorDictBase:
+ """Checks the done keys of the input tensordict and, if needed, resets the environment where it is done.
+
+ Args:
+ tensordict (TensorDictBase): a tensordict coming from the output of :func:`~torchrl.envs.utils.step_mdp`.
+
+ Returns:
+ A tensordict that is identical to the input where the environment was
+ not reset and contains the new reset data where the environment was reset.
+
+ """
  any_done = _terminated_or_truncated(
- tensordict_,
+ tensordict,
  full_done_spec=self.output_spec["full_done_spec"],
  key="_reset",
  )
  if any_done:
- tensordict_ = self.reset(tensordict_)
- return tensordict, tensordict_
+ tensordict = self.reset(tensordict)
+ return tensordict
 
  def empty_cache(self):
  """Erases all the cached values.