vmas 1.2.6

Author: matteobettini

setup.py

@@ -6,7 +6,7 @@
 
 setup(
     name="vmas",
-    version="1.2.5",
+    version="1.2.6",
     description="Vectorized Multi-Agent Simulator",
     url="https://github.com/proroklab/VectorizedMultiAgentSimulator",
     license="GPLv3",

vmas/simulator/environment/environment.py

@@ -164,16 +164,13 @@ def step(self, actions: List):
         rewards = []
         infos = []
         for agent in self.agents:
-            rewards.append(self.scenario.reward(agent).clone())
-            obs.append(self.scenario.observation(agent).clone())
+            rewards.append(self.scenario.reward(agent))
+            obs.append(self.scenario.observation(agent))
             # A dictionary per agent
             infos.append(self.scenario.info(agent))
 
-        dones = self.scenario.done().clone()
-
         self.steps += 1
-        if self.max_steps is not None:
-            dones += self.steps >= self.max_steps
+        dones = self.done()
         # print("\nStep results in unwrapped environment")
         # print(
         #     f"Actions len (n_agents): {len(actions)}, "
@@ -192,6 +189,12 @@ def step(self, actions: List):
         # print(f"Info len (n_agents): {len(infos)}, info[0] (infos agent 0): {infos[0]}")
         return obs, rewards, dones, infos
 
+    def done(self):
+        dones = self.scenario.done()
+        if self.max_steps is not None:
+            dones += self.steps >= self.max_steps
+        return dones
+
     def get_agent_action_size(self, agent: Agent):
         return (
             self.world.dim_p