diff --git a/test/transforms/test_env_transforms.py b/test/transforms/test_env_transforms.py index ce3b61d07ed..2a97ca7e230 100644 --- a/test/transforms/test_env_transforms.py +++ b/test/transforms/test_env_transforms.py @@ -501,6 +501,61 @@ def test_stepcounter_marl_truncation(self): : max_steps - 1 ].any(), "Root done should be False before max_steps" + @pytest.mark.parametrize("truncating_group", ["left", "right"]) + @pytest.mark.parametrize("update_done", [False, True]) + def test_stepcounter_independent_nested_groups(self, truncating_group, update_done): + class GroupedEnv(EnvBase): + def __init__(self): + super().__init__() + self.observation_spec = Composite(observation=Unbounded(shape=(1,))) + self.action_spec = Categorical(2, shape=(1,)) + self.reward_spec = Unbounded(shape=(1,)) + flags = Composite( + { + key: Categorical(2, shape=(1,), dtype=torch.bool) + for key in ("done", "terminated", "truncated") + } + ) + group = flags.clone() + group["agents"] = flags.expand(3) + self.done_spec = Composite(left=group.clone(), right=group.clone()) + + def _reset(self, tensordict=None, **kwargs): + return self.observation_spec.zero().update(self.full_done_spec.zero()) + + def _step(self, tensordict): + result = self._reset().update(self.full_reward_spec.zero()) + result["right", "agents", "truncated"][0] = True + result["right", "agents", "done"][0] = True + return result + + def _set_seed(self, seed): + pass + + env = TransformedEnv(GroupedEnv(), StepCounter(3, update_done=update_done)) + try: + td = env.reset() + td[truncating_group, "step_count"].fill_(2) + rollout = env.rollout(1, tensordict=td, auto_reset=False) + for group in ("left", "right"): + truncated = group == truncating_group + torch.testing.assert_close( + rollout["next", group, "truncated"], torch.tensor([[truncated]]) + ) + agent_done = torch.tensor([group == "right", False, False]).reshape( + 1, 3, 1 + ) + torch.testing.assert_close( + rollout["next", group, "agents", "truncated"], + agent_done | truncated, + ) + torch.testing.assert_close( + rollout["next", group, "agents", "done"], + agent_done | (truncated and update_done), + ) + finally: + env.close() + class _AgentNestedCountingEnv(EnvBase): """Counting env whose done keys live under the ``agent`` nest.""" diff --git a/torchrl/envs/transforms/_env.py b/torchrl/envs/transforms/_env.py index e20b1e99077..ff880efc3a4 100644 --- a/torchrl/envs/transforms/_env.py +++ b/torchrl/envs/transforms/_env.py @@ -862,6 +862,8 @@ class StepCounter(Transform): .. note:: To ensure compatibility with environments that have multiple done_key(s), this transform will write a step_count entry for every done entry within the tensordict. + Truncation is propagated to nested done entries only from their own + parent group, not from independent sibling groups. Examples: >>> import gymnasium @@ -1143,6 +1145,11 @@ def _propagate_to_nested_keys(self, next_tensordict: TensorDictBase) -> None: continue # Find a parent truncated key to propagate from for parent_key in self.truncated_keys: + if ( + isinstance(parent_key, tuple) + and nested_key[: len(parent_key) - 1] != parent_key[:-1] + ): + continue parent_truncated = next_tensordict.get(parent_key, None) if parent_truncated is not None: # Insert extra dims (e.g. agent dims) so the parent is @@ -1164,6 +1171,11 @@ def _propagate_to_nested_keys(self, next_tensordict: TensorDictBase) -> None: continue # Find a parent done key to propagate from for parent_key in self.done_keys: + if ( + isinstance(parent_key, tuple) + and nested_key[: len(parent_key) - 1] != parent_key[:-1] + ): + continue parent_done = next_tensordict.get(parent_key, None) if parent_done is not None: parent_val = parent_done