[BugFix] Vmas expanded specs (#942)

Author: matteobettini

torchrl/envs/libs/vmas.py

@@ -2,6 +2,7 @@
 
 import torch
 from tensordict.tensordict import TensorDict, TensorDictBase
+
 from torchrl.data import CompositeSpec, DEVICE_TYPING, UnboundedContinuousTensorSpec
 from torchrl.envs.common import _EnvWrapper, EnvBase
 from torchrl.envs.libs.gym import _gym_to_torchrl_spec_transform
@@ -107,25 +108,6 @@ def __init__(
                 raise TypeError("Env device is different from vmas device")
             kwargs["device"] = str(env.device)
         super().__init__(**kwargs)
-        if len(self.batch_size) == 0:
-            # Batch size not set
-            self.batch_size = torch.Size((self.num_envs,))
-        elif len(self.batch_size) == 1:
-            # Batch size is set
-            if not self.batch_size[0] == self.num_envs:
-                raise TypeError(
-                    "Batch size used in constructor does not match vmas batch size."
-                )
-        else:
-            raise TypeError(
-                "Batch size used in constructor is not compatible with vmas."
-            )
-        self.batch_size = torch.Size([self.n_agents, *self.batch_size])
-        self.input_spec = self.input_spec.expand(self.batch_size)
-        self.observation_spec = self.observation_spec.expand(self.batch_size)
-        self.reward_spec = self.reward_spec.expand(
-            [*self.batch_size, *self.reward_spec.shape]
-        )
 
     @property
     def lib(self):
@@ -144,6 +126,22 @@ def _build_env(
         if self.from_pixels:
             raise NotImplementedError("vmas rendering not yet implemented")
 
+        # Adjust batch size
+        if len(self.batch_size) == 0:
+            # Batch size not set
+            self.batch_size = torch.Size((env.num_envs,))
+        elif len(self.batch_size) == 1:
+            # Batch size is set
+            if not self.batch_size[0] == env.num_envs:
+                raise TypeError(
+                    "Batch size used in constructor does not match vmas batch size."
+                )
+        else:
+            raise TypeError(
+                "Batch size used in constructor is not compatible with vmas."
+            )
+        self.batch_size = torch.Size([env.n_agents, *self.batch_size])
+
         return env
 
     def _make_specs(