[BugFix] Fix behavior or partial, nested dones in PEnv and TEnv

[vmoens]

Stack from ghstack (oldest at bottom):

• -> [BugFix] Fix behavior or partial, nested dones in PEnv and TEnv #2959

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/2959

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 4 New Failures, 6 Pending, 2 Unrelated Failures

As of commit 9185c7a with merge base 36f34da ():

NEW FAILURES - The following jobs have failed:

• Continuous Benchmark (PR) / GPU Pytest benchmark (gh)
  Process completed with exit code 1.
• Habitat Tests on Linux / tests (3.9, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t fa826afc96321a2eec2a2a4866c71272520a64af7f2887599d91088a8797a399 /exec failed with exit code 1
• Libs Tests on Linux / unittests-jumanji (3.10, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t 97c691eda2fed4677515c585cdda67fa3cba238f982235753397e8798ff49d38 /exec failed with exit code 245
• Unit-tests on Linux / tests-optdeps (3.11, 12.8) / linux-job (gh)
  test/test_transforms.py::TestVecNormV2::test_vecnorm_parallel_auto[5]

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Libs Tests on Linux / unittests-gym (3.9, 12.8) / linux-job (gh) (trunk failure)
  test/test_libs.py::TestGym::test_gym_fake_td[True-False-3-HalfCheetah-v2]
• Unit-tests on Windows / unittests-cpu (3.10, windows.4xlarge, cpu) / windows-job (gh) (trunk failure)
  test/test_transforms.py::TestTimer::test_transform_env

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[thomasbbrunner]

Nice! Thanks for the quick fix.

[thomasbbrunner]

I think there's a mismatch in the asserts and comments and some asserts are duplicated. Maybe this helps:

Suggested change

	if done_at_root:
	assert parallel_env._simple_done
	assert transformed_env._simple_done
	# We expect each env to have reached a done state once.
	assert parallel_td["next", "done"].sum().item() == 2
	# We expect env[0] to have been reset and executed 2 steps.
	# We expect env[1] to have just been reset (0 steps).
	assert parallel_env._counter() == [2, 0]
	assert parallel_td["next", "done"].sum().item() == 2
	# We expect each env to have reached a done state once.
	assert transformed_td["next", "done"].sum().item() == 2
	assert_allclose_td(transformed_td, parallel_td, intersection=True)
	# We expect env[0] to have been reset and executed 2 steps.
	# We expect env[1] to have just been reset (0 steps).
	# We only expect env[0] to have reached a done state.
	else:
	assert not parallel_env._simple_done
	assert not transformed_env._simple_done
	assert ("next", "done") not in parallel_td
	assert ("next", "done") not in transformed_td
	assert parallel_td["next", "agent_1", "done"].sum().item() == 2
	assert parallel_env._counter() == [2, 0]
	assert parallel_td["next", "agent_1", "done"].sum().item() == 2
	assert transformed_td["next", "agent_1", "done"].sum().item() == 2
	assert_allclose_td(transformed_td, parallel_td, intersection=True)
	assert transformed_env._counter() == [2, 0]
	# We expect env[0] to have been reset and executed 2 steps.
	# We expect env[1] to have just been reset (0 steps).
	assert parallel_env._counter() == [2, 0]
	assert transformed_env._counter() == [2, 0]
	if done_at_root:
	assert parallel_env._simple_done
	assert transformed_env._simple_done
	# We expect each env to have reached a done state once.
	assert parallel_td["next", "done"].sum().item() == 2
	assert transformed_td["next", "done"].sum().item() == 2
	assert_allclose_td(transformed_td, parallel_td, intersection=True)
	else:
	assert not parallel_env._simple_done
	assert not transformed_env._simple_done
	assert ("next", "done") not in parallel_td
	assert ("next", "done") not in transformed_td
	assert parallel_td["next", "agent_1", "done"].sum().item() == 2
	assert transformed_td["next", "agent_1", "done"].sum().item() == 2
	assert_allclose_td(transformed_td, parallel_td, intersection=True)