[BugFix] Remove reset on last step of a rollout

torchrl/envs/common.py

@@ -2388,6 +2388,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 +2496,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 +2563,18 @@ 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:
  any_done = _terminated_or_truncated(
  tensordict_,
  full_done_spec=self.output_spec["full_done_spec"],
  key="_reset",
  )
  if any_done:
  tensordict_ = self.reset(tensordict_)
- return tensordict, tensordict_
+ return tensordict_
 
  def empty_cache(self):
  """Erases all the cached values.