Update

[ghstack-poisoned]

sota-implementations/multiagent/qmix_vdn.py

@@ -110,7 +110,7 @@ def train(cfg: "DictConfig"):  # noqa: F821
     if cfg.loss.mixer_type == "qmix":
         mixer = TensorDictModule(
             module=QMixer(
-                state_shape=env.unbatched_observation_spec[
+                state_shape=env.observation_spec_unbatched[
                     "agents", "observation"
                 ].shape,
                 mixing_embed_dim=32,

test/mocking_classes.py

@@ -1040,7 +1040,9 @@ def _step(
         action = tensordict.get(self.action_key)
         self.count += action.to(
             dtype=torch.int,
-            device=self.action_spec.device if self.device is None else self.device,
+            device=self.full_action_spec[self.action_key].device
+            if self.device is None
+            else self.device,
         )
         tensordict = TensorDict(
             source={

torchrl/envs/common.py

@@ -172,10 +172,11 @@ def __call__(cls, *args, **kwargs):
         # we create the done spec by adding a done/terminated entry if one is missing
         instance._create_done_specs()
         # we access lazy attributed to make sure they're built properly.
-        # This isn't done in `__init__` because we don't know if supre().__init__
+        # This isn't done in `__init__` because we don't know if super().__init__
         # will be called before or after the specs, batch size etc are set.
         _ = instance.done_spec
-        _ = instance.reward_spec
+        _ = instance.reward_keys
+        _ = instance.action_keys
         _ = instance.state_spec
         if auto_reset:
             from torchrl.envs.transforms.transforms import (