[Feature] Re-enable cache for specs

[vmoens]

Stack from ghstack (oldest at bottom):

• -> [Feature] Re-enable cache for specs #2730
• [Feature] lock_ / unlock_ graphs #2729
• [BugFix] Temporarily remove unsafe caching in envs #2728

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 7 New Failures, 1 Cancelled Job, 7 Pending, 5 Unrelated Failures

As of commit 2d84e84 with merge base cd4f359 ():

NEW FAILURES - The following jobs have failed:

• Build Windows Wheels / pytorch/rl / upload / wheel-py3_9-cuda12_4 (gh)
  Unable to download artifact(s): Artifact not found for name: pytorch_rl__3.9_cu124_
• 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.
• Libs Tests on Linux / unittests-minari (3.9, 12.4) / linux-job (gh)
  test/test_libs.py::TestMinari::test_load[minigrid/BabyAI-KeyCorridorS6R3/optimal-v0-True]
• Libs Tests on Linux / unittests-vd4rl (3.9, 12.4) / linux-job (gh)
  test/test_libs.py::TestVD4RL::test_vd4rl_preproc
• Unit-tests on Windows / unittests-cpu / windows-job (gh)
  Process completed with exit code 1.
• Wheels / build-wheel-windows (3.11, 3.11) (gh)

CANCELLED JOB - The following job was cancelled. Please retry:

• Wheels / build-wheel-windows (3.10, 3.10.3) (gh)
  ##[error]The operation was canceled.

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-cuda11_8 (gh) (trunk failure)
  Unable to download artifact(s): Artifact not found for name: pytorch_rl__3.9_cu118_
• 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.4) / linux-job (gh) (trunk failure)
  AttributeError: _ARRAY_API not found
• Libs Tests on Linux / unittests-chess (3.9, 12.4) / linux-job (gh) (trunk failure)
  test/test_env.py::TestChessEnv::test_chess_rendering[True]

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

[vmoens]

This is a good clean-up PR:

• Now specs are locked recursively if the env is locked. That was always the intention but not enforced, which is dangerous because without that your cached keys can be out of sync.

• To change the specs you must unlock the env, which clears the cache.

• The lock is similar to tensordict's and uses weakrefs to sub-specs to build a graph. You cannot unlock a node in the middle of the graph.

• The locking occurs in the metaclass, which skips the super().__init__ and means that you can pass that kwarg even if it's not in your __init__.

• I also made it possible to just pass a regular spec to the observation_spec in which case we assume it's a "observation" entry.