Update

[ghstack-poisoned]
pytorch · Nov 20, 2024 · f9c4e00 · f9c4e00
1 parent b59d5de
commit f9c4e00
Show file tree

Hide file tree

Showing 5 changed files with 17 additions and 17 deletions.
diff --git a/sota-implementations/multiagent/iql.py b/sota-implementations/multiagent/iql.py
@@ -91,7 +91,7 @@ def train(cfg: "DictConfig"):  # noqa: F821
             ("agents", "action_value"),
             ("agents", "chosen_action_value"),
         ],
-        spec=env.action_spec_unbatched,
+        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.action_spec_unbatched,
+            spec=env.full_action_spec_unbatched,
         ),
     )
 

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

diff --git a/sota-implementations/multiagent/mappo_ippo.py b/sota-implementations/multiagent/mappo_ippo.py
@@ -92,13 +92,13 @@ def train(cfg: "DictConfig"):  # noqa: F821
     )
     policy = ProbabilisticActor(
         module=policy_module,
-        spec=env.action_spec_unbatched,
+        spec=env.full_action_spec_unbatched,
         in_keys=[("agents", "loc"), ("agents", "scale")],
         out_keys=[env.action_key],
         distribution_class=TanhNormal,
         distribution_kwargs={
-            "low": env.action_spec_unbatched[("agents", "action")].space.low,
-            "high": env.action_spec_unbatched[("agents", "action")].space.high,
+            "low": env.full_action_spec_unbatched[("agents", "action")].space.low,
+            "high": env.full_action_spec_unbatched[("agents", "action")].space.high,
         },
         return_log_prob=True,
     )

diff --git a/sota-implementations/multiagent/qmix_vdn.py b/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.action_spec_unbatched,
+        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.action_spec_unbatched,
+            spec=env.full_action_spec_unbatched,
         ),
     )
 

diff --git a/sota-implementations/multiagent/sac.py b/sota-implementations/multiagent/sac.py
@@ -96,13 +96,13 @@ def train(cfg: "DictConfig"):  # noqa: F821
 
         policy = ProbabilisticActor(
             module=policy_module,
-            spec=env.action_spec_unbatched,
+            spec=env.full_action_spec_unbatched,
             in_keys=[("agents", "loc"), ("agents", "scale")],
             out_keys=[env.action_key],
             distribution_class=TanhNormal,
             distribution_kwargs={
-                "low": env.action_spec_unbatched[("agents", "action")].space.low,
-                "high": env.action_spec_unbatched[("agents", "action")].space.high,
+                "low": env.full_action_spec_unbatched[("agents", "action")].space.low,
+                "high": env.full_action_spec_unbatched[("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.action_spec_unbatched,
+            spec=env.full_action_spec_unbatched,
             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.action_spec_unbatched,
+            action_spec=env.full_action_spec_unbatched,
         )
         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.action_spec_unbatched,
+            action_space=env.full_action_spec_unbatched,
         )
         loss_module.set_keys(
             action_value=("agents", "action_value"),