[Differentiable] Further remove inplace operations and introduce world.zero_grad()

Author: matteobettini

vmas/simulator/core.py

@@ -282,12 +282,21 @@ def rot(self, rot: Tensor):
         self._rot = rot.to(self._device)
 
     def _reset(self, env_index: typing.Optional[int]):
-        for attr in [self.pos, self.rot, self.vel, self.ang_vel]:
+        for attr_name in ["pos", "rot", "vel", "ang_vel"]:
+            attr = self.__getattribute__(attr_name)
             if attr is not None:
                 if env_index is None:
-                    attr[:] = 0.0
+                    self.__setattr__(attr_name, torch.zeros_like(attr))
                 else:
-                    attr[env_index] = 0.0
+                    self.__setattr__(
+                        attr_name, TorchUtils.where_from_index(env_index, 0, attr)
+                    )
+
+    def zero_grad(self):
+        for attr_name in ["pos", "rot", "vel", "ang_vel"]:
+            attr = self.__getattribute__(attr_name)
+            if attr is not None:
+                self.__setattr__(attr_name, attr.detach())
 
     def _spawn(self, dim_c: int, dim_p: int):
         self.pos = torch.zeros(
@@ -363,14 +372,25 @@ def torque(self, value):
 
     @override(EntityState)
     def _reset(self, env_index: typing.Optional[int]):
-        for attr in [self.c, self.force, self.torque]:
+        for attr_name in ["c", "force", "torque"]:
+            attr = self.__getattribute__(attr_name)
             if attr is not None:
                 if env_index is None:
-                    attr[:] = 0.0
+                    self.__setattr__(attr_name, torch.zeros_like(attr))
                 else:
-                    attr[env_index] = 0.0
+                    self.__setattr__(
+                        attr_name, TorchUtils.where_from_index(env_index, 0, attr)
+                    )
         super()._reset(env_index)
 
+    @override(EntityState)
+    def zero_grad(self):
+        for attr_name in ["c", "force", "torque"]:
+            attr = self.__getattribute__(attr_name)
+            if attr is not None:
+                self.__setattr__(attr_name, attr.detach())
+        super().zero_grad()
+
     @override(EntityState)
     def _spawn(self, dim_c: int, dim_p: int):
         if dim_c > 0:
@@ -492,12 +512,21 @@ def _to_tensor(self, value):
         )
 
     def _reset(self, env_index: typing.Optional[int]):
-        for attr in [self.u, self.c]:
+        for attr_name in ["u", "c"]:
+            attr = self.__getattribute__(attr_name)
             if attr is not None:
                 if env_index is None:
-                    attr[:] = 0.0
+                    self.__setattr__(attr_name, torch.zeros_like(attr))
                 else:
-                    attr[env_index] = 0.0
+                    self.__setattr__(
+                        attr_name, TorchUtils.where_from_index(env_index, 0, attr)
+                    )
+
+    def zero_grad(self):
+        for attr_name in ["u", "c"]:
+            attr = self.__getattribute__(attr_name)
+            if attr is not None:
+                self.__setattr__(attr_name, attr.detach())
 
 
 # properties and state of physical world entity
@@ -690,6 +719,9 @@ def _spawn(self, dim_c: int, dim_p: int):
     def _reset(self, env_index: int):
         self.state._reset(env_index)
 
+    def zero_grad(self):
+        self.state.zero_grad()
+
     def set_pos(self, pos: Tensor, batch_index: int):
         self._set_state_property(EntityState.pos, self.state, pos, batch_index)
 
@@ -988,6 +1020,11 @@ def _reset(self, env_index: int):
         self.dynamics.reset(env_index)
         super()._reset(env_index)
 
+    def zero_grad(self):
+        self.action.zero_grad()
+        self.dynamics.zero_grad()
+        super().zero_grad()
+
     @override(Entity)
     def to(self, device: torch.device):
         super().to(device)
@@ -1112,6 +1149,10 @@ def reset(self, env_index: int):
         for e in self.entities:
             e._reset(env_index)
 
+    def zero_grad(self):
+        for e in self.entities:
+            e.zero_grad()
+
     @property
     def agents(self) -> List[Agent]:
         return self._agents

vmas/simulator/dynamics/common.py

@@ -17,6 +17,9 @@ def __init__(
     def reset(self, index: Union[Tensor, int] = None):
         pass
 
+    def zero_grad(self):
+        pass
+
     @property
     def agent(self):
         if self._agent is None:

vmas/simulator/dynamics/drone.py

@@ -49,7 +49,12 @@ def reset(self, index: Union[Tensor, int] = None):
                 device=self.world.device,
             )
         else:
-            self.drone_state[index] = 0.0
+            self.drone_state = vmas.TorchUtils.where_from_index(
+                index, 0.0, self.drone_state
+            )
+
+    def zero_grad(self):
+        self.drone_state = self.drone_state.detach()
 
     def needs_reset(self) -> Tensor:
         # Constraint roll and pitch within +-30 degrees

vmas/simulator/utils.py

@@ -218,6 +218,12 @@ def recursive_require_grad_(value: Union[Dict[str, Tensor], Tensor, List[Tensor]
             for val in value:
                 TorchUtils.recursive_require_grad_(val)
 
+    @staticmethod
+    def where_from_index(env_index, new_value, old_value):
+        mask = torch.zeros_like(old_value, dtype=torch.bool, device=old_value.device)
+        mask[env_index] = True
+        return torch.where(mask, new_value, old_value)
+
 
 class ScenarioUtils:
     @staticmethod