[BugFix] Use single_action_spec whenever necessary

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

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.single_action_spec.space.low,
+            "high": env.single_action_spec.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.single_action_spec.space.low.to(device),
+            "high": proof_environment.single_action_spec.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.single_action_spec.space.low.to(device),
+        "high": proof_environment.single_action_spec.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.single_action_spec
 
     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.single_action_spec
     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.single_action_spec
     for key, value in proof_environment.observation_spec.items():
         if key == "observation":
             state_dim = value.shape[-1]

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.single_action_spec.space.low,
+        "high": proof_environment.single_action_spec.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.single_action_spec
     actor_net, q_net, value_net = make_iql_modules_state(model_cfg, eval_env)
 
     out_keys = ["loc", "scale"]

sota-implementations/multiagent/iql.py

@@ -91,7 +91,7 @@ def train(cfg: "DictConfig"):  # noqa: F821
             ("agents", "action_value"),
             ("agents", "chosen_action_value"),
         ],
-        spec=env.unbatched_action_spec,
+        spec=env.single_action_spec,
         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.single_action_spec,
         ),
     )
 

sota-implementations/multiagent/maddpg_iddpg.py

@@ -91,21 +91,21 @@ def train(cfg: "DictConfig"):  # noqa: F821
     )
     policy = ProbabilisticActor(
         module=policy_module,
-        spec=env.unbatched_action_spec,
+        spec=env.single_action_spec,
         in_keys=[("agents", "param")],
         out_keys=[env.action_key],
         distribution_class=TanhDelta,
         distribution_kwargs={
-            "low": env.unbatched_action_spec[("agents", "action")].space.low,
-            "high": env.unbatched_action_spec[("agents", "action")].space.high,
+            "low": env.single_action_spec[("agents", "action")].space.low,
+            "high": env.single_action_spec[("agents", "action")].space.high,
         },
         return_log_prob=False,
     )
 
     policy_explore = TensorDictSequential(
         policy,
         AdditiveGaussianModule(
-            spec=env.unbatched_action_spec,
+            spec=env.single_action_spec,
             annealing_num_steps=int(cfg.collector.total_frames * (1 / 2)),
             action_key=env.action_key,
             device=cfg.train.device,

sota-implementations/multiagent/mappo_ippo.py

@@ -92,13 +92,13 @@ def train(cfg: "DictConfig"):  # noqa: F821
     )
     policy = ProbabilisticActor(
         module=policy_module,
-        spec=env.unbatched_action_spec,
+        spec=env.single_action_spec,
         in_keys=[("agents", "loc"), ("agents", "scale")],
         out_keys=[env.action_key],
         distribution_class=TanhNormal,
         distribution_kwargs={
-            "low": env.unbatched_action_spec[("agents", "action")].space.low,
-            "high": env.unbatched_action_spec[("agents", "action")].space.high,
+            "low": env.single_action_spec[("agents", "action")].space.low,
+            "high": env.single_action_spec[("agents", "action")].space.high,
         },
         return_log_prob=True,
     )

sota-implementations/multiagent/qmix_vdn.py

@@ -91,7 +91,7 @@ def train(cfg: "DictConfig"):  # noqa: F821
             ("agents", "action_value"),
             ("agents", "chosen_action_value"),
         ],
-        spec=env.unbatched_action_spec,
+        spec=env.single_action_spec,
         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.single_action_spec,
         ),
     )
 

sota-implementations/multiagent/sac.py

@@ -96,13 +96,13 @@ def train(cfg: "DictConfig"):  # noqa: F821
 
         policy = ProbabilisticActor(
             module=policy_module,
-            spec=env.unbatched_action_spec,
+            spec=env.single_action_spec,
             in_keys=[("agents", "loc"), ("agents", "scale")],
             out_keys=[env.action_key],
             distribution_class=TanhNormal,
             distribution_kwargs={
-                "low": env.unbatched_action_spec[("agents", "action")].space.low,
-                "high": env.unbatched_action_spec[("agents", "action")].space.high,
+                "low": env.single_action_spec[("agents", "action")].space.low,
+                "high": env.single_action_spec[("agents", "action")].space.high,
             },
             return_log_prob=True,
         )
@@ -146,7 +146,7 @@ def train(cfg: "DictConfig"):  # noqa: F821
         )
         policy = ProbabilisticActor(
             module=policy_module,
-            spec=env.unbatched_action_spec,
+            spec=env.single_action_spec,
             in_keys=[("agents", "logits")],
             out_keys=[env.action_key],
             distribution_class=OneHotCategorical
@@ -194,7 +194,7 @@ def train(cfg: "DictConfig"):  # noqa: F821
             actor_network=policy,
             qvalue_network=value_module,
             delay_qvalue=True,
-            action_spec=env.unbatched_action_spec,
+            action_spec=env.single_action_spec,
         )
         loss_module.set_keys(
             state_action_value=("agents", "state_action_value"),
@@ -209,7 +209,7 @@ def train(cfg: "DictConfig"):  # noqa: F821
             qvalue_network=value_module,
             delay_qvalue=True,
             num_actions=env.action_spec.space.n,
-            action_space=env.unbatched_action_spec,
+            action_space=env.single_action_spec,
         )
         loss_module.set_keys(
             action_value=("agents", "action_value"),

sota-implementations/ppo/utils_atari.py

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

sota-implementations/ppo/utils_mujoco.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.single_action_spec.space.low,
+        "high": proof_environment.single_action_spec.space.high,
         "tanh_loc": False,
     }
 

sota-implementations/redq/utils.py

@@ -410,7 +410,7 @@ def make_redq_model(
     default_policy_scale = cfg.network.default_policy_scale
     gSDE = cfg.exploration.gSDE
 
-    action_spec = proof_environment.action_spec
+    action_spec = proof_environment.single_action_spec
 
     if actor_net_kwargs is None:
         actor_net_kwargs = {}

sota-implementations/sac/utils.py

@@ -161,9 +161,7 @@ def make_sac_agent(cfg, train_env, eval_env, device):
     """Make SAC 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.single_action_spec
     actor_net_kwargs = {
         "num_cells": cfg.network.hidden_sizes,
         "out_features": 2 * action_spec.shape[-1],

test/mocking_classes.py

@@ -1394,7 +1394,7 @@ def _make_specs(self):
             device=self.device,
         )
 
-        self.unbatched_action_spec = Composite(
+        self.single_action_spec = Composite(
             lazy=action_specs,
             device=self.device,
         )
@@ -1423,8 +1423,8 @@ def _make_specs(self):
             device=self.device,
         )
 
-        self.action_spec = self.unbatched_action_spec.expand(
-            *self.batch_size, *self.unbatched_action_spec.shape
+        self.action_spec = self.single_action_spec.expand(
+            *self.batch_size, *self.single_action_spec.shape
         )
         self.observation_spec = self.unbatched_observation_spec.expand(
             *self.batch_size, *self.unbatched_observation_spec.shape
@@ -1610,8 +1610,8 @@ def __init__(self, max_steps: int = 5, start_val: int = 0, **kwargs):
 
         self.make_specs()
 
-        self.action_spec = self.unbatched_action_spec.expand(
-            *self.batch_size, *self.unbatched_action_spec.shape
+        self.action_spec = self.single_action_spec.expand(
+            *self.batch_size, *self.single_action_spec.shape
         )
         self.observation_spec = self.unbatched_observation_spec.expand(
             *self.batch_size, *self.unbatched_observation_spec.shape
@@ -1642,7 +1642,7 @@ def make_specs(self):
             ),
         )
 
-        self.unbatched_action_spec = Composite(
+        self.single_action_spec = Composite(
             nested_1=Composite(
                 action=Categorical(n=2, shape=(self.nested_dim_1,)),
                 shape=(self.nested_dim_1,),

test/test_libs.py

@@ -2253,6 +2253,11 @@ def test_vmas_batch_size(self, scenario_name, num_envs, n_agents):
             max_steps=n_rollout_samples,
             return_contiguous=False if env.het_specs else True,
         )
+        assert env.single_full_action_spec.shape == env.unbatched_action_spec.shape, (
+            env.action_spec,
+            env.batch_size,
+        )
+
         env.close()
 
         if env.het_specs:

tutorials/sphinx-tutorials/coding_ppo.py

@@ -431,8 +431,8 @@
     in_keys=["loc", "scale"],
     distribution_class=TanhNormal,
     distribution_kwargs={
-        "low": env.action_spec.space.low,
-        "high": env.action_spec.space.high,
+        "low": env.single_action_spec.space.low,
+        "high": env.single_action_spec.space.high,
     },
     return_log_prob=True,
     # we'll need the log-prob for the numerator of the importance weights

tutorials/sphinx-tutorials/multiagent_competitive_ddpg.py

@@ -486,8 +486,8 @@
         out_keys=[(group, "action")],
         distribution_class=TanhDelta,
         distribution_kwargs={
-            "low": env.full_action_spec[group, "action"].space.low,
-            "high": env.full_action_spec[group, "action"].space.high,
+            "low": env.single_full_action_spec[group, "action"].space.low,
+            "high": env.single_full_action_spec[group, "action"].space.high,
         },
         return_log_prob=False,
     )

tutorials/sphinx-tutorials/multiagent_ppo.py

@@ -445,13 +445,13 @@
 
 policy = ProbabilisticActor(
     module=policy_module,
-    spec=env.unbatched_action_spec,
+    spec=env.single_action_spec,
     in_keys=[("agents", "loc"), ("agents", "scale")],
     out_keys=[env.action_key],
     distribution_class=TanhNormal,
     distribution_kwargs={
-        "low": env.unbatched_action_spec[env.action_key].space.low,
-        "high": env.unbatched_action_spec[env.action_key].space.high,
+        "low": env.single_action_spec[env.action_key].space.low,
+        "high": env.single_action_spec[env.action_key].space.high,
     },
     return_log_prob=True,
     log_prob_key=("agents", "sample_log_prob"),