[BugFix] patch rand_action in TransformedEnv to read the base_env method

[vmoens]

Stack from ghstack (oldest at bottom):

• [Feature] ConditionalPolicySwitch transform #2711
• [Example] Self-play chess PPO example #2709
• [Feature] Add deterministic_sample to masked categorical #2708
• [Feature,Refactor] Chess improvements: fen, pgn, pixels, san #2702
• [Feature] UnaryTransform for input entries #2700
• -> [BugFix] patch rand_action in TransformedEnv to read the base_env method #2699

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 13 New Failures, 5 Unrelated Failures

As of commit 1024d61 with merge base 319bb68 ():

NEW FAILURES - The following jobs have failed:

• Build Windows Wheels / pytorch/rl / upload / wheel-py3_9-cuda11_8 (gh)
  Unable to download artifact(s): Artifact not found for name: pytorch_rl__3.9_cu118_
• Continuous Benchmark (PR) / CPU Pytest benchmark (gh)
  Process completed with exit code 1.
• Continuous Benchmark (PR) / GPU Pytest benchmark (gh)
  Process completed with exit code 1.
• SOTA Tests on Linux / tests (3.9, 12.1) / linux-job (gh)
  RuntimeError: Command docker exec -t 1ea5775e048494341073d1697f99ce3c738640567d3cefc8351794e8c08c63ea /exec failed with exit code 1
• Unit-tests on Linux / tests-cpu (3.10) / linux-job (gh)
  test/test_transforms.py::TestActionDiscretizer::test_transform_env[env_cls3-SamplingStrategy.RANDOM-False-True]
• Unit-tests on Linux / tests-cpu (3.11) / linux-job (gh)
  test/test_transforms.py::TestActionDiscretizer::test_transform_env[env_cls3-3-False-True]
• Unit-tests on Linux / tests-cpu (3.12) / linux-job (gh)
  test/test_transforms.py::TestActionDiscretizer::test_transform_env[env_cls3-3-False-True]
• Unit-tests on Linux / tests-cpu (3.9) / linux-job (gh)
  test/test_transforms.py::TestActionDiscretizer::test_transform_env[env_cls3-SamplingStrategy.RANDOM-False-True]
• Unit-tests on Linux / tests-cpu-oldget (3.12) / linux-job (gh)
  test/test_transforms.py::TestActionDiscretizer::test_transform_env[env_cls3-3-False-True]
• Unit-tests on Linux / tests-gpu (3.11, 12.1) / linux-job (gh)
  test/test_transforms.py::TestActionDiscretizer::test_transform_env[env_cls3-3-False-True]
• Unit-tests on Linux / tests-olddeps (3.8, 11.6) / linux-job (gh)
  test/test_transforms.py::TestActionDiscretizer::test_transform_env[env_cls3-SamplingStrategy.RANDOM-False-True]
• Unit-tests on Linux / tests-optdeps (3.11, 12.1) / linux-job (gh)
  test/test_transforms.py::TestActionDiscretizer::test_transform_compose
• Unit-tests on Linux / tests-stable-gpu (3.10, 11.8) / linux-job (gh)
  test/test_transforms.py::TestActionDiscretizer::test_transform_env[env_cls3-SamplingStrategy.RANDOM-False-True]

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

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

• Build Windows Wheels / pytorch/rl / upload / wheel-py3_9-cpu (gh) (trunk failure)
  Unable to download artifact(s): Artifact not found for name: pytorch_rl__3.9_cpu_
• Build Windows Wheels / pytorch/rl / upload / wheel-py3_9-cuda12_4 (gh) (trunk failure)
• Build Windows Wheels / pytorch/rl / upload / wheel-py3_9-cuda12_6 (gh) (trunk failure)
  Unable to download artifact(s): Artifact not found for name: pytorch_rl__3.9_cu126_
• Habitat Tests on Linux / tests (3.9, 12.1) / linux-job (gh) (trunk failure)
  AttributeError: _ARRAY_API not found
• Unit-tests on Windows / unittests-cpu / windows-job (gh) (trunk failure)
  Process completed with exit code 1.

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

[vmoens]

@kurtamohler I'm not super happy with this, as the comment says it's far from accounting for transforms that do some kind of inverse mapping but the point is that without that, a transformed chess env cannot generate random actions (because the rand_action is not the same as the one in EnvBase that would otherwise be used by TransformedEnv).

[kurtamohler]

Is there a reason we couldn't use self.action_spec.rand()?

>>> import torchrl
>>> env = torchrl.envs.PendulumEnv().append_transform(torchrl.envs.ActionDiscretizer(num_intervals=4))
>>> env.action_spec.rand()
tensor([3])
>>> env.action_spec.rand().dtype
torch.int64

[vmoens]

yes, what I meant with that comment is that if your base env redefines the rand_action then the action you'll get won't be transformed

import torchrl
Pendulum = torchrl.envs.PendulumEnv
rand_action = Pendulum.rand_action
Pendulum.rand_action = lambda *args, **kwargs: rand_action(*args, **kwargs)

env = Pendulum().append_transform(torchrl.envs.ActionDiscretizer(num_intervals=4))

print(env.action_spec.rand())
print(env.action_spec.rand().dtype)
print(env.rand_action())

which prints

tensor([3])
torch.int64
TensorDict(
    fields={
        action: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)