[BugFix] Casting cached vals to device after device change

Author: matteobettini

vmas/scenarios/debug/circle_trajectory.py

@@ -5,7 +5,6 @@
 
 import torch
 from torch import Tensor
-
 from vmas import render_interactively
 from vmas.simulator.core import Agent, Sphere, World
 from vmas.simulator.scenario import BaseScenario
@@ -127,13 +126,21 @@ def get_tangent_to_circle(self, agent: Agent, closest_point=None):
             torch.linalg.vector_norm(agent.state.pos, dim=1) < self.desired_radius,
         )
 
-        rotated_vector = TorchUtils.rotate_vector(
-            distance_to_circle, torch.tensor(torch.pi / 2, device=self.world.device)
+        angle_90 = torch.tensor(torch.pi / 2, device=self.world.device).expand(
+            self.world.batch_dim
+        )
+
+        rotated_vector_90 = TorchUtils.rotate_vector(
+            distance_to_circle,
+            angle_90,
         )
-        rotated_vector[inside_circle] = TorchUtils.rotate_vector(
-            distance_to_circle[inside_circle],
-            torch.tensor(-torch.pi / 2, device=self.world.device),
+        rotated_vector_neg_90 = TorchUtils.rotate_vector(
+            distance_to_circle,
+            -angle_90,
         )
+        rotated_vector = rotated_vector_90
+        rotated_vector[inside_circle] = rotated_vector_neg_90[inside_circle]
+
         angle = rotated_vector / torch.linalg.vector_norm(
             rotated_vector, dim=1
         ).unsqueeze(-1)

vmas/simulator/core.py

@@ -11,7 +11,6 @@
 
 import torch
 from torch import Tensor
-
 from vmas.simulator.joints import Joint
 from vmas.simulator.physics import (
     _get_closest_point_line,
@@ -1047,9 +1046,6 @@ def __init__(
         ]
         # Map to save entity indexes
         self.entity_index_map = {}
-        self._normal_vector = torch.tensor(
-            [1.0, 0.0], dtype=torch.float32, device=self.device
-        ).repeat(self._batch_dim, 1)
 
     def add_agent(self, agent: Agent):
         """Only way to add agents to the world"""

vmas/simulator/joints.py

@@ -158,6 +158,9 @@ def _delta_anchor_tensor(self, entity):
                 .expand(entity.state.pos.shape)
             )
             self._delta_anchor_tensor_map[entity] = delta_anchor_tensor
+        self._delta_anchor_tensor_map[entity] = self._delta_anchor_tensor_map[
+            entity
+        ].to(entity.state.pos.device)
         return self._delta_anchor_tensor_map[entity]
 
     def get_delta_anchor(self, entity: vmas.simulator.core.Entity):