[BugFix] action_spec_unbatched whenever necessary

ghstack-source-id: ec87794
Pull Request resolved: #2592

Author: Vincent Moens

examples/distributed/collectors/multi_nodes/ray_train.py

@@ -85,8 +85,8 @@
         in_keys=["loc", "scale"],
         distribution_class=TanhNormal,
         distribution_kwargs={
-            "low": env.action_spec.space.low,
-            "high": env.action_spec.space.high,
+            "low": env.action_spec_unbatched.space.low,
+            "high": env.action_spec_unbatched.space.high,
         },
         return_log_prob=True,
     )

sota-implementations/a2c/utils_atari.py

@@ -101,8 +101,8 @@ def make_ppo_modules_pixels(proof_environment, device):
         num_outputs = proof_environment.action_spec.shape
         distribution_class = TanhNormal
         distribution_kwargs = {
-            "low": proof_environment.action_spec.space.low.to(device),
-            "high": proof_environment.action_spec.space.high.to(device),
+            "low": proof_environment.action_spec_unbatched.space.low.to(device),
+            "high": proof_environment.action_spec_unbatched.space.high.to(device),
         }
 
     # Define input keys

sota-implementations/a2c/utils_mujoco.py

@@ -57,8 +57,8 @@ def make_ppo_models_state(proof_environment, device, *, compile: bool = False):
     num_outputs = proof_environment.action_spec.shape[-1]
     distribution_class = TanhNormal
     distribution_kwargs = {
-        "low": proof_environment.action_spec.space.low.to(device),
-        "high": proof_environment.action_spec.space.high.to(device),
+        "low": proof_environment.action_spec_unbatched.space.low.to(device),
+        "high": proof_environment.action_spec_unbatched.space.high.to(device),
         "tanh_loc": False,
         "safe_tanh": True,
     }

sota-implementations/cql/utils.py

@@ -191,7 +191,7 @@ def make_offline_replay_buffer(rb_cfg):
 def make_cql_model(cfg, train_env, eval_env, device="cpu"):
     model_cfg = cfg.model
 
-    action_spec = train_env.action_spec
+    action_spec = train_env.action_spec_unbatched
 
     actor_net, q_net = make_cql_modules_state(model_cfg, eval_env)
     in_keys = ["observation"]

sota-implementations/crossq/utils.py

@@ -147,9 +147,7 @@ def make_crossQ_agent(cfg, train_env, device):
     """Make CrossQ agent."""
     # Define Actor Network
     in_keys = ["observation"]
-    action_spec = train_env.action_spec
-    if train_env.batch_size:
-        action_spec = action_spec[(0,) * len(train_env.batch_size)]
+    action_spec = train_env.action_spec_unbatched
     actor_net_kwargs = {
         "num_cells": cfg.network.actor_hidden_sizes,
         "out_features": 2 * action_spec.shape[-1],

sota-implementations/decision_transformer/utils.py

@@ -393,7 +393,7 @@ def make_dt_model(cfg):
         make_base_env(env_cfg), env_cfg, obs_loc=0, obs_std=1
     )
 
-    action_spec = proof_environment.action_spec
+    action_spec = proof_environment.action_spec_unbatched
     for key, value in proof_environment.observation_spec.items():
         if key == "observation":
             state_dim = value.shape[-1]

sota-implementations/dreamer/dreamer.py

@@ -20,7 +20,7 @@
 )
 
 # mixed precision training
-from torch.cuda.amp import GradScaler
+from torch.amp import GradScaler
 from torch.nn.utils import clip_grad_norm_
 from torchrl._utils import logger as torchrl_logger, timeit
 from torchrl.envs.utils import ExplorationType, set_exploration_type
@@ -321,6 +321,14 @@ def compile_rssms(module):
 
         t_collect_init = time.time()
 
+    test_env.close()
+    train_env.close()
+    collector.shutdown()
+
+    del test_env
+    del train_env
+    del collector
+
 
 if __name__ == "__main__":
     main()

sota-implementations/gail/ppo_utils.py

@@ -52,8 +52,8 @@ def make_ppo_models_state(proof_environment):
     num_outputs = proof_environment.action_spec.shape[-1]
     distribution_class = TanhNormal
     distribution_kwargs = {
-        "low": proof_environment.action_spec.space.low,
-        "high": proof_environment.action_spec.space.high,
+        "low": proof_environment.action_spec_unbatched.space.low,
+        "high": proof_environment.action_spec_unbatched.space.high,
         "tanh_loc": False,
     }
 

sota-implementations/iql/utils.py

@@ -195,9 +195,7 @@ def make_iql_model(cfg, train_env, eval_env, device="cpu"):
     model_cfg = cfg.model
 
     in_keys = ["observation"]
-    action_spec = train_env.action_spec
-    if train_env.batch_size:
-        action_spec = action_spec[(0,) * len(train_env.batch_size)]
+    action_spec = train_env.action_spec_unbatched
     actor_net, q_net, value_net = make_iql_modules_state(model_cfg, eval_env)
 
     out_keys = ["loc", "scale"]

sota-implementations/multiagent/iql.py

@@ -72,7 +72,7 @@ def train(cfg: "DictConfig"):  # noqa: F821
     # Policy
     net = MultiAgentMLP(
         n_agent_inputs=env.observation_spec["agents", "observation"].shape[-1],
-        n_agent_outputs=env.action_spec.space.n,
+        n_agent_outputs=env.full_action_spec["agents", "action"].space.n,
         n_agents=env.n_agents,
         centralised=False,
         share_params=cfg.model.shared_parameters,
@@ -91,7 +91,7 @@ def train(cfg: "DictConfig"):  # noqa: F821
             ("agents", "action_value"),
             ("agents", "chosen_action_value"),
         ],
-        spec=env.unbatched_action_spec,
+        spec=env.full_action_spec_unbatched,
         action_space=None,
     )
     qnet = SafeSequential(module, value_module)
@@ -103,7 +103,7 @@ def train(cfg: "DictConfig"):  # noqa: F821
             eps_end=0,
             annealing_num_steps=int(cfg.collector.total_frames * (1 / 2)),
             action_key=env.action_key,
-            spec=env.unbatched_action_spec,
+            spec=env.full_action_spec_unbatched,
         ),
     )