a

app.py

@@ -149,7 +149,7 @@ def create_demo():
         inputs = input_file
         outputs = output_file
         gr.Examples(
-            examples=[os.path.join(os.path.dirname(__file__), "./gradio_inter/20231104_017.pkl")],
+            examples=[os.path.join(os.path.dirname(__file__), "./data/102_grab_all_data.npy")],
             inputs=inputs,
             fn=predict,
             outputs=outputs,

models/dyn_model_act_v2.py

@@ -147,7 +147,7 @@ class Inertial:
         self.mass = mass
         self.inertia = inertia
         if torch.sum(self.inertia).item() < 1e-4:
-            self.inertia = self.inertia + torch.eye(3, dtype=torch.float32).cuda()
+            self.inertia = self.inertia + torch.eye(3, dtype=torch.float32)
         pass
     
 class Visual:
@@ -189,8 +189,8 @@ class Visual:
         vertices = mesh.vertices
         faces = mesh.faces
         
-        vertices = torch.from_numpy(vertices).float().cuda()
-        faces =torch.from_numpy(faces).long().cuda()
+        vertices = torch.from_numpy(vertices).float()
+        faces =torch.from_numpy(faces).long()
         
         vertices = vertices * self.geometry_mesh_scale.unsqueeze(0) + self.visual_xyz.unsqueeze(0)
         
@@ -244,7 +244,7 @@ class Visual:
                 # sample from the circile with cur_pts as thejcenter and the radius as expand_r
                 # (-r, r) # sample the offset vector in the size of (nn_expand_pts, 3)
                 offset_dist = Uniform(-1. * expand_r, expand_r)
-                offset_vec = offset_dist.sample((nn_expand_pts, 3)).cuda()
+                offset_vec = offset_dist.sample((nn_expand_pts, 3))
                 cur_expanded_pts = cur_pts + offset_vec
                 cur_expanded_visual_pts.append(cur_expanded_pts)
             cur_expanded_visual_pts = torch.cat(cur_expanded_visual_pts, dim=0)
@@ -252,7 +252,7 @@ class Visual:
         else:
             print(f"Loading visual pts from {expand_save_fn}") # load from the fn #
             cur_expanded_visual_pts = np.load(expand_save_fn, allow_pickle=True)
-            cur_expanded_visual_pts = torch.from_numpy(cur_expanded_visual_pts).float().cuda()
+            cur_expanded_visual_pts = torch.from_numpy(cur_expanded_visual_pts).float()
         self.cur_expanded_visual_pts = cur_expanded_visual_pts # expanded visual pts #
         return self.cur_expanded_visual_pts
 
@@ -343,7 +343,7 @@ class Link_urdf:
                 try:
                     cur_child_inertia = cur_child_struct.cur_inertia
                 except:
-                    cur_child_inertia = torch.eye(3, dtype=torch.float32).cuda()
+                    cur_child_inertia = torch.eye(3, dtype=torch.float32)
                 
                 
                 if cur_joint.type in ['revolute'] and (cur_joint_name not in ['WRJ2', 'WRJ1']):
@@ -428,10 +428,10 @@ class Link_urdf:
                         joint_link_idxes = torch.cat(joint_link_idxes, dim=-1) ### joint_link idxes ###
                         cur_joint_idx = cur_child.link_idx
                         joint_link_idxes = torch.cat(
-                            [joint_link_idxes, torch.tensor([cur_joint_idx], dtype=torch.long).cuda()], dim=-1
+                            [joint_link_idxes, torch.tensor([cur_joint_idx], dtype=torch.long)], dim=-1
                         )
                     else:
-                        joint_link_idxes = torch.tensor([cur_child.link_idx], dtype=torch.long).cuda().view(1,)
+                        joint_link_idxes = torch.tensor([cur_child.link_idx], dtype=torch.long).view(1,)
                     
                     # joint link idxes #
                     
@@ -446,7 +446,7 @@ class Link_urdf:
                     
             # joint_origin_xyz = self.joint.origin_xyz # c ## get forces from the expanded point set ## 
         else:
-            joint_origin_xyz = torch.tensor([0., 0., 0.], dtype=torch.float32).cuda()
+            joint_origin_xyz = torch.tensor([0., 0., 0.], dtype=torch.float32)
         # self.parent_rot_mtx = parent_rot
         # self.parent_trans_vec = parent_trans + joint_origin_xyz
         
@@ -455,13 +455,13 @@ class Link_urdf:
             # ## get init visual meshes ## ## --
             init_visual_meshes = self.visual.get_init_visual_meshes(parent_rot, parent_trans, init_visual_meshes, expanded_pts=expanded_pts)
             cur_visual_mesh_pts_nn = self.visual.vertices.size(0)
-            cur_link_idxes = torch.zeros((cur_visual_mesh_pts_nn, ), dtype=torch.long).cuda()+ self.link_idx
+            cur_link_idxes = torch.zeros((cur_visual_mesh_pts_nn, ), dtype=torch.long)+ self.link_idx
             init_visual_meshes['link_idxes'].append(cur_link_idxes)
             
             # self.link_idx # 
             if joint_idxes is not None:
                 cur_idxes = [self.link_idx for _ in range(cur_visual_mesh_pts_nn)]
-                cur_idxes = torch.tensor(cur_idxes, dtype=torch.long).cuda()
+                cur_idxes = torch.tensor(cur_idxes, dtype=torch.long)
                 joint_idxes.append(cur_idxes)
                 
         
@@ -473,7 +473,7 @@ class Link_urdf:
     # calculate inerti
     def calculate_inertia(self, link_name_to_visited, link_name_to_link_struct):
         link_name_to_visited[self.name] = 1
-        self.cur_inertia = torch.zeros((3, 3), dtype=torch.float32).cuda()
+        self.cur_inertia = torch.zeros((3, 3), dtype=torch.float32)
         
         if self.joint is not None:
             for joint_nm in self.joint:
@@ -634,7 +634,7 @@ class Link_urdf:
                 try:
                     cur_child_inertia = cur_child_struct.cur_inertia
                 except:
-                    cur_child_inertia = torch.eye(3, dtype=torch.float32).cuda()
+                    cur_child_inertia = torch.eye(3, dtype=torch.float32)
                 
                 
                 if cur_joint.type in ['revolute']:
@@ -663,18 +663,18 @@ class Link_urdf:
             self.joint_angle = init_states[self.joint.joint_idx]
             joint_axis = self.joint.axis
             self.rot_vec = self.joint_angle * joint_axis
-            self.joint.state = torch.tensor([1, 0, 0, 0], dtype=torch.float32).cuda()
+            self.joint.state = torch.tensor([1, 0, 0, 0], dtype=torch.float32)
             self.joint.state = self.joint.state + update_quaternion(self.rot_vec, self.joint.state)
             self.joint.timestep_to_states[0] = self.joint.state.detach()
-            self.joint.timestep_to_vels[0] = torch.zeros((3,), dtype=torch.float32).cuda().detach() ## velocity ##
+            self.joint.timestep_to_vels[0] = torch.zeros((3,), dtype=torch.float32).detach() ## velocity ##
         for cur_link in self.children:
             cur_link.set_init_states_target_value(init_states)
     
     # should forward for one single step -> use the action #
     def set_init_states(self, ):
-        self.joint.state = torch.tensor([1, 0, 0, 0], dtype=torch.float32).cuda()
+        self.joint.state = torch.tensor([1, 0, 0, 0], dtype=torch.float32)
         self.joint.timestep_to_states[0] = self.joint.state.detach()
-        self.joint.timestep_to_vels[0] = torch.zeros((3,), dtype=torch.float32).cuda().detach() ## velocity ##
+        self.joint.timestep_to_vels[0] = torch.zeros((3,), dtype=torch.float32).detach() ## velocity ##
         for cur_link in self.children:
             cur_link.set_init_states() 
     
@@ -737,31 +737,31 @@ class Joint_urdf: #
         
         #### only for the current state #### # joint urdf #
         self.state = nn.Parameter(
-            torch.tensor([1., 0., 0., 0.], dtype=torch.float32, requires_grad=True).cuda(), requires_grad=True
+            torch.tensor([1., 0., 0., 0.], dtype=torch.float32, requires_grad=True), requires_grad=True
         )
         self.action = nn.Parameter(
-            torch.zeros((1,), dtype=torch.float32, requires_grad=True).cuda(), requires_grad=True
+            torch.zeros((1,), dtype=torch.float32, requires_grad=True), requires_grad=True
         )
         # self.rot_mtx = np.eye(3, dtypes=np.float32)
         # self.trans_vec = np.zeros((3,), dtype=np.float32) ## rot m
-        self.rot_mtx = nn.Parameter(torch.eye(n=3, dtype=torch.float32, requires_grad=True).cuda(), requires_grad=True)
-        self.trans_vec  = nn.Parameter(torch.zeros((3,), dtype=torch.float32,  requires_grad=True).cuda(), requires_grad=True)
+        self.rot_mtx = nn.Parameter(torch.eye(n=3, dtype=torch.float32, requires_grad=True), requires_grad=True)
+        self.trans_vec  = nn.Parameter(torch.zeros((3,), dtype=torch.float32,  requires_grad=True), requires_grad=True)
         
     def set_initial_state(self, state):
         # joint angle as the state value #
-        self.timestep_to_vels[0] = torch.zeros((3,), dtype=torch.float32).cuda().detach() ## velocity ##
+        self.timestep_to_vels[0] = torch.zeros((3,), dtype=torch.float32).detach() ## velocity ##
         delta_rot_vec = self.axis_xyz * state
         # self.timestep_to_states[0] = state.detach()
-        cur_state = torch.tensor([1., 0., 0., 0.], dtype=torch.float32).cuda()
+        cur_state = torch.tensor([1., 0., 0., 0.], dtype=torch.float32)
         init_state = cur_state + update_quaternion(delta_rot_vec, cur_state)
         self.timestep_to_states[0] = init_state.detach()
         self.state = init_state
         
     def set_delta_state_and_update(self, state, cur_timestep):
-        self.timestep_to_vels[cur_timestep] = torch.zeros((3,), dtype=torch.float32).cuda().detach() 
+        self.timestep_to_vels[cur_timestep] = torch.zeros((3,), dtype=torch.float32).detach() 
         delta_rot_vec = self.axis_xyz * state
         if cur_timestep == 0:
-            prev_state = torch.tensor([1., 0., 0., 0.], dtype=torch.float32).cuda()
+            prev_state = torch.tensor([1., 0., 0., 0.], dtype=torch.float32)
         else:
             # prev_state = self.timestep_to_states[cur_timestep - 1].detach()
             prev_state = self.timestep_to_states[cur_timestep - 1] # .detach() # not detach? #
@@ -771,7 +771,7 @@ class Joint_urdf: #
     
     
     def set_delta_state_and_update_v2(self, delta_state, cur_timestep):
-        self.timestep_to_vels[cur_timestep] = torch.zeros((3,), dtype=torch.float32).cuda().detach() 
+        self.timestep_to_vels[cur_timestep] = torch.zeros((3,), dtype=torch.float32).detach() 
         
         if cur_timestep == 0:
             cur_state = delta_state
@@ -787,12 +787,12 @@ class Joint_urdf: #
         
         cur_rot_vec = self.axis_xyz * cur_state ### cur_state #### # 
         # angle to the quaternion ? # 
-        init_state = torch.tensor([1., 0., 0., 0.], dtype=torch.float32).cuda()
+        init_state = torch.tensor([1., 0., 0., 0.], dtype=torch.float32)
         cur_quat_state = init_state + update_quaternion(cur_rot_vec, init_state)
         self.state = cur_quat_state
         
         # if cur_timestep == 0:
-        #     prev_state = torch.tensor([1., 0., 0., 0.], dtype=torch.float32).cuda()
+        #     prev_state = torch.tensor([1., 0., 0., 0.], dtype=torch.float32)
         # else:
         #     # prev_state = self.timestep_to_states[cur_timestep - 1].detach()
         #     prev_state = self.timestep_to_states[cur_timestep - 1] # .detach() # not detach? #
@@ -816,8 +816,8 @@ class Joint_urdf: #
             self.rot_mtx = rot_mtx
             self.trans_vec = trans_vec
         elif self.type == "fixed":
-            rot_mtx = torch.eye(3, dtype=torch.float32).cuda()
-            trans_vec = torch.zeros((3,), dtype=torch.float32).cuda()
+            rot_mtx = torch.eye(3, dtype=torch.float32)
+            trans_vec = torch.zeros((3,), dtype=torch.float32)
             # trans_vec = self.origin_xyz
             self.rot_mtx = rot_mtx
             self.trans_vec = trans_vec # 
@@ -840,7 +840,7 @@ class Joint_urdf: #
             torque = self.action * self.axis_xyz
             
             # # Compute inertia matrix #
-            # inertial = torch.zeros((3, 3), dtype=torch.float32).cuda()
+            # inertial = torch.zeros((3, 3), dtype=torch.float32)
             # for i_pts in range(self.visual_pts.size(0)):
             #     cur_pts = self.visual_pts[i_pts]
             #     cur_pts_mass = self.visual_pts_mass[i_pts]
@@ -848,7 +848,7 @@ class Joint_urdf: #
             #     # cur_vert = init_passive_mesh[i_v]
             #     # cur_r = cur_vert - init_passive_mesh_center
             #     dot_r_r = torch.sum(cur_r * cur_r)
-            #     cur_eye_mtx = torch.eye(3, dtype=torch.float32).cuda()
+            #     cur_eye_mtx = torch.eye(3, dtype=torch.float32)
             #     r_mult_rT = torch.matmul(cur_r.unsqueeze(-1), cur_r.unsqueeze(0))
             #     inertial += (dot_r_r * cur_eye_mtx - r_mult_rT) * cur_pts_mass
             # m = torch.sum(self.visual_pts_mass)
@@ -935,7 +935,7 @@ class Joint_urdf: #
             torque = self.action * self.axis_xyz
             
             # # Compute inertia matrix #
-            # inertial = torch.zeros((3, 3), dtype=torch.float32).cuda()
+            # inertial = torch.zeros((3, 3), dtype=torch.float32)
             # for i_pts in range(self.visual_pts.size(0)):
             #     cur_pts = self.visual_pts[i_pts]
             #     cur_pts_mass = self.visual_pts_mass[i_pts]
@@ -943,7 +943,7 @@ class Joint_urdf: #
             #     # cur_vert = init_passive_mesh[i_v]
             #     # cur_r = cur_vert - init_passive_mesh_center
             #     dot_r_r = torch.sum(cur_r * cur_r)
-            #     cur_eye_mtx = torch.eye(3, dtype=torch.float32).cuda()
+            #     cur_eye_mtx = torch.eye(3, dtype=torch.float32)
             #     r_mult_rT = torch.matmul(cur_r.unsqueeze(-1), cur_r.unsqueeze(0))
             #     inertial += (dot_r_r * cur_eye_mtx - r_mult_rT) * cur_pts_mass
             # m = torch.sum(self.visual_pts_mass)
@@ -957,7 +957,7 @@ class Joint_urdf: #
             
             # inertia_inv = torch.linalg.inv(cur_inertia).detach()
             
-            inertia_inv = torch.eye(n=3, dtype=torch.float32).cuda()
+            inertia_inv = torch.eye(n=3, dtype=torch.float32)
             
             
             
@@ -994,14 +994,14 @@ class Joint_urdf: #
             
             # if cur_timestep 
             if cur_timestep == 0:
-                self.timestep_to_states[cur_timestep] = torch.zeros((1,), dtype=torch.float32).cuda()
+                self.timestep_to_states[cur_timestep] = torch.zeros((1,), dtype=torch.float32)
             cur_state = self.timestep_to_states[cur_timestep].detach()
             nex_state  = cur_state + delta_state
             # nex_state = nex_state + penetration_delta_state
             ## state rot vector along axis ## ## get the pentrated froces -- calulaterot qj
             state_rot_vec_along_axis = nex_state * self.axis_xyz
             ### state in the rotation vector -> state in quaternion ###
-            state_rot_quat = torch.tensor([1., 0., 0., 0.], dtype=torch.float32).cuda() + update_quaternion(state_rot_vec_along_axis, torch.tensor([1., 0., 0., 0.], dtype=torch.float32).cuda())
+            state_rot_quat = torch.tensor([1., 0., 0., 0.], dtype=torch.float32) + update_quaternion(state_rot_vec_along_axis, torch.tensor([1., 0., 0., 0.], dtype=torch.float32))
             ### state 
             self.state = state_rot_quat
             ### get states? ##
@@ -1060,7 +1060,7 @@ class Joint_urdf: #
                     # penetration_delta_state = forces_torques_dot_axis
                 else:
                     penetration_delta_state = 0.0
-                    cur_joint_maximal_forces = torch.zeros((6,), dtype=torch.float32).cuda()
+                    cur_joint_maximal_forces = torch.zeros((6,), dtype=torch.float32)
                     cur_joint_idx = joint_idx
                     joint_penetration_forces[cur_joint_idx][:] = cur_joint_maximal_forces[:].clone()
                     
@@ -1095,7 +1095,7 @@ class Robot_urdf:
         # # 
         self.act_joint_idxes = list(self.actions_joint_name_to_joint_idx.values())
         self.act_joint_idxes = sorted(self.act_joint_idxes, reverse=False)
-        self.act_joint_idxes = torch.tensor(self.act_joint_idxes, dtype=torch.long).cuda()[2:]
+        self.act_joint_idxes = torch.tensor(self.act_joint_idxes, dtype=torch.long)[2:]
         
         self.real_actions_joint_name_to_joint_idx = real_actions_joint_name_to_joint_idx
         
@@ -1136,8 +1136,8 @@ class Robot_urdf:
         init_visual_meshes = {
             'vertices': [], 'faces': [], 'link_idxes': [], 'transformed_joint_pos': [], 'link_idxes': [],  'transformed_joint_pos': [], 'joint_link_idxes': []
         }
-        init_parent_rot = torch.eye(3, dtype=torch.float32).cuda()
-        init_parent_trans = torch.zeros((3,), dtype=torch.float32).cuda()
+        init_parent_rot = torch.eye(3, dtype=torch.float32)
+        init_parent_trans = torch.zeros((3,), dtype=torch.float32)
         
         palm_idx = self.link_name_to_link_idxes["palm"]
         palm_link = self.links[palm_idx]
@@ -1174,8 +1174,8 @@ class Robot_urdf:
         action_joint_name_to_joint_idx = self.real_actions_joint_name_to_joint_idx
         # print(f"action_joint_name_to_joint_idx: {action_joint_name_to_joint_idx}")
         
-        parent_rot = torch.eye(3, dtype=torch.float32).cuda()
-        parent_trans = torch.zeros((3,), dtype=torch.float32).cuda()
+        parent_rot = torch.eye(3, dtype=torch.float32)
+        parent_trans = torch.zeros((3,), dtype=torch.float32)
         
         # cur_timestep, time_cons, action_joint_name_to_joint_idx, link_name_to_visited, link_name_to_link_struct, parent_rot, parent_trans,  penetration_forces, sampled_visual_pts_joint_idxes, joint_penetration_forces):
         
@@ -1251,8 +1251,8 @@ class Robot_urdf:
         
         link_name_to_visited = {}
         
-        # parent_rot = torch.eye(3, dtype=torch.float32).cuda()
-        # parent_trans = torch.zeros((3,), dtype=torch.float32).cuda()
+        # parent_rot = torch.eye(3, dtype=torch.float32)
+        # parent_trans = torch.zeros((3,), dtype=torch.float32)
         ## set actions ## #
         # set_actions_and_update_states_v2(self, action, cur_timestep, time_cons, cur_inertia):
         # self, actions, cur_timestep, time_cons, action_joint_name_to_joint_idx, link_name_to_visited, link_name_to_link_struct ## set and update states ##
@@ -1271,8 +1271,8 @@ class Robot_urdf:
         
         link_name_to_visited = {}
         
-        parent_rot = torch.eye(3, dtype=torch.float32).cuda()
-        parent_trans = torch.zeros((3,), dtype=torch.float32).cuda()
+        parent_rot = torch.eye(3, dtype=torch.float32)
+        parent_trans = torch.zeros((3,), dtype=torch.float32)
         ## set actions ## #
         # set_actions_and_update_states_v2(self, action, cur_timestep, time_cons, cur_inertia):
         # self, actions, cur_timestep, time_cons, action_joint_name_to_joint_idx, link_name_to_visited, link_name_to_link_struct ## set and update states ##
@@ -1363,7 +1363,7 @@ def parse_nparray_from_string(strr, args=None):
     vals = np.array(vals, dtype=np.float32)
     vals = torch.from_numpy(vals).float()
     ## vals ##
-    vals = nn.Parameter(vals.cuda(), requires_grad=True)
+    vals = nn.Parameter(vals, requires_grad=True)
     
     return vals
 
@@ -1486,7 +1486,7 @@ def parse_link_data_urdf(link):
     inertial_izz = inertial_inertia.attrib["izz"]
     inertial_izz = float(inertial_izz)
     
-    inertial_inertia_mtx = torch.zeros((3, 3), dtype=torch.float32).cuda()
+    inertial_inertia_mtx = torch.zeros((3, 3), dtype=torch.float32)
     inertial_inertia_mtx[0, 0] = inertial_ixx
     inertial_inertia_mtx[0, 1] = inertial_ixy
     inertial_inertia_mtx[0, 2] = inertial_ixz
@@ -1547,7 +1547,7 @@ def parse_joint_data_urdf(joint):
         origin_xyz_string = joint_origin.attrib["xyz"]
         origin_xyz = parse_nparray_from_string(origin_xyz_string)
     except:
-        origin_xyz = torch.tensor([0., 0., 0.], dtype=torch.float32).cuda()
+        origin_xyz = torch.tensor([0., 0., 0.], dtype=torch.float32)
         origin_xyz_string = ""
         
     joint_axis = joint.find("./axis")
@@ -1555,7 +1555,7 @@ def parse_joint_data_urdf(joint):
         joint_axis = joint_axis.attrib["xyz"]
         joint_axis = parse_nparray_from_string(joint_axis)
     else:
-        joint_axis = torch.tensor([1, 0., 0.], dtype=torch.float32).cuda()
+        joint_axis = torch.tensor([1, 0., 0.], dtype=torch.float32)
     
     joint_limit = joint.find("./limit")
     if joint_limit is not None:
@@ -1746,13 +1746,13 @@ class RobotAgent: # robot and the robot #
         
         self.time_constant = nn.Embedding(
             num_embeddings=3, embedding_dim=1
-        ).cuda()
+        )
         torch.nn.init.ones_(self.time_constant.weight) #
         self.time_constant.weight.data = self.time_constant.weight.data * 0.2 ### time_constant data #
         
         self.optimizable_actions = nn.Embedding(
             num_embeddings=100, embedding_dim=60,
-        ).cuda()
+        )
         torch.nn.init.zeros_(self.optimizable_actions.weight) #
         
         self.learning_rate = 5e-4
@@ -1770,24 +1770,24 @@ class RobotAgent: # robot and the robot #
         
         
     def set_init_states_target_value(self, init_states):
-        # glb_rot = torch.eye(n=3, dtype=torch.float32).cuda()
-        # glb_trans = torch.zeros((3,), dtype=torch.float32).cuda() ### glb_trans #### and the rot 3##
+        # glb_rot = torch.eye(n=3, dtype=torch.float32)
+        # glb_trans = torch.zeros((3,), dtype=torch.float32) ### glb_trans #### and the rot 3##
 
         # tot_init_states = {}
         # tot_init_states['glb_rot'] = glb_rot;
         # tot_init_states['glb_trans'] = glb_trans; 
         # tot_init_states['links_init_states'] = init_states
         # self.active_robot.set_init_states_target_value(tot_init_states)
-        # init_joint_states = torch.zeros((60, ), dtype=torch.float32).cuda()
+        # init_joint_states = torch.zeros((60, ), dtype=torch.float32)
         self.active_robot.set_initial_state(init_states)
     
     def set_init_states(self):
-        # glb_rot = torch.eye(n=3, dtype=torch.float32).cuda()
-        # glb_trans = torch.zeros((3,), dtype=torch.float32).cuda() ### glb_trans #### and the rot 3##
+        # glb_rot = torch.eye(n=3, dtype=torch.float32)
+        # glb_trans = torch.zeros((3,), dtype=torch.float32) ### glb_trans #### and the rot 3##
         
         # ### random rotation ###
         # # glb_rot_np = R.random().as_matrix()
-        # # glb_rot = torch.from_numpy(glb_rot_np).float().cuda()
+        # # glb_rot = torch.from_numpy(glb_rot_np).float()
         # ### random rotation ###
         
         # # glb_rot, glb_trans #
@@ -1796,7 +1796,7 @@ class RobotAgent: # robot and the robot #
         # init_states['glb_trans'] = glb_trans; 
         # self.active_robot.set_init_states(init_states)
         
-        init_joint_states = torch.zeros((60, ), dtype=torch.float32).cuda()
+        init_joint_states = torch.zeros((60, ), dtype=torch.float32)
         self.active_robot.set_initial_state(init_joint_states)
     
     # cur_verts, joint_idxes =  get_init_state_visual_pts(expanded_pts=False, ret_joint_idxes=True)
@@ -1821,13 +1821,13 @@ class RobotAgent: # robot and the robot #
     
     def set_actions_and_update_states(self, actions, cur_timestep):
         # 
-        time_cons = self.time_constant(torch.zeros((1,), dtype=torch.long).cuda()) ### time constant of the system ##
+        time_cons = self.time_constant(torch.zeros((1,), dtype=torch.long)) ### time constant of the system ##
         self.active_robot.set_actions_and_update_states(actions, cur_timestep, time_cons) ### 
         pass
     
     def set_actions_and_update_states_v2(self, actions, cur_timestep, penetration_forces=None, sampled_visual_pts_joint_idxes=None):
         # 
-        time_cons = self.time_constant(torch.zeros((1,), dtype=torch.long).cuda()) ### time constant of the system ##
+        time_cons = self.time_constant(torch.zeros((1,), dtype=torch.long)) ### time constant of the system ##
         self.active_robot.set_actions_and_update_states_v2(actions, cur_timestep, time_cons, penetration_forces=penetration_forces, sampled_visual_pts_joint_idxes=sampled_visual_pts_joint_idxes) ### 
         pass
     
@@ -1841,9 +1841,9 @@ class RobotAgent: # robot and the robot #
         timestep_to_visual_pts = {}
         for i_step in range(50):
             actions = {}
-            actions['delta_glb_rot']  = torch.eye(3, dtype=torch.float32).cuda()
-            actions['delta_glb_trans'] = torch.zeros((3,), dtype=torch.float32).cuda()
-            actions_link_actions = torch.ones((22, ), dtype=torch.float32).cuda()
+            actions['delta_glb_rot']  = torch.eye(3, dtype=torch.float32)
+            actions['delta_glb_trans'] = torch.zeros((3,), dtype=torch.float32)
+            actions_link_actions = torch.ones((22, ), dtype=torch.float32)
             # actions_link_actions = actions_link_actions * 0.2
             actions_link_actions = actions_link_actions * -1. # 
             actions['link_actions'] = actions_link_actions
@@ -1862,7 +1862,7 @@ class RobotAgent: # robot and the robot #
         # TODO: load reference points #
         self.ts_to_reference_pts = np.load(reference_pts_dict, allow_pickle=True).item() #### 
         self.ts_to_reference_pts = {
-            ts // 2 + 1: torch.from_numpy(self.ts_to_reference_pts[ts]).float().cuda() for ts in self.ts_to_reference_pts
+            ts // 2 + 1: torch.from_numpy(self.ts_to_reference_pts[ts]).float() for ts in self.ts_to_reference_pts
         }
     
     
@@ -1879,9 +1879,9 @@ class RobotAgent: # robot and the robot #
             for cur_ts in range(self.n_timesteps):
                 # print(f"iter: {i_iter}, cur_ts: {cur_ts}")
                 # actions = {}
-                # actions['delta_glb_rot']  = torch.eye(3, dtype=torch.float32).cuda()
-                # actions['delta_glb_trans'] = torch.zeros((3,), dtype=torch.float32).cuda()
-                actions_link_actions = self.optimizable_actions(torch.zeros((1,), dtype=torch.long).cuda() + cur_ts).squeeze(0)
+                # actions['delta_glb_rot']  = torch.eye(3, dtype=torch.float32)
+                # actions['delta_glb_trans'] = torch.zeros((3,), dtype=torch.float32)
+                actions_link_actions = self.optimizable_actions(torch.zeros((1,), dtype=torch.long) + cur_ts).squeeze(0)
                 # actions_link_actions = actions_link_actions * 0.2
                 # actions_link_actions = actions_link_actions * -1. # 
                 # actions['link_actions'] = actions_link_actions
@@ -2240,39 +2240,39 @@ def get_shadow_GT_states_data_from_ckpt(ckpt_fn):
     # robot actions # # 
     # robot_actions = nn.Embedding(
     #     num_embeddings=num_steps, embedding_dim=22,
-    # ).cuda()
+    # )
     # torch.nn.init.zeros_(robot_actions.weight)
     # # params_to_train += list(robot_actions.parameters())
     
     # robot_delta_states = nn.Embedding(
     #     num_embeddings=num_steps, embedding_dim=60,
-    # ).cuda()
+    # )
     # torch.nn.init.zeros_(robot_delta_states.weight)
     # # params_to_train += list(robot_delta_states.parameters())
     
     
     robot_states = nn.Embedding(
         num_embeddings=num_steps, embedding_dim=60,
-    ).cuda()
+    )
     torch.nn.init.zeros_(robot_states.weight)
     # params_to_train += list(robot_states.parameters())
     
     # robot_init_states = nn.Embedding(
     #     num_embeddings=1, embedding_dim=22,
-    # ).cuda()
+    # )
     # torch.nn.init.zeros_(robot_init_states.weight)
     # # params_to_train += list(robot_init_states.parameters())
     
     robot_glb_rotation = nn.Embedding(
         num_embeddings=num_steps, embedding_dim=4
-    ).cuda()
+    )
     robot_glb_rotation.weight.data[:, 0] = 1.
     robot_glb_rotation.weight.data[:, 1:] = 0.
     
     
     robot_glb_trans = nn.Embedding(
         num_embeddings=num_steps, embedding_dim=3
-    ).cuda()
+    )
     torch.nn.init.zeros_(robot_glb_trans.weight)
         
     ''' Load optimized MANO hand actions and states '''
@@ -2531,7 +2531,7 @@ if __name__=='__main__': # # #
     tot_transformed_pts = []
     for i_ts in range(len(blended_states)):
         cur_blended_states = blended_states[i_ts]
-        cur_blended_states = torch.from_numpy(cur_blended_states).float().cuda()
+        cur_blended_states = torch.from_numpy(cur_blended_states).float()
         robot_agent.active_robot.set_delta_state_and_update_v2(cur_blended_states, 0)
         cur_pts = robot_agent.get_init_state_visual_pts().detach().cpu().numpy()
         
@@ -2568,7 +2568,7 @@ if __name__=='__main__': # # #
     shadow_hand_sv_fn = "/home/xueyi/diffsim/NeuS/raw_data/scaled_redmax_hand_rescaled_grab.obj"
     shadow_hand_mesh.export(shadow_hand_sv_fn)
     
-    init_joint_states = torch.randn((60, ), dtype=torch.float32).cuda()
+    init_joint_states = torch.randn((60, ), dtype=torch.float32)
     robot_agent.set_initial_state(init_joint_states)
     
     
@@ -2676,15 +2676,15 @@ if __name__=='__main__': # # #
     # mesh_obj.export(f"hand_urdf.ply")
     
     ##### Test the set initial state function #####
-    init_joint_states = torch.zeros((60, ), dtype=torch.float32).cuda()
+    init_joint_states = torch.zeros((60, ), dtype=torch.float32)
     cur_robot.set_initial_state(init_joint_states)
     ##### Test the set initial state function #####
     
     
     
     
-    cur_zeros_actions = torch.zeros((60, ), dtype=torch.float32).cuda()
-    cur_ones_actions = torch.ones((60, ), dtype=torch.float32).cuda() #  * 100
+    cur_zeros_actions = torch.zeros((60, ), dtype=torch.float32)
+    cur_ones_actions = torch.ones((60, ), dtype=torch.float32) #  * 100
     
     ts_to_mesh_verts = {}
     for i_ts in range(50):

pre-requirements.txt

@@ -1,2 +1,2 @@
 pip==23.3.2
-torch==2.2.0
+# torch==2.2.0

requirements.txt

@@ -1,5 +1,6 @@
 -f https://download.pytorch.org/whl/cpu/torch_stable.html
--f https://data.pyg.org/whl/torch-2.2.0%2Bcpu.html
+# -f https://data.pyg.org/whl/torch-2.2.0%2Bcpu.html
+-i https://download.pytorch.org/whl/cpu
 # pip==20.2.4
 torch==2.2.0
 # torchvision==0.13.1

scripts_demo/train_grab_pointset_points_dyn_s1.sh

@@ -18,11 +18,15 @@ export conf=dyn_grab_pointset_points_dyn_s1.conf
 export conf_root="./confs_new"
 
 
-export data_path="./data/102_grab_all_data.npy"
-# bash scripts_new/train_grab_pointset_points_dyn_s1.sh
+export data_fn="./data/102_grab_all_data.npy"
+
+
+# bash scripts_demo/train_grab_pointset_points_dyn_s1.sh
 
 export cuda_ids="0"
 
 
-CUDA_VISIBLE_DEVICES=${cuda_ids} python ${trainer} --mode ${mode} --conf ${conf_root}/${conf} --case ${data_case} --data_path=${data_path}
+# CUDA_VISIBLE_DEVICES=${cuda_ids} 
+
+python ${trainer} --mode ${mode} --conf ${conf_root}/${conf} --case ${data_case} --data_fn=${data_fn}