From 5e5eabb8a0245c526cc7528a2b17aa86e5f237fe Mon Sep 17 00:00:00 2001
From: StoneT2000 <stonezt2019@gmail.com>
Date: Sun, 21 Jan 2024 14:33:13 -0800
Subject: [PATCH] fix bug with to_tensor calls

---
 mani_skill2/utils/sapien_utils.py         | 11 +++--------
 mani_skill2/utils/structs/articulation.py |  6 +++---
 mani_skill2/utils/structs/base.py         |  4 ++--
 mani_skill2/utils/structs/joint.py        |  5 ++---
 4 files changed, 10 insertions(+), 16 deletions(-)

diff --git a/mani_skill2/utils/sapien_utils.py b/mani_skill2/utils/sapien_utils.py
index 49b722da2..2989f45af 100644
--- a/mani_skill2/utils/sapien_utils.py
+++ b/mani_skill2/utils/sapien_utils.py
@@ -20,10 +20,7 @@
 
 def to_tensor(array: Union[torch.Tensor, np.array, Sequence]):
     """
-    Maps any given sequence to the appropriate tensor for the appropriate backend
-
-    Note that all torch tensors are always moved to the GPU. There is generally no reason for them to ever be on the CPU as
-    GPU simulation is the only time torch is used in ManiSkill and thus all tensors from the simulation are on cuda devices
+    Maps any given sequence to a torch tensor on the CPU/GPU. If physx gpu is not enabled then we use CPU, otherwise GPU.
     """
     if get_backend_name() == "torch":
         if isinstance(array, np.ndarray):
@@ -37,11 +34,9 @@ def to_tensor(array: Union[torch.Tensor, np.array, Sequence]):
             return torch.Tensor(array).cuda()
     elif get_backend_name() == "numpy":
         if isinstance(array, np.ndarray):
-            return array
-        elif isinstance(array, torch.Tensor):
-            return array.cpu().numpy()
+            return torch.from_numpy(array)
         else:
-            return np.array(array)
+            return torch.tensor(array)
 
 
 def to_numpy(array: Union[Array, Sequence]):
diff --git a/mani_skill2/utils/structs/articulation.py b/mani_skill2/utils/structs/articulation.py
index 6cc6bab7b..94744965c 100644
--- a/mani_skill2/utils/structs/articulation.py
+++ b/mani_skill2/utils/structs/articulation.py
@@ -271,8 +271,8 @@ def qf(self):
 
     @qf.setter
     def qf(self, arg1):
-        arg1 = to_tensor(arg1)
         if physx.is_gpu_enabled():
+            arg1 = to_tensor(arg1)
             self.px.cuda_articulation_qf[self._data_index, : self.dof] = arg1
         else:
             self._objs[0].qf = arg1
@@ -294,8 +294,8 @@ def qpos(self):
 
     @qpos.setter
     def qpos(self, arg1):
-        arg1 = to_tensor(arg1)
         if physx.is_gpu_enabled():
+            arg1 = to_tensor(arg1)
             self.px.cuda_articulation_qpos[self._data_index, : self.dof] = arg1
         else:
             self._objs[0].qpos = arg1
@@ -309,8 +309,8 @@ def qvel(self):
 
     @qvel.setter
     def qvel(self, arg1):
-        arg1 = to_tensor(arg1)
         if physx.is_gpu_enabled():
+            arg1 = to_tensor(arg1)
             self.px.cuda_articulation_qvel[self._data_index, : self.dof] = arg1
         else:
             self._objs[0].qvel = arg1
diff --git a/mani_skill2/utils/structs/base.py b/mani_skill2/utils/structs/base.py
index 01fbfed3c..fd4490fcb 100644
--- a/mani_skill2/utils/structs/base.py
+++ b/mani_skill2/utils/structs/base.py
@@ -229,8 +229,8 @@ def angular_velocity(self) -> torch.Tensor:
 
     @angular_velocity.setter
     def angular_velocity(self, arg1: Array):
-        arg1 = to_tensor(arg1)
         if physx.is_gpu_enabled():
+            arg1 = to_tensor(arg1)
             self._body_data[self._body_data_index, 10:13] = arg1
         else:
             self._bodies[0].angular_velocity = arg1
@@ -292,8 +292,8 @@ def linear_velocity(self) -> torch.Tensor:
 
     @linear_velocity.setter
     def linear_velocity(self, arg1: Array):
-        arg1 = to_tensor(arg1)
         if physx.is_gpu_enabled():
+            arg1 = to_tensor(arg1)
             self._body_data[self._body_data_index, 7:10] = arg1
         else:
             self._bodies[0].linear_velocity = arg1
diff --git a/mani_skill2/utils/structs/joint.py b/mani_skill2/utils/structs/joint.py
index 47bacd8a0..65deaeded 100644
--- a/mani_skill2/utils/structs/joint.py
+++ b/mani_skill2/utils/structs/joint.py
@@ -187,8 +187,8 @@ def drive_target(self) -> torch.Tensor:
 
     @drive_target.setter
     def drive_target(self, arg1: Array) -> None:
-        arg1 = to_tensor(arg1)
         if physx.is_gpu_enabled():
+            arg1 = to_tensor(arg1)
             raise NotImplementedError(
                 "Setting drive targets of individual joints is not implemented yet."
             )
@@ -210,8 +210,8 @@ def drive_velocity_target(self) -> torch.Tensor:
 
     @drive_velocity_target.setter
     def drive_velocity_target(self, arg1: Array) -> None:
-        arg1 = to_tensor(arg1)
         if physx.is_gpu_enabled():
+            arg1 = to_tensor(arg1)
             raise NotImplementedError(
                 "Cannot set drive velocity targets at the moment in GPU simulation"
             )
@@ -248,7 +248,6 @@ def limits(self) -> torch.Tensor:
 
     @limits.setter
     def limits(self, arg1: Array) -> None:
-        arg1 = to_tensor(arg1)
         for joint in self._objs:
             joint.limits = arg1