models/dyn_model_act_v2.py

import math
# import torch
# from ..utils import Timer
import numpy as np
# import torch.nn.functional as F
import os

import argparse

from xml.etree.ElementTree import ElementTree

import trimesh
import torch
import torch.nn as nn
# import List
# class link; joint; body
### 

from scipy.spatial.transform import Rotation as R
from torch.distributions.uniform import Uniform
# deformable articulated objects with the articulated models #

DAMPING = 1.0
DAMPING = 0.3

urdf_fn = ""

def plane_rotation_matrix_from_angle_xz(angle):
    sin_ = torch.sin(angle)
    cos_ = torch.cos(angle)
    zero_padding = torch.zeros_like(cos_)
    one_padding = torch.ones_like(cos_)
    col_a = torch.stack(
        [cos_, zero_padding, sin_], dim=0
    )
    col_b = torch.stack(
        [zero_padding, one_padding, zero_padding], dim=0
    )
    col_c = torch.stack(
        [-1. * sin_, zero_padding, cos_], dim=0
    )
    rot_mtx = torch.stack(
        [col_a, col_b, col_c], dim=-1
    )
    return rot_mtx

def plane_rotation_matrix_from_angle(angle):
    ## angle of 
    sin_ = torch.sin(angle)
    cos_ = torch.cos(angle)
    col_a = torch.stack(
        [cos_, sin_], dim=0 ### col of the rotation matrix
    )
    col_b = torch.stack(
        [-1. * sin_, cos_], dim=0 ## cols of the rotation matrix
    )
    rot_mtx = torch.stack(
        [col_a, col_b], dim=-1 ### rotation matrix
    )
    return rot_mtx

def rotation_matrix_from_axis_angle(axis, angle): # rotation_matrix_from_axis_angle -> 
        # sin_ = np.sin(angle) #  ti.math.sin(angle)
        # cos_ = np.cos(angle) #  ti.math.cos(angle)
        sin_ = torch.sin(angle) #  ti.math.sin(angle)
        cos_ = torch.cos(angle) #  ti.math.cos(angle)
        u_x, u_y, u_z = axis[0], axis[1], axis[2]
        u_xx = u_x * u_x
        u_yy = u_y * u_y
        u_zz = u_z * u_z
        u_xy = u_x * u_y
        u_xz = u_x * u_z
        u_yz = u_y * u_z
        
        row_a = torch.stack(
            [cos_ + u_xx * (1 - cos_), u_xy * (1. - cos_) + u_z * sin_, u_xz * (1. - cos_) - u_y * sin_], dim=0
        )
        # print(f"row_a: {row_a.size()}")
        row_b = torch.stack(
            [u_xy * (1. - cos_) - u_z * sin_, cos_ + u_yy * (1. - cos_), u_yz * (1. - cos_) + u_x * sin_], dim=0
        )
        # print(f"row_b: {row_b.size()}")
        row_c = torch.stack(
            [u_xz * (1. - cos_) + u_y * sin_, u_yz * (1. - cos_) - u_x * sin_, cos_ + u_zz * (1. - cos_)], dim=0
        )
        # print(f"row_c: {row_c.size()}")
        
        ### rot_mtx for the rot_mtx ###
        rot_mtx = torch.stack(
            [row_a, row_b, row_c], dim=-1 ### rot_matrix of he matrix ##
        )
        
        return rot_mtx


def update_quaternion(delta_angle, prev_quat):
    s1 = 0
    s2 = prev_quat[0]
    v2 = prev_quat[1:]
    v1 = delta_angle / 2
    new_v = s1 * v2 + s2 * v1 + torch.cross(v1, v2)
    new_s = s1 * s2 - torch.sum(v1 * v2)
    new_quat = torch.cat([new_s.unsqueeze(0), new_v], dim=0)
    return new_quat


def quaternion_to_matrix(quaternions: torch.Tensor) -> torch.Tensor:
    """
    Convert rotations given as quaternions to rotation matrices.

    Args:
        quaternions: quaternions with real part first,
            as tensor of shape (..., 4).

    Returns:
        Rotation matrices as tensor of shape (..., 3, 3).
    """
    r, i, j, k = torch.unbind(quaternions, -1) # -1 for the quaternion matrix # 
    # pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
    two_s = 2.0 / (quaternions * quaternions).sum(-1)

    o = torch.stack(
        (
            1 - two_s * (j * j + k * k),
            two_s * (i * j - k * r),
            two_s * (i * k + j * r),
            two_s * (i * j + k * r),
            1 - two_s * (i * i + k * k),
            two_s * (j * k - i * r),
            two_s * (i * k - j * r),
            two_s * (j * k + i * r),
            1 - two_s * (i * i + j * j),
        ),
        -1,
    )
    
    return o.reshape(quaternions.shape[:-1] + (3, 3))

    


class Inertial:
    def __init__(self, origin_rpy, origin_xyz, mass, inertia) -> None:
        self.origin_rpy = origin_rpy
        self.origin_xyz = origin_xyz
        self.mass = mass
        self.inertia = inertia
        if torch.sum(self.inertia).item() < 1e-4:
            self.inertia = self.inertia + torch.eye(3, dtype=torch.float32)
        pass
    
class Visual:
    def __init__(self, visual_xyz, visual_rpy, geometry_mesh_fn, geometry_mesh_scale) -> None:
        # self.visual_origin = visual_origin
        self.visual_xyz = visual_xyz
        self.visual_rpy = visual_rpy
        self.mesh_nm = geometry_mesh_fn.split("/")[-1].split(".")[0]
        mesh_root = "./rsc/mano"
        # if not os.path.exists(mesh_root):
        #     mesh_root = "/data/xueyi/diffsim/NeuS/rsc/mano"
        # if "shadow" in urdf_fn and "left" in urdf_fn:
        #     mesh_root = "/home/xueyi/diffsim/NeuS/rsc/shadow_hand_description_left"
        #     if not os.path.exists(mesh_root):
        #         mesh_root = "/root/diffsim/quasi-dyn/rsc/shadow_hand_description_left"
        # elif "shadow" in urdf_fn:
        #     mesh_root = "/home/xueyi/diffsim/NeuS/rsc/shadow_hand_description"
        #     if not os.path.exists(mesh_root):
        #         mesh_root = "/root/diffsim/quasi-dyn/rsc/shadow_hand_description"
        # elif "redmax" in urdf_fn:
        #     mesh_root = "/home/xueyi/diffsim/NeuS/rsc/redmax_hand"
        #     if not os.path.exists(mesh_root):
        #         mesh_root = "/root/diffsim/quasi-dyn/rsc/redmax_hand"
            
        self.mesh_root = mesh_root
        geometry_mesh_fn = geometry_mesh_fn.replace(".dae", ".obj")
        self.geometry_mesh_fn = os.path.join(mesh_root, geometry_mesh_fn)
        
        self.geometry_mesh_scale = geometry_mesh_scale
        # tranformed by xyz #
        self.vertices, self.faces = self.load_geoemtry_mesh()
        self.cur_expanded_visual_pts = None
        pass
    
    def load_geoemtry_mesh(self, ):
        # mesh_root = 
        # if self.geometry_mesh_fn.end
        mesh = trimesh.load_mesh(self.geometry_mesh_fn)
        vertices = mesh.vertices
        faces = mesh.faces
        
        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)
        
        return vertices, faces
    
    # init_visual_meshes = get_init_visual_meshes(self, parent_rot, parent_trans, init_visual_meshes)
    def get_init_visual_meshes(self, parent_rot, parent_trans, init_visual_meshes, expanded_pts=False):
        # cur_vertices = torch.matmul(parent_rot, self.vertices.transpose(1, 0)).contiguous().transpose(1, 0).contiguous() + parent_trans.unsqueeze(0)
        
        if not expanded_pts:
            cur_vertices = self.vertices
            # print(f"adding mesh loaded from {self.geometry_mesh_fn}")
            init_visual_meshes['vertices'].append(cur_vertices) # cur vertices # trans # 
            init_visual_meshes['faces'].append(self.faces)
        else:
            ## expanded visual meshes ##
            cur_vertices = self.cur_expanded_visual_pts
            init_visual_meshes['vertices'].append(cur_vertices) 
            init_visual_meshes['faces'].append(self.faces)
        return init_visual_meshes

    def expand_visual_pts(self, ):
        # expand_factor = 0.2
        # nn_expand_pts = 20
        
        # expand_factor = 0.4
        # nn_expand_pts = 40 ### number of the expanded points ### ## points ##
        
        # expand_factor = 0.2
        # nn_expand_pts = 20 ##
        
        expand_factor = 0.1
        nn_expand_pts = 10 ##
        expand_save_fn = f"{self.mesh_nm}_expanded_pts_factor_{expand_factor}_nnexp_{nn_expand_pts}_new.npy"
        expand_save_fn = os.path.join(self.mesh_root, expand_save_fn) # 
        
        if not os.path.exists(expand_save_fn):
            cur_expanded_visual_pts = []
            if self.cur_expanded_visual_pts is None:
                cur_src_pts = self.vertices
            else:
                cur_src_pts = self.cur_expanded_visual_pts
            maxx_verts, _ = torch.max(cur_src_pts, dim=0)
            minn_verts, _ = torch.min(cur_src_pts, dim=0)
            extent_verts = maxx_verts - minn_verts ## (3,)-dim vecotr
            norm_extent_verts = torch.norm(extent_verts, dim=-1).item() ## (1,)-dim vector
            expand_r = norm_extent_verts * expand_factor
            # nn_expand_pts = 5 # expand the vertices to 5 times of the original vertices 
            for i_pts in range(self.vertices.size(0)):
                cur_pts = cur_src_pts[i_pts]
                # 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))
                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)
            np.save(expand_save_fn, cur_expanded_visual_pts.detach().cpu().numpy())
        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()
        self.cur_expanded_visual_pts = cur_expanded_visual_pts # expanded visual pts #
        return self.cur_expanded_visual_pts

## epand 
## link urdf ## expand the visual pts to form the expanded visual grids pts #
# use get_name_to_visual_pts_faces to get the transformed visual pts and faces #
class Link_urdf:
    def __init__(self, name, inertial: Inertial, visual: Visual=None) -> None:
        
        self.name = name
        self.inertial = inertial
        self.visual = visual # vsiual meshes #
        
        # self.joint = joint
        # self.body = body
        # self.children = children
        # self.name = name
        
        self.link_idx = ...
        
        # self.args = args
        
        self.joint = None # joint name to struct
        # self.join
        self.children = ...
        self.children = {} # joint name to child sruct
        
    def expand_visual_pts(self, expanded_visual_pts, link_name_to_visited, link_name_to_link_struct):
        link_name_to_visited[self.name] = 1
        if self.visual is not None:
            cur_expanded_visual_pts = self.visual.expand_visual_pts()
            expanded_visual_pts.append(cur_expanded_visual_pts)
        
        for cur_link in self.children:
            cur_link_struct = link_name_to_link_struct[self.children[cur_link]]
            cur_link_name = cur_link_struct.name
            if cur_link_name in link_name_to_visited:
                continue
            ## expanded visual pts for the expand visual ptsS ## 
            ## link name to visited ## 
            expanded_visual_pts = cur_link_struct.expand_visual_pts(expanded_visual_pts, link_name_to_visited, link_name_to_link_struct)
        return expanded_visual_pts
    
    def set_initial_state(self, states, action_joint_name_to_joint_idx, link_name_to_visited, link_name_to_link_struct):
        
        link_name_to_visited[self.name] = 1
        
        if self.joint is not None:
            for cur_joint_name in self.joint:
                cur_joint = self.joint[cur_joint_name]
                cur_joint_name = cur_joint.name
                cur_child = self.children[cur_joint_name]
                cur_child_struct = link_name_to_link_struct[cur_child]
                cur_child_name = cur_child_struct.name
                
                if cur_child_name in link_name_to_visited:
                    continue
                if cur_joint.type in ['revolute']:
                    cur_joint_idx = action_joint_name_to_joint_idx[cur_joint_name] # action joint name to joint idx #
                    # cur_joint_idx = action_joint_name_to_joint_idx[cur_joint_name] # 
                    # cur_joint = self.joint[cur_joint_name]
                    cur_state = states[cur_joint_idx] ### joint state ###
                    cur_joint.set_initial_state(cur_state)
                cur_child_struct.set_initial_state(states, action_joint_name_to_joint_idx, link_name_to_visited, link_name_to_link_struct)
        
        
        
    def set_penetration_forces(self,  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):
        link_name_to_visited[self.name] = 1
        
        # the current joint of  the  # update state #
        if self.joint is not None:
            for cur_joint_name in self.joint:
                
                cur_joint = self.joint[cur_joint_name] # joint model 
                
                cur_child = self.children[cur_joint_name] # child model #
                
                cur_child_struct = link_name_to_link_struct[cur_child]
                
                cur_child_name = cur_child_struct.name
                
                cur_child_link_idx = cur_child_struct.link_idx
                
                if cur_child_name in link_name_to_visited:
                    continue
                
                try:
                    cur_child_inertia = cur_child_struct.cur_inertia
                except:
                    cur_child_inertia = torch.eye(3, dtype=torch.float32)
                
                
                if cur_joint.type in ['revolute'] and (cur_joint_name not in ['WRJ2', 'WRJ1']):
                    cur_joint_idx = action_joint_name_to_joint_idx[cur_joint_name]
                    # cur_action = actions[cur_joint_idx]
                    ### get the child struct ###
                    # set_actions_and_update_states(self, action, cur_timestep, time_cons, cur_inertia):
                    # set actions and update states #
                    cur_joint_rot, cur_joint_trans = cur_joint.compute_transformation_from_current_state(n_grad=True)
                    cur_joint_tot_rot = torch.matmul(parent_rot, cur_joint_rot) ## R_p (R_j p + t_j) + t_p
                    cur_joint_tot_trans = torch.matmul(parent_rot, cur_joint_trans.unsqueeze(-1)).squeeze(-1) + parent_trans
                    
                    # cur_joint.set_actions_and_update_states_v2(cur_action, cur_timestep, time_cons, cur_child_inertia.detach(), parent_rot, parent_trans + cur_joint.origin_xyz, penetration_forces=penetration_forces, link_idx=cur_child_link_idx)
                    
                    # cur_timestep, time_cons, cur_inertia, cur_joint_tot_rot=None, cur_joint_tot_trans=None, penetration_forces=None, sampled_visual_pts_joint_idxes=None, joint_idx=None
                    
                    
                    cur_joint.set_penetration_forces(cur_child_inertia.detach(), cur_joint_tot_rot, cur_joint_tot_trans, link_idx=cur_child_link_idx, penetration_forces=penetration_forces, sampled_visual_pts_joint_idxes=sampled_visual_pts_joint_idxes, joint_idx=cur_joint_idx - 2, joint_penetration_forces=joint_penetration_forces)
                else:
                    cur_joint_tot_rot = parent_rot
                    cur_joint_tot_trans = parent_trans
                    
                
                cur_child_struct.set_penetration_forces(action_joint_name_to_joint_idx, link_name_to_visited, link_name_to_link_struct,  parent_rot=cur_joint_tot_rot, parent_trans=cur_joint_tot_trans, penetration_forces=penetration_forces, sampled_visual_pts_joint_idxes=sampled_visual_pts_joint_idxes, joint_penetration_forces=joint_penetration_forces)


    
        
    
    def get_init_visual_meshes(self, parent_rot, parent_trans, init_visual_meshes, link_name_to_link_struct, link_name_to_visited, expanded_pts=False, joint_idxes=None, state_vals=None):
        link_name_to_visited[self.name] = 1
        
        # 'transformed_joint_pos': [], 'link_idxes': []
        if self.joint is not None: # get init visual meshes #
            # for i_ch, (cur_joint, cur_child) in enumerate(zip(self.joint, self.children)):
            #     print(f"joint: {cur_joint.name}, child: {cur_child.name}, parent: {self.name}, child_visual: {cur_child.visual is not None}")
            #     joint_origin_xyz = cur_joint.origin_xyz
            #     init_visual_meshes = cur_child.get_init_visual_meshes(parent_rot, parent_trans + joint_origin_xyz, init_visual_meshes)
            # print(f"name: {self.name}, keys: {self.joint.keys()}")
            for cur_joint_name in self.joint: # 
                cur_joint = self.joint[cur_joint_name]
                
                # if state_vals is not None:
                #     cur_joint_idx = cur_joint.joint_idx
                #     state_vals[cur_joint_idx] = cur_joint.state.detach().cpu().numpy()
                
                cur_child_name = self.children[cur_joint_name]
                cur_child = link_name_to_link_struct[cur_child_name]
                # print(f"joint: {cur_joint.name}, child: {cur_child_name}, parent: {self.name}, child_visual: {cur_child.visual is not None}")
                # print(f"joint: {cur_joint.name}, child: {cur_child_name}, parent: {self.name}, child_visual: {cur_child.visual is not None}")
                joint_origin_xyz = cur_joint.origin_xyz
                if cur_child_name in link_name_to_visited:
                    continue
                cur_child_visual_pts = {'vertices': [], 'faces': [], 'link_idxes': [], 'transformed_joint_pos': [], 'joint_link_idxes': []}
                
                # joint idxes #
                cur_child_visual_pts, joint_idxes = cur_child.get_init_visual_meshes(parent_rot, parent_trans + joint_origin_xyz, cur_child_visual_pts, link_name_to_link_struct, link_name_to_visited, expanded_pts=expanded_pts, joint_idxes=joint_idxes)
                
                cur_child_verts, cur_child_faces = cur_child_visual_pts['vertices'], cur_child_visual_pts['faces']
                cur_child_link_idxes = cur_child_visual_pts['link_idxes']
                cur_transformed_joint_pos = cur_child_visual_pts['transformed_joint_pos']
                joint_link_idxes = cur_child_visual_pts['joint_link_idxes']
                
                if len(cur_child_verts) > 0:
                    cur_child_verts, cur_child_faces = merge_meshes(cur_child_verts, cur_child_faces)
                    cur_child_verts = cur_child_verts + cur_joint.origin_xyz.unsqueeze(0)
                    cur_joint_rot, cur_joint_trans = cur_joint.compute_transformation_from_current_state()
                    cur_child_verts = torch.matmul(cur_joint_rot, cur_child_verts.transpose(1, 0).contiguous()).transpose(1, 0).contiguous() + cur_joint_trans.unsqueeze(0)
                    
                    if len(cur_transformed_joint_pos) > 0:
                        cur_transformed_joint_pos = torch.cat(cur_transformed_joint_pos, dim=0)
                        cur_transformed_joint_pos = cur_transformed_joint_pos + cur_joint.origin_xyz.unsqueeze(0)
                        cur_transformed_joint_pos = torch.matmul(cur_joint_rot, cur_transformed_joint_pos.transpose(1, 0).contiguous()).transpose(1, 0).contiguous() + cur_joint_trans.unsqueeze(0)
                        cur_joint_pos = cur_joint_trans.unsqueeze(0).clone()
                        cur_transformed_joint_pos = torch.cat(
                            [cur_transformed_joint_pos, cur_joint_pos], dim=0 ##### joint poses #####
                        )
                    else:
                        cur_transformed_joint_pos = cur_joint_trans.unsqueeze(0).clone()
                    
                    if len(joint_link_idxes) > 0:
                        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)], dim=-1
                        )
                    else:
                        joint_link_idxes = torch.tensor([cur_child.link_idx], dtype=torch.long).view(1,)
                    
                    # joint link idxes #
                    
                    # cur_child_verts = cur_child_verts + # transformed joint pos #
                    cur_child_link_idxes = torch.cat(cur_child_link_idxes, dim=-1)
                    # joint_link_idxes = torch.cat(joint_link_idxes, dim=-1)
                    init_visual_meshes['vertices'].append(cur_child_verts)
                    init_visual_meshes['faces'].append(cur_child_faces)
                    init_visual_meshes['link_idxes'].append(cur_child_link_idxes)
                    init_visual_meshes['transformed_joint_pos'].append(cur_transformed_joint_pos)
                    init_visual_meshes['joint_link_idxes'].append(joint_link_idxes)
                    
            # 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)
        # self.parent_rot_mtx = parent_rot
        # self.parent_trans_vec = parent_trans + joint_origin_xyz
        
        
        if self.visual is not None:
            # ## 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)+ 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)
                joint_idxes.append(cur_idxes)
                
        
        
        # for cur_link in self.children: # 
        #     init_visual_meshes = cur_link.get_init_visual_meshes(self.parent_rot_mtx, self.parent_trans_vec, init_visual_meshes)
        return init_visual_meshes, joint_idxes ## init visual meshes ## 
    
    # 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)
        
        if self.joint is not None:
            for joint_nm in self.joint:
                cur_joint = self.joint[joint_nm]
                cur_child = self.children[joint_nm]
                cur_child_struct = link_name_to_link_struct[cur_child]
                cur_child_name = cur_child_struct.name
                if cur_child_name in link_name_to_visited:
                    continue
                joint_rot, joint_trans = cur_joint.compute_transformation_from_current_state(n_grad=True)
                # cur_parent_rot = torch.matmul(parent_rot, joint_rot) # 
                # cur_parent_trans = torch.matmul(parent_rot, joint_trans.unsqueeze(-1)).squeeze(-1) + parent_trans # 
                child_inertia = cur_child_struct.calculate_inertia(link_name_to_visited, link_name_to_link_struct)
                child_inertia = torch.matmul(
                    joint_rot.detach(), torch.matmul(child_inertia, joint_rot.detach().transpose(1, 0).contiguous())
                ).detach()
                self.cur_inertia += child_inertia
        # if self.visual is not None:
        # self.cur_inertia += self.visual.inertia
        self.cur_inertia += self.inertial.inertia.detach()
        return self.cur_inertia
    
    
    def set_delta_state_and_update(self, states, cur_timestep, link_name_to_visited, action_joint_name_to_joint_idx, link_name_to_link_struct):
        
        link_name_to_visited[self.name] = 1
        
        if self.joint is not None:
            for cur_joint_name in self.joint:
                
                cur_joint = self.joint[cur_joint_name] # joint model 
                
                cur_child = self.children[cur_joint_name] # child model #
                
                cur_child_struct = link_name_to_link_struct[cur_child]
                
                cur_child_name = cur_child_struct.name
                
                if cur_child_name in link_name_to_visited:
                    continue
                
                ## cur child inertia ##
                # cur_child_inertia = cur_child_struct.cur_inertia
                
                
                if cur_joint.type in ['revolute']:
                    cur_joint_idx = action_joint_name_to_joint_idx[cur_joint_name]
                    cur_state = states[cur_joint_idx]
                    ### get the child struct ###
                    # set_actions_and_update_states(self, action, cur_timestep, time_cons, cur_inertia):
                    # set actions and update states #
                    cur_joint.set_delta_state_and_update(cur_state, cur_timestep)
                
                cur_child_struct.set_delta_state_and_update(states, cur_timestep, link_name_to_visited, action_joint_name_to_joint_idx, link_name_to_link_struct)

    def set_delta_state_and_update_v2(self, states, cur_timestep, link_name_to_visited, action_joint_name_to_joint_idx, link_name_to_link_struct):
        link_name_to_visited[self.name] = 1
        
        if self.joint is not None:
            for cur_joint_name in self.joint:
                cur_joint = self.joint[cur_joint_name] # joint model 
                cur_child = self.children[cur_joint_name] # child model #
                cur_child_struct = link_name_to_link_struct[cur_child]
                cur_child_name = cur_child_struct.name
                if cur_child_name in link_name_to_visited:
                    continue
                # cur_child_inertia = cur_child_struct.cur_inertia
                if cur_joint.type in ['revolute']:
                    cur_joint_idx = action_joint_name_to_joint_idx[cur_joint_name]
                    cur_state = states[cur_joint_idx]
                    ### get the child struct ###
                    # set_actions_and_update_states(self, action, cur_timestep, time_cons, cur_inertia):
                    # set actions and update states #
                    cur_joint.set_delta_state_and_update_v2(cur_state, cur_timestep)
                cur_child_struct.set_delta_state_and_update_v2(states, cur_timestep, link_name_to_visited, action_joint_name_to_joint_idx, link_name_to_link_struct)
        
    # get_joint_state(self, cur_ts, state_vals):
    def get_joint_state(self, cur_ts, state_vals, link_name_to_visited, link_name_to_link_struct, action_joint_name_to_joint_idx):
        link_name_to_visited[self.name] = 1
        if self.joint is not None:
            for cur_joint_name in self.joint:
                
                cur_joint = self.joint[cur_joint_name] # joint model 
                
                cur_child = self.children[cur_joint_name] # child model #
                
                cur_child_struct = link_name_to_link_struct[cur_child]
                
                cur_child_name = cur_child_struct.name
                
                if cur_child_name in link_name_to_visited:
                    continue
                
                if cur_joint.type in ['revolute']:
                    cur_joint_idx = action_joint_name_to_joint_idx[cur_joint_name]
                    state_vals[cur_joint_idx] = cur_joint.timestep_to_states[cur_ts + 1] # .state.detach().cpu().numpy()
                    # state_vals = cur_joint.get_joint_state(cur_ts, state_vals)
                
                state_vals = cur_child_struct.get_joint_state(cur_ts, state_vals, link_name_to_visited, link_name_to_link_struct, action_joint_name_to_joint_idx)
        return state_vals
        
    # the joint #
    # set_actions_and_update_states(actions, cur_timestep, time_cons, action_joint_name_to_joint_idx, link_name_to_visited, self.link_name_to_link_struct)
    def set_actions_and_update_states(self, actions, cur_timestep, time_cons, action_joint_name_to_joint_idx, link_name_to_visited, link_name_to_link_struct):
        
        link_name_to_visited[self.name] = 1
        
        # the current joint of  the  # update state #
        if self.joint is not None:
            for cur_joint_name in self.joint:
                
                cur_joint = self.joint[cur_joint_name] # joint model 
                
                cur_child = self.children[cur_joint_name] # child model #
                
                cur_child_struct = link_name_to_link_struct[cur_child]
                
                cur_child_name = cur_child_struct.name
                
                if cur_child_name in link_name_to_visited:
                    continue
                
                cur_child_inertia = cur_child_struct.cur_inertia
                
                
                if cur_joint.type in ['revolute']:
                    cur_joint_idx = action_joint_name_to_joint_idx[cur_joint_name]
                    cur_action = actions[cur_joint_idx]
                    ### get the child struct ###
                    # set_actions_and_update_states(self, action, cur_timestep, time_cons, cur_inertia):
                    # set actions and update states #
                    cur_joint.set_actions_and_update_states(cur_action, cur_timestep, time_cons, cur_child_inertia.detach())
                
                cur_child_struct.set_actions_and_update_states(actions, cur_timestep, time_cons, action_joint_name_to_joint_idx, link_name_to_visited, link_name_to_link_struct)


    def set_actions_and_update_states_v2(self, actions, 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=None, sampled_visual_pts_joint_idxes=None):
        
        link_name_to_visited[self.name] = 1
        
        # the current joint of  the  # update state #
        if self.joint is not None:
            for cur_joint_name in self.joint:
                
                cur_joint = self.joint[cur_joint_name] # joint model 
                
                cur_child = self.children[cur_joint_name] # child model #
                
                cur_child_struct = link_name_to_link_struct[cur_child]
                
                cur_child_name = cur_child_struct.name
                
                cur_child_link_idx = cur_child_struct.link_idx
                
                if cur_child_name in link_name_to_visited:
                    continue
                
                try:
                    cur_child_inertia = cur_child_struct.cur_inertia
                except:
                    cur_child_inertia = torch.eye(3, dtype=torch.float32)
                
                
                if cur_joint.type in ['revolute']:
                    cur_joint_idx = action_joint_name_to_joint_idx[cur_joint_name]
                    cur_action = actions[cur_joint_idx]
                    ### get the child struct ###
                    # set_actions_and_update_states(self, action, cur_timestep, time_cons, cur_inertia):
                    # set actions and update states #
                    cur_joint_rot, cur_joint_trans = cur_joint.compute_transformation_from_current_state(n_grad=True)
                    cur_joint_tot_rot = torch.matmul(parent_rot, cur_joint_rot) ## R_p (R_j p + t_j) + t_p
                    cur_joint_tot_trans = torch.matmul(parent_rot, cur_joint_trans.unsqueeze(-1)).squeeze(-1) + parent_trans
                    
                    # cur_joint.set_actions_and_update_states_v2(cur_action, cur_timestep, time_cons, cur_child_inertia.detach(), parent_rot, parent_trans + cur_joint.origin_xyz, penetration_forces=penetration_forces, link_idx=cur_child_link_idx)
                    
                    cur_joint.set_actions_and_update_states_v2(cur_action, cur_timestep, time_cons, cur_child_inertia.detach(), cur_joint_tot_rot, cur_joint_tot_trans, penetration_forces=penetration_forces, link_idx=cur_child_link_idx, sampled_visual_pts_joint_idxes=sampled_visual_pts_joint_idxes)
                else:
                    cur_joint_tot_rot = parent_rot
                    cur_joint_tot_trans = parent_trans
                    
                
                cur_child_struct.set_actions_and_update_states_v2(actions, cur_timestep, time_cons,  action_joint_name_to_joint_idx, link_name_to_visited, link_name_to_link_struct,  parent_rot=cur_joint_tot_rot, parent_trans=cur_joint_tot_trans, penetration_forces=penetration_forces, sampled_visual_pts_joint_idxes=sampled_visual_pts_joint_idxes)


    def set_init_states_target_value(self, init_states):
        if self.joint.type == 'revolute':
            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)
            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).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)
        self.joint.timestep_to_states[0] = self.joint.state.detach()
        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() 
    
    
    def get_visual_pts(self, visual_pts_list):
        visual_pts_list = self.body.get_visual_pts(visual_pts_list)
        for cur_link in self.children:
            visual_pts_list = cur_link.get_visual_pts(visual_pts_list)
        visual_pts_list = torch.cat(visual_pts_list, dim=0)
        return visual_pts_list
    
    def get_visual_faces_list(self, visual_faces_list):
        visual_faces_list = self.body.get_visual_faces_list(visual_faces_list)
        for cur_link in self.children:
            visual_faces_list = cur_link.get_visual_faces_list(visual_faces_list)
        return visual_faces_list
        # pass

    
    def set_state(self, name_to_state):
        self.joint.set_state(name_to_state=name_to_state)
        for child_link in self.children:
            child_link.set_state(name_to_state)
            
    def set_state_via_vec(self, state_vec):
        self.joint.set_state_via_vec(state_vec)
        for child_link in self.children:
            child_link.set_state_via_vec(state_vec)




class Joint_Limit:
    def __init__(self, effort, lower, upper, velocity) -> None:
        self.effort = effort
        self.lower = lower
        self.velocity = velocity
        self.upper = upper
        pass

# Joint_urdf(name, joint_type, parent_link, child_link, origin_xyz, axis_xyz, limit: Joint_Limit)
class Joint_urdf: # 
    
    def __init__(self, name, joint_type, parent_link, child_link, origin_xyz, axis_xyz, limit: Joint_Limit, origin_xyz_string="") -> None:
        self.name = name
        self.type = joint_type
        self.parent_link = parent_link
        self.child_link = child_link
        self.origin_xyz = origin_xyz
        self.axis_xyz = axis_xyz
        self.limit = limit
        
        self.origin_xyz_string = origin_xyz_string
        
        # joint angle; joint state #
        self.timestep_to_vels = {}
        self.timestep_to_states = {}
        
        self.init_pos = self.origin_xyz.clone()
        
        #### only for the current state #### # joint urdf #
        self.state = nn.Parameter(
            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), 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), 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).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)
        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).detach() 
        delta_rot_vec = self.axis_xyz * state
        if cur_timestep == 0:
            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? #
        cur_state = prev_state + update_quaternion(delta_rot_vec, prev_state)
        self.timestep_to_states[cur_timestep] = cur_state.detach()
        self.state = cur_state
    
    
    def set_delta_state_and_update_v2(self, delta_state, cur_timestep):
        self.timestep_to_vels[cur_timestep] = torch.zeros((3,), dtype=torch.float32).detach() 
        
        if cur_timestep == 0:
            cur_state = delta_state
        else:
            # prev_state = self.timestep_to_states[cur_timestep - 1].detach()
            # prev_state = self.timestep_to_states[cur_timestep - 1]
            cur_state = self.timestep_to_states[cur_timestep - 1].detach() + delta_state
        ## cur_state ## # 
        self.timestep_to_states[cur_timestep] = cur_state  # .detach()
        
        
        # delta_rot_vec = self.axis_xyz * state #
        
        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)
        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)
        # else:
        #     # prev_state = self.timestep_to_states[cur_timestep - 1].detach()
        #     prev_state = self.timestep_to_states[cur_timestep - 1] # .detach() # not detach? #
        # cur_state = prev_state + update_quaternion(delta_rot_vec, prev_state)
        # self.timestep_to_states[cur_timestep] = cur_state.detach()
        # self.state = cur_state
    
    
    def compute_transformation_from_current_state(self, n_grad=False):
        # together with the parent rot mtx and the parent trans vec #
        # cur_joint_state = self.state
        if self.type == "revolute":
            # rot_mtx = rotation_matrix_from_axis_angle(self.axis, cur_joint_state)
            # trans_vec = self.pos - np.matmul(rot_mtx, self.pos.reshape(3, 1)).reshape(3)
            if n_grad:
                rot_mtx = quaternion_to_matrix(self.state.detach())
            else:
                rot_mtx = quaternion_to_matrix(self.state)
            # trans_vec = self.pos - torch.matmul(rot_mtx, self.pos.view(3, 1)).view(3).contiguous()
            trans_vec = self.origin_xyz - torch.matmul(rot_mtx, self.origin_xyz.view(3, 1)).view(3).contiguous()
            self.rot_mtx = rot_mtx
            self.trans_vec = trans_vec
        elif self.type == "fixed":
            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 # 
        else:
            pass
        return self.rot_mtx, self.trans_vec

    
    # set actions # set actions and udpate states #
    def set_actions_and_update_states(self, action, cur_timestep, time_cons, cur_inertia):
        
        # timestep_to_vels, timestep_to_states, state #
        if self.type in ['revolute']:
            
            self.action = action
            # 
            # visual_pts and visual_pts_mass #
            # cur_joint_pos = self.joint.pos #
            # TODO: check whether the following is correct # # set 
            torque = self.action * self.axis_xyz
            
            # # Compute inertia matrix #
            # 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]
            #     cur_r = cur_pts - cur_joint_pos # r_i
            #     # 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)
            #     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)
            # # Use torque to update angular velocity -> state #
            # inertia_inv = torch.linalg.inv(inertial)
            
            # axis-angle of 
            # inertia_inv = self.cur_inertia_inv
            # print(f"updating actions and states for the joint {self.name} with type {self.type}")
            inertia_inv = torch.linalg.inv(cur_inertia).detach()
            
            delta_omega = torch.matmul(inertia_inv, torque.unsqueeze(-1)).squeeze(-1)
            
            # delta_omega = torque / 400 # # axis_xyz #
            
            ## actions -> with the dynamic information -> time cons -> angular acc -> delta angular vel -> delta angle
            # TODO: dt should be an optim#izable constant? should it be the same value as that optimized for the passive object? # 
            delta_angular_vel = delta_omega * time_cons #  * self.args.dt
            delta_angular_vel = delta_angular_vel.squeeze(0)
            if cur_timestep > 0: ## cur_timestep - 1 ##
                prev_angular_vel = self.timestep_to_vels[cur_timestep - 1].detach()
                # cur_angular_vel = prev_angular_vel + delta_angular_vel * DAMPING
                cur_angular_vel = prev_angular_vel * DAMPING + delta_angular_vel # p
            else:
                cur_angular_vel = delta_angular_vel # angular vel #
            
            self.timestep_to_vels[cur_timestep] = cur_angular_vel.detach()

            cur_delta_quat = cur_angular_vel * time_cons #  * self.args.dt
            cur_delta_quat = cur_delta_quat.squeeze(0) # delta quat # 
            cur_state = self.timestep_to_states[cur_timestep].detach() # quaternion #
            # print(f"cur_delta_quat: {cur_delta_quat.size()}, cur_state: {cur_state.size()}")
            nex_state = cur_state + update_quaternion(cur_delta_quat, cur_state)
            self.timestep_to_states[cur_timestep + 1] = nex_state.detach()
            self.state = nex_state  # set the joint state #


    def set_actions_and_update_states_v2(self, action, cur_timestep, time_cons, cur_inertia, cur_joint_tot_rot=None, cur_joint_tot_trans=None, penetration_forces=None, link_idx=None, sampled_visual_pts_joint_idxes=None):
        
        # timestep_to_vels, timestep_to_states, state #
        if self.type in ['revolute']:
            
            self.action = action ## strategy 2 
            # 
            # visual_pts and visual_pts_mass #
            # cur_joint_pos = self.joint.pos #
            # TODO: check whether the following is correct # # set 
            
            if penetration_forces is not None:
                penetration_forces_values = penetration_forces['penetration_forces'].detach()
                penetration_forces_points = penetration_forces['penetration_forces_points'].detach()
                
                ####### use a part of peentration points and forces #######
                if sampled_visual_pts_joint_idxes is not None:
                    selected_forces_mask = sampled_visual_pts_joint_idxes == link_idx ## select the current link's penetrated points 
                else:
                    selected_forces_mask = torch.ones_like(penetration_forces_values[:, 0]).bool()
                ####### use a part of peentration points and forces #######
                
                if torch.sum(selected_forces_mask.float()) > 0.5: ## has penetrated points in this link ##
                    
                    penetration_forces_values = penetration_forces_values[selected_forces_mask]
                    penetration_forces_points = penetration_forces_points[selected_forces_mask]
                    # tot_rot_mtx, tot_trans_vec
                    # cur_joint_rot =  self.tot_rot_mtx
                    # cur_joint_trans = self.tot_trans_vec
                    cur_joint_rot =  cur_joint_tot_rot.detach()
                    cur_joint_trans = cur_joint_tot_trans.detach() ## total rot; total trans ##
                    local_frame_penetration_forces_values = torch.matmul(cur_joint_rot.transpose(1, 0), penetration_forces_values.transpose(1, 0)).transpose(1, 0)
                    local_frame_penetration_forces_points = torch.matmul(cur_joint_rot.transpose(1, 0), (penetration_forces_points - cur_joint_trans.unsqueeze(0)).transpose(1, 0)).transpose(1, 0)
                    
                    joint_pos_to_forces_points = local_frame_penetration_forces_points - self.axis_xyz.unsqueeze(0)
                    forces_torques = torch.cross(joint_pos_to_forces_points, local_frame_penetration_forces_values) # forces values of the local frame #
                    forces_torques = torch.sum(forces_torques, dim=0)
                    
                    forces_torques_dot_axis = torch.sum(self.axis_xyz * forces_torques)
                    penetration_delta_state = forces_torques_dot_axis
                else:
                    penetration_delta_state = 0.0
            else:
                penetration_delta_state = 0.0
                
            
            torque = self.action * self.axis_xyz
            
            # # Compute inertia matrix #
            # 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]
            #     cur_r = cur_pts - cur_joint_pos # r_i
            #     # 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)
            #     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)
            # # Use torque to update angular velocity -> state #
            # inertia_inv = torch.linalg.inv(inertial)
            
            # axis-angle of 
            # inertia_inv = self.cur_inertia_inv
            # print(f"updating actions and states for the joint {self.name} with type {self.type}")
            
            
            # inertia_inv = torch.linalg.inv(cur_inertia).detach()
            
            inertia_inv = torch.eye(n=3, dtype=torch.float32)
            
            
            
            delta_omega = torch.matmul(inertia_inv, torque.unsqueeze(-1)).squeeze(-1)
            
            # delta_omega = torque / 400
            
            
            # TODO: dt should be an optim#izable constant? should it be the same value as that optimized for the passive object? # 
            delta_angular_vel = delta_omega * time_cons #  * self.args.dt
            delta_angular_vel = delta_angular_vel.squeeze(0)
            if cur_timestep > 0: ## cur_timestep - 1 ##
                prev_angular_vel = self.timestep_to_vels[cur_timestep - 1].detach()
                # cur_angular_vel = prev_angular_vel + delta_angular_vel * DAMPING
                cur_angular_vel = prev_angular_vel * DAMPING + delta_angular_vel # p
                # cur_angular_vel = prev_angular_vel + delta_angular_vel # p
            else:
                cur_angular_vel = delta_angular_vel # angular vel #
            
            self.timestep_to_vels[cur_timestep] = cur_angular_vel.detach()

            cur_delta_angle = cur_angular_vel * time_cons #  * self.args.dt
            # cur_delta_quat = cur_delta_angle.squeeze(0) # delta quat # 
            # cur_state = self.timestep_to_states[cur_timestep].detach() # quaternion #
            # # print(f"cur_delta_quat: {cur_delta_quat.size()}, cur_state: {cur_state.size()}")
            # nex_state = cur_state + update_quaternion(cur_delta_quat, cur_state)
            
            ### strategy 2 ###
            dot_cur_delta_angle_w_axis = torch.sum( ## delta angle with axises ##
                cur_delta_angle * self.axis_xyz, dim=-1
            )
            ## dot cur deltawith the 
            delta_state = dot_cur_delta_angle_w_axis ## delta angle w axieses ##
            
            # if cur_timestep 
            if cur_timestep == 0:
                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) + update_quaternion(state_rot_vec_along_axis, torch.tensor([1., 0., 0., 0.], dtype=torch.float32))
            ### state 
            self.state = state_rot_quat
            ### get states? ##
            self.timestep_to_states[cur_timestep + 1] = nex_state # .detach()
            # self.state = nex_state  # set the joint state #
    
    
    
    def set_penetration_forces(self, cur_inertia, cur_joint_tot_rot=None, cur_joint_tot_trans=None, link_idx=None, penetration_forces=None, sampled_visual_pts_joint_idxes=None, joint_idx=None, joint_penetration_forces=None):
        
        # timestep_to_vels, timestep_to_states, state #
        if self.type in ['revolute'] :
            
            # self.action = action ## strategy 2 
            # 
            # visual_pts and visual_pts_mass #
            # cur_joint_pos = self.joint.pos #
            # TODO: check whether the following is correct # # set 
            
            if penetration_forces is not None:
                penetration_forces_values = penetration_forces['penetration_forces'].detach()
                penetration_forces_points = penetration_forces['penetration_forces_points'].detach()
                
                ####### use a part of peentration points and forces #######
                if sampled_visual_pts_joint_idxes is not None:
                    selected_forces_mask = sampled_visual_pts_joint_idxes == link_idx ## select the current link's penetrated points 
                else:
                    selected_forces_mask = torch.ones_like(penetration_forces_values[:, 0]).bool()
                ####### use a part of peentration points and forces #######
                
                if torch.sum(selected_forces_mask.float()) > 0.5: ## has penetrated points in this link ##
                    
                    penetration_forces_values = penetration_forces_values[selected_forces_mask]
                    penetration_forces_points = penetration_forces_points[selected_forces_mask]
                    # tot_rot_mtx, tot_trans_vec
                    # cur_joint_rot =  self.tot_rot_mtx
                    # cur_joint_trans = self.tot_trans_vec
                    cur_joint_rot =  cur_joint_tot_rot.detach()
                    cur_joint_trans = cur_joint_tot_trans.detach() ## total rot; total trans ##
                    local_frame_penetration_forces_values = torch.matmul(cur_joint_rot.transpose(1, 0), penetration_forces_values.transpose(1, 0)).transpose(1, 0)
                    local_frame_penetration_forces_points = torch.matmul(cur_joint_rot.transpose(1, 0), (penetration_forces_points - cur_joint_trans.unsqueeze(0)).transpose(1, 0)).transpose(1, 0)
                    
                    joint_pos_to_forces_points = local_frame_penetration_forces_points - self.axis_xyz.unsqueeze(0)
                    forces_torques = torch.cross(joint_pos_to_forces_points, local_frame_penetration_forces_values) # forces values of the local frame #
                    forces_torques = torch.sum(forces_torques, dim=0)
                    
                    forces = torch.sum(local_frame_penetration_forces_values, dim=0)
                
                    cur_joint_maximal_forces = torch.cat(
                        [forces, forces_torques], dim=0
                    )
                    cur_joint_idx = joint_idx
                    joint_penetration_forces[cur_joint_idx][:] = cur_joint_maximal_forces[:].clone()
                        
                    # forces_torques_dot_axis = torch.sum(self.axis_xyz * forces_torques)
                    # penetration_delta_state = forces_torques_dot_axis
                else:
                    penetration_delta_state = 0.0
                    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()
                    
            else:
                penetration_delta_state = 0.0
                cur_joint_idx = joint_idx
                joint_penetration_forces[cur_joint_idx][:] = cur_joint_maximal_forces[:].clone()
                
       
            
            
    
    
    def get_joint_state(self, cur_ts, state_vals):
        cur_joint_state = self.timestep_to_states[cur_ts + 1]
        state_vals[self.joint_idx] = cur_joint_state
        return state_vals


class Robot_urdf:
    def __init__(self, links, link_name_to_link_idxes, link_name_to_link_struct, joint_name_to_joint_idx, actions_joint_name_to_joint_idx, tot_joints=None, real_actions_joint_name_to_joint_idx=None) -> None:
        self.links = links
        self.link_name_to_link_idxes = link_name_to_link_idxes
        self.link_name_to_link_struct = link_name_to_link_struct
        
        # joint_name_to_joint_idx, actions_joint_name_to_joint_idx
        self.joint_name_to_joint_idx = joint_name_to_joint_idx
        self.actions_joint_name_to_joint_idx = actions_joint_name_to_joint_idx
        
        self.tot_joints = tot_joints
        # # 
        # # 
        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)[2:]
        
        self.real_actions_joint_name_to_joint_idx = real_actions_joint_name_to_joint_idx
        
        
        self.init_vertices, self.init_faces = self.get_init_visual_pts()
        
        joint_name_to_joint_idx_sv_fn = "mano_joint_name_to_joint_idx.npy"
        np.save(joint_name_to_joint_idx_sv_fn, self.joint_name_to_joint_idx)
        
        actions_joint_name_to_joint_idx_sv_fn = "mano_actions_joint_name_to_joint_idx.npy"
        np.save(actions_joint_name_to_joint_idx_sv_fn, self.actions_joint_name_to_joint_idx)
        
        tot_joints = len(self.joint_name_to_joint_idx)
        tot_actions_joints = len(self.actions_joint_name_to_joint_idx)
        
        print(f"tot_joints: {tot_joints}, tot_actions_joints: {tot_actions_joints}")
        
        pass
    
    # robot.expande
    def expand_visual_pts(self, ):
        link_name_to_visited = {}
        # transform the visual pts #
        # action_joint_name_to_joint_idx = self.actions_joint_name_to_joint_idx
        
        palm_idx = self.link_name_to_link_idxes["palm"]
        palm_link = self.links[palm_idx]
        expanded_visual_pts = []
        # expanded the visual pts # # transformed viusal pts # or the translations of the visual pts #
        expanded_visual_pts = palm_link.expand_visual_pts(expanded_visual_pts, link_name_to_visited, self.link_name_to_link_struct) 
        expanded_visual_pts = torch.cat(expanded_visual_pts, dim=0)
        # pass 
        return expanded_visual_pts
    
    
    ### samping issue? --- TODO`    `
    def get_init_visual_pts(self, expanded_pts=False, joint_idxes=None):
        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)
        init_parent_trans = torch.zeros((3,), dtype=torch.float32)
        
        palm_idx = self.link_name_to_link_idxes["palm"]
        palm_link = self.links[palm_idx]
        
        link_name_to_visited = {}
        
        ### from the palm linke ##
        init_visual_meshes, joint_idxes = palm_link.get_init_visual_meshes(init_parent_rot, init_parent_trans, init_visual_meshes, self.link_name_to_link_struct, link_name_to_visited, expanded_pts=expanded_pts, joint_idxes=joint_idxes)
        
        self.link_idxes = torch.cat(init_visual_meshes['link_idxes'], dim=-1)
        self.transformed_joint_pos = torch.cat(init_visual_meshes['transformed_joint_pos'], dim=0)
        self.joint_link_idxes = torch.cat(init_visual_meshes['joint_link_idxes'], dim=-1) ### 
        
        
        if joint_idxes is not None:
            joint_idxes = torch.cat(joint_idxes, dim=0)
            
        # for cur_link in self.links:
        #     init_visual_meshes = cur_link.get_init_visual_meshes(init_parent_rot, init_parent_trans, init_visual_meshes, self.link_name_to_link_struct, link_name_to_visited)

        init_vertices, init_faces = merge_meshes(init_visual_meshes['vertices'], init_visual_meshes['faces'])
        
        if joint_idxes is not None:
            return init_vertices, init_faces, joint_idxes
        else:
            return init_vertices, init_faces
    
    def set_penetration_forces(self, penetration_forces, sampled_visual_pts_joint_idxes, joint_penetration_forces):
        palm_idx = self.link_name_to_link_idxes["palm"]
        palm_link = self.links[palm_idx]
        
        link_name_to_visited = {}
        
        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)
        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):
        
        palm_link.set_penetration_forces(action_joint_name_to_joint_idx, link_name_to_visited, self.link_name_to_link_struct, parent_rot, parent_trans,  penetration_forces, sampled_visual_pts_joint_idxes, joint_penetration_forces)
    
    def set_delta_state_and_update(self, states, cur_timestep):
        link_name_to_visited = {}
        
        action_joint_name_to_joint_idx = self.actions_joint_name_to_joint_idx
        
        palm_idx = self.link_name_to_link_idxes["palm"]
        palm_link = self.links[palm_idx]
        
        link_name_to_visited = {}
        
        palm_link.set_delta_state_and_update(states, cur_timestep, link_name_to_visited, action_joint_name_to_joint_idx, self.link_name_to_link_struct)


    def set_delta_state_and_update_v2(self, states, cur_timestep, use_real_act_joint=False):
        link_name_to_visited = {}
        
        if use_real_act_joint:
            action_joint_name_to_joint_idx = self.real_actions_joint_name_to_joint_idx
        else:
            action_joint_name_to_joint_idx = self.actions_joint_name_to_joint_idx
        
        palm_idx = self.link_name_to_link_idxes["palm"]
        palm_link = self.links[palm_idx]
        
        link_name_to_visited = {}
        
        palm_link.set_delta_state_and_update_v2(states, cur_timestep, link_name_to_visited, action_joint_name_to_joint_idx, self.link_name_to_link_struct)
        
    
    
    # cur_joint.set_actions_and_update_states(cur_action, cur_timestep, time_cons, cur_child_inertia)
    def set_actions_and_update_states(self, actions, cur_timestep, time_cons,):
        # self.actions_joint_name_to_joint_idx as the action joint name to joint idx 
        link_name_to_visited = {}
        ## to joint idx ##
        action_joint_name_to_joint_idx = self.actions_joint_name_to_joint_idx
        
        palm_idx = self.link_name_to_link_idxes["palm"]
        palm_link = self.links[palm_idx]
        
        link_name_to_visited = {}
        
        ## set actions ##
        palm_link.set_actions_and_update_states(actions, cur_timestep, time_cons, action_joint_name_to_joint_idx, link_name_to_visited, self.link_name_to_link_struct)
        
        # for cur_joint in 
        
        # for cur_link in self.links:
        #     if cur_link.joint is not None:
        #         for cur_joint_nm in cur_link.joint:
        #             if cur_link.joint[cur_joint_nm].type in ['revolute']:
        #                 cur_link_joint_name = cur_link.joint[cur_joint_nm].name
        #                 cur_link_joint_idx = self.actions_joint_name_to_joint_idx[cur_link_joint_name]
                        
        
        # for cur_link in self.links:
        #     cur_link.set_actions_and_update_states(actions, cur_timestep, time_cons, action_joint_name_to_joint_idx, link_name_to_visited, self.link_name_to_link_struct)
        
        
    
    def get_joint_state(self, cur_ts, state_vals):
        # link_name_to_visited = {}
        ## to joint idx ##
        # action_joint_name_to_joint_idx = self.actions_joint_name_to_joint_idx
        
        palm_idx = self.link_name_to_link_idxes["palm"]
        palm_link = self.links[palm_idx]
        
        link_name_to_visited = {}
        
        # 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 ##
        state_vals = palm_link.get_joint_state(cur_ts, state_vals, link_name_to_visited, self.link_name_to_link_struct, self.actions_joint_name_to_joint_idx)
        return state_vals
        
        
    def set_actions_and_update_states_v2(self, actions, cur_timestep, time_cons, penetration_forces=None, sampled_visual_pts_joint_idxes=None):
        # self.actions_joint_name_to_joint_idx as the action joint name to joint idx 
        link_name_to_visited = {}
        ## to joint idx ##
        action_joint_name_to_joint_idx = self.actions_joint_name_to_joint_idx
        
        palm_idx = self.link_name_to_link_idxes["palm"]
        palm_link = self.links[palm_idx]
        
        link_name_to_visited = {}
        
        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 ##
        palm_link.set_actions_and_update_states_v2(actions, cur_timestep, time_cons, action_joint_name_to_joint_idx, link_name_to_visited, self.link_name_to_link_struct,  parent_rot=parent_rot, parent_trans=parent_trans, penetration_forces=penetration_forces, sampled_visual_pts_joint_idxes=sampled_visual_pts_joint_idxes)
        
        
    ### TODO: add the contact torque when calculating the nextstep states ###
    ### TODO: not an accurate implementation since differen joints should be considered ###
    ### TODO: the articulated force modle is not so easy as this one .... ###
    def set_contact_forces(self, hard_selected_forces, hard_selected_manipulating_points, hard_selected_sampled_input_pts_idxes):
        # transformed_joint_pos, joint_link_idxes, link_idxes #
        selected_pts_link_idxes = self.link_idxes[hard_selected_sampled_input_pts_idxes]
        # use the selected link idxes # 
        # selected pts idxes #
        
        # self.joint_link_idxes, transformed_joint_pos #
        self.link_idx_to_transformed_joint_pos = {}
        for i_link in range(self.transformed_joint_pos.size(0)):
            cur_link_idx = self.link_idxes[i_link].item()
            cur_link_pos = self.transformed_joint_pos[i_link]
            # if cur_link_idx not in self.link_idx_to_transformed_joint_pos:
            self.link_idx_to_transformed_joint_pos[cur_link_idx] = cur_link_pos
            # self.link_idx_to_transformed_joint_pos[cur_link_idx].append(cur_link_pos)
        
        # from the 
        self.link_idx_to_contact_forces = {}
        for i_c_pts in range(hard_selected_forces.size(0)):
            cur_contact_force = hard_selected_forces[i_c_pts] ## 
            cur_link_idx = selected_pts_link_idxes[i_c_pts].item()
            cur_link_pos = self.link_idx_to_transformed_joint_pos[cur_link_idx]
            cur_link_action_pos = hard_selected_manipulating_points[i_c_pts]
            # (action_pos - link_pos) x (-contact_force) #
            cur_contact_torque = torch.cross(
                cur_link_action_pos - cur_link_pos, -cur_contact_force
            )
            if cur_link_idx not in self.link_idx_to_contact_forces:
                self.link_idx_to_contact_forces[cur_link_idx] = [cur_contact_torque]
            else:
                self.link_idx_to_contact_forces[cur_link_idx].append(cur_contact_torque)
        for link_idx in self.link_idx_to_contact_forces:
            self.link_idx_to_contact_forces[link_idx] = torch.stack(self.link_idx_to_contact_forces[link_idx], dim=0)
            self.link_idx_to_contact_forces[link_idx] = torch.sum(self.link_idx_to_contact_forces[link_idx] , dim=0)
        for link_idx, link_struct in enumerate(self.links):
            if link_idx in self.link_idx_to_contact_forces:
                cur_link_contact_force = self.link_idx_to_contact_forces[link_idx]
                link_struct.contact_torque = cur_link_contact_force
            else:
                link_struct.contact_torque = None
        
    
    # def se ### from the optimizable initial states ###
    def set_initial_state(self, states):
        action_joint_name_to_joint_idx = self.actions_joint_name_to_joint_idx
        link_name_to_visited = {}
        
        palm_idx = self.link_name_to_link_idxes["palm"]
        palm_link = self.links[palm_idx]
        
        link_name_to_visited = {}
        
        palm_link.set_initial_state(states, action_joint_name_to_joint_idx, link_name_to_visited, self.link_name_to_link_struct)
        
        # for cur_link in self.links:
        #     cur_link.set_initial_state(states, action_joint_name_to_joint_idx, link_name_to_visited, self.link_name_to_link_struct)
            
    ### after each timestep -> re-calculate the inertial matrix using the current simulated states and the set the new actiosn and forward the simulation # 
    def calculate_inertia(self):
        link_name_to_visited = {}
        
        palm_idx = self.link_name_to_link_idxes["palm"]
        palm_link = self.links[palm_idx]
        
        link_name_to_visited = {}
        
        palm_link.calculate_inertia(link_name_to_visited, self.link_name_to_link_struct)
        
        # for cur_link in self.links:
        #     cur_link.calculate_inertia(link_name_to_visited, self.link_name_to_link_struct)
            
    ### 
    
    


def parse_nparray_from_string(strr, args=None):
    vals = strr.split(" ")
    vals = [float(val) for val in vals]
    vals = np.array(vals, dtype=np.float32)
    vals = torch.from_numpy(vals).float()
    ## vals ##
    vals = nn.Parameter(vals, requires_grad=True)
    
    return vals


### parse link data ###
def parse_link_data(link, args):
    
    link_name = link.attrib["name"]
    # print(f"parsing link: {link_name}") ## joints body meshes #
    
    joint = link.find("./joint")
    
    joint_name = joint.attrib["name"]
    joint_type = joint.attrib["type"]
    if joint_type in ["revolute"]: ## a general xml parser here? 
        axis = joint.attrib["axis"]
        axis = parse_nparray_from_string(axis, args=args)
    else:
        axis = None
    pos = joint.attrib["pos"] # 
    pos = parse_nparray_from_string(pos, args=args)
    quat = joint.attrib["quat"]
    quat = parse_nparray_from_string(quat, args=args)
    
    try:
        frame = joint.attrib["frame"]
    except:
        frame = "WORLD"
    
    if joint_type not in ["fixed"]:
        damping = joint.attrib["damping"]
        damping = float(damping)
    else:
        damping = 0.0
    
    cur_joint = Joint(joint_name, joint_type, axis, pos, quat, frame, damping, args=args)
    
    body = link.find("./body")
    body_name = body.attrib["name"]
    body_type = body.attrib["type"]
    if body_type == "mesh":
        filename = body.attrib["filename"]
    else:
        filename = ""
    
    if body_type == "sphere":
        radius = body.attrib["radius"]
        radius = float(radius)
    else:
        radius = 0.
    
    pos = body.attrib["pos"]
    pos = parse_nparray_from_string(pos, args=args)
    quat = body.attrib["quat"]
    quat = joint.attrib["quat"]
    try:
        transform_type = body.attrib["transform_type"]
    except:
        transform_type = "OBJ_TO_WORLD"
    density = body.attrib["density"]
    density = float(density)
    mu = body.attrib["mu"]
    mu = float(mu)
    try: ## rgba ##
        rgba = body.attrib["rgba"]
        rgba = parse_nparray_from_string(rgba, args=args)
    except:
        rgba = np.zeros((4,), dtype=np.float32)
    
    cur_body = Body(body_name, body_type, filename, pos, quat, transform_type, density, mu, rgba, radius, args=args)
    
    children_link = []
    links = link.findall("./link")
    for child_link in links: # 
        cur_child_link = parse_link_data(child_link, args=args)
        children_link.append(cur_child_link)
    
    link_name = link.attrib["name"]
    link_obj = Link(link_name, joint=cur_joint, body=cur_body, children=children_link, args=args)
    return link_obj


### parse link data ###
def parse_link_data_urdf(link):
    
    link_name = link.attrib["name"]
    # print(f"parsing link: {link_name}") ## joints body meshes #
    
    inertial = link.find("./inertial")
    
    origin = inertial.find("./origin")

    if origin is not None:
        inertial_pos = origin.attrib["xyz"]
        try:
            inertial_rpy = origin.attrib["rpy"]
        except:
            inertial_rpy = "0.0 0.0 0.0"
    else:
        inertial_pos = "0.0 0.0 0.0"
        inertial_rpy = "0.0 0.0 0.0"
    inertial_pos = parse_nparray_from_string(inertial_pos)

    inertial_rpy = parse_nparray_from_string(inertial_rpy)
    
    inertial_mass = inertial.find("./mass")
    inertial_mass = inertial_mass.attrib["value"]
    
    inertial_inertia = inertial.find("./inertia")
    inertial_ixx = inertial_inertia.attrib["ixx"]
    inertial_ixx = float(inertial_ixx)
    inertial_ixy = inertial_inertia.attrib["ixy"]
    inertial_ixy = float(inertial_ixy)
    inertial_ixz = inertial_inertia.attrib["ixz"]
    inertial_ixz = float(inertial_ixz)
    inertial_iyy = inertial_inertia.attrib["iyy"]
    inertial_iyy = float(inertial_iyy)
    inertial_iyz = inertial_inertia.attrib["iyz"]
    inertial_iyz = float(inertial_iyz)
    inertial_izz = inertial_inertia.attrib["izz"]
    inertial_izz = float(inertial_izz)
    
    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
    inertial_inertia_mtx[1, 0] = inertial_ixy
    inertial_inertia_mtx[1, 1] = inertial_iyy
    inertial_inertia_mtx[1, 2] = inertial_iyz
    inertial_inertia_mtx[2, 0] = inertial_ixz
    inertial_inertia_mtx[2, 1] = inertial_iyz
    inertial_inertia_mtx[2, 2] = inertial_izz
    
    # [xx, xy, xz] # 
    # [0, yy, yz] #
    # [0, 0, zz] #
    
    # a strange inertia value ... #
    # TODO: how to compute the inertia matrix? #
    
    visual = link.find("./visual")
    
    if visual is not None:
        origin = visual.find("./origin")
        visual_pos = origin.attrib["xyz"]
        visual_pos = parse_nparray_from_string(visual_pos)
        visual_rpy = origin.attrib["rpy"]
        visual_rpy = parse_nparray_from_string(visual_rpy)
        geometry = visual.find("./geometry")
        geometry_mesh = geometry.find("./mesh")
        if geometry_mesh is None:
            visual = None
        else:
            mesh_fn = geometry_mesh.attrib["filename"]

            try:
                mesh_scale = geometry_mesh.attrib["scale"]
            except:
                mesh_scale = "1 1 1"
            
            mesh_scale = parse_nparray_from_string(mesh_scale)
            mesh_fn = str(mesh_fn)
        
    
    link_struct = Link_urdf(name=link_name, inertial=Inertial(origin_rpy=inertial_rpy, origin_xyz=inertial_pos, mass=inertial_mass, inertia=inertial_inertia_mtx), visual=Visual(visual_rpy=visual_rpy, visual_xyz=visual_pos, geometry_mesh_fn=mesh_fn, geometry_mesh_scale=mesh_scale) if visual is not None else None)

    return link_struct

def parse_joint_data_urdf(joint):
    joint_name = joint.attrib["name"]
    joint_type = joint.attrib["type"]
    
    parent = joint.find("./parent")
    child = joint.find("./child")
    parent_name = parent.attrib["link"]
    child_name = child.attrib["link"]
    
    joint_origin = joint.find("./origin")
    # if joint_origin.
    try:
        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)
        origin_xyz_string = ""
        
    joint_axis = joint.find("./axis")
    if joint_axis is not None:
        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)
    
    joint_limit = joint.find("./limit")
    if joint_limit is not None:
        joint_lower = joint_limit.attrib["lower"]
        joint_lower = float(joint_lower)
        joint_upper = joint_limit.attrib["upper"]
        joint_upper = float(joint_upper)
        joint_effort = joint_limit.attrib["effort"]
        joint_effort = float(joint_effort)
        if "velocity" in joint_limit.attrib:
            joint_velocity = joint_limit.attrib["velocity"]
            joint_velocity = float(joint_velocity)
        else:
            joint_velocity = 0.5
    else:
        joint_lower = -0.5000
        joint_upper = 1.57
        joint_effort = 1000
        joint_velocity = 0.5
    
    # cosntruct the joint data #      
    joint_limit = Joint_Limit(effort=joint_effort, lower=joint_lower, upper=joint_upper, velocity=joint_velocity)
    cur_joint_struct = Joint_urdf(joint_name, joint_type, parent_name, child_name, origin_xyz, joint_axis, joint_limit, origin_xyz_string)
    return cur_joint_struct
    


def parse_data_from_urdf(xml_fn):
    
    tree = ElementTree()
    tree.parse(xml_fn)
    print(f"{xml_fn}")
    ### get total robots ###
    # robots = tree.findall("link")
    cur_robot = tree
    # i_robot = 0
    # tot_robots = []
    # for cur_robot in robots:
    # print(f"Getting robot: {i_robot}")
    # i_robot += 1
    # print(f"len(robots): {len(robots)}")
    # cur_robot = robots[0]
    cur_links = cur_robot.findall("./link")
    # curlinks
    # i_link = 0
    link_name_to_link_idxes = {}
    cur_robot_links = []
    link_name_to_link_struct = {}
    for i_link_idx, cur_link in enumerate(cur_links):
        cur_link_struct = parse_link_data_urdf(cur_link)
        print(f"Adding link {cur_link_struct.name}, link_idx: {i_link_idx}")
        cur_link_struct.link_idx = i_link_idx
        cur_robot_links.append(cur_link_struct)
        
        link_name_to_link_idxes[cur_link_struct.name] = i_link_idx
        link_name_to_link_struct[cur_link_struct.name] = cur_link_struct
    # for cur_link in cur_links:
    #     cur_robot_links.append(parse_link_data_urdf(cur_link, args=args))
    
    print(f"link_name_to_link_struct: {len(link_name_to_link_struct)}, ")
    
    tot_robot_joints = []
    
    joint_name_to_joint_idx = {}
    
    actions_joint_name_to_joint_idx = {}
    
    cur_joints = cur_robot.findall("./joint")
    
    real_actions_joint_name_to_joint_idx = {}
    
    act_joint_idx = 0
    for i_joint, cur_joint in enumerate(cur_joints):
        cur_joint_struct = parse_joint_data_urdf(cur_joint)
        cur_joint_parent_link = cur_joint_struct.parent_link
        cur_joint_child_link = cur_joint_struct.child_link
        
        cur_joint_idx = len(tot_robot_joints)
        cur_joint_name = cur_joint_struct.name
        
        joint_name_to_joint_idx[cur_joint_name] = cur_joint_idx
        
        print(f"cur_joint_name: {cur_joint_name}, cur_joint_idx: {cur_joint_idx}, axis: {cur_joint_struct.axis_xyz}, origin: {cur_joint_struct.origin_xyz}")
        
        cur_joint_type = cur_joint_struct.type
        if cur_joint_type in ['revolute']:
            actions_joint_name_to_joint_idx[cur_joint_name] = cur_joint_idx
            # actions_joint_name_to_joint_idx[cur_joint_name] = act_joint_idx
            # act_joint_idx = act_joint_idx + 1
            
            real_actions_joint_name_to_joint_idx[cur_joint_name] = act_joint_idx
            act_joint_idx = act_joint_idx + 1

        
        #### add the current joint to tot joints ###
        tot_robot_joints.append(cur_joint_struct)
        
        parent_link_idx = link_name_to_link_idxes[cur_joint_parent_link]
        cur_parent_link_struct = cur_robot_links[parent_link_idx]
        
        
        child_link_idx = link_name_to_link_idxes[cur_joint_child_link]
        cur_child_link_struct = cur_robot_links[child_link_idx]
        # parent link struct #
        if link_name_to_link_struct[cur_joint_parent_link].joint is not None:
            link_name_to_link_struct[cur_joint_parent_link].joint[cur_joint_struct.name] = cur_joint_struct
            link_name_to_link_struct[cur_joint_parent_link].children[cur_joint_struct.name] = cur_child_link_struct.name
            # cur_child_link_struct
            # cur_parent_link_struct.joint.append(cur_joint_struct)
            # cur_parent_link_struct.children.append(cur_child_link_struct)
        else:
            link_name_to_link_struct[cur_joint_parent_link].joint = {
                cur_joint_struct.name: cur_joint_struct
            }
            link_name_to_link_struct[cur_joint_parent_link].children = {
                cur_joint_struct.name: cur_child_link_struct.name
                    # cur_child_link_struct
            }
            # cur_parent_link_struct.joint = [cur_joint_struct]
            # cur_parent_link_struct.children.append(cur_child_link_struct)
        # pass
    
    print(f"actions_joint_name_to_joint_idx: {len(actions_joint_name_to_joint_idx)}")
    print(f"real_actions_joint_name_to_joint_idx: {len(real_actions_joint_name_to_joint_idx)}")
    cur_robot_obj = Robot_urdf(cur_robot_links, link_name_to_link_idxes, link_name_to_link_struct, joint_name_to_joint_idx, actions_joint_name_to_joint_idx, tot_robot_joints, real_actions_joint_name_to_joint_idx=real_actions_joint_name_to_joint_idx)
        # tot_robots.append(cur_robot_obj)
    
    print(f"Actions joint idxes:")
    print(list(actions_joint_name_to_joint_idx.keys()))
    
    actions_joint_idxes = list(actions_joint_name_to_joint_idx.values())
    actions_joint_idxes = sorted(actions_joint_idxes)
    print(f"joint indexes: {actions_joint_idxes}")
    
    # for the joint robots #
    # for every joint 
    # tot_actuators = []
    # actuators = tree.findall("./actuator/motor")
    # joint_nm_to_joint_idx = {}
    # i_act = 0
    # for cur_act in actuators:
    #     cur_act_joint_nm  = cur_act.attrib["joint"]
    #     joint_nm_to_joint_idx[cur_act_joint_nm] = i_act
    #     i_act += 1 ### add the act ###
    
    # tot_robots[0].set_joint_idx(joint_nm_to_joint_idx) ### set joint idx here ### # tot robots #
    # tot_robots[0].get_nn_pts()
    # tot_robots[1].get_nn_pts()
    
    return cur_robot_obj


def get_name_to_state_from_str(states_str):
    tot_states = states_str.split(" ")
    tot_states = [float(cur_state) for cur_state in tot_states]
    joint_name_to_state = {}
    for i in range(len(tot_states)):
        cur_joint_name = f"joint{i + 1}"
        cur_joint_state = tot_states[i]
        joint_name_to_state[cur_joint_name] = cur_joint_state
    return joint_name_to_state


def merge_meshes(verts_list, faces_list):
    nn_verts = 0
    tot_verts_list = []
    tot_faces_list = []
    for i_vv, cur_verts in enumerate(verts_list):
        cur_verts_nn = cur_verts.size(0)
        tot_verts_list.append(cur_verts)
        tot_faces_list.append(faces_list[i_vv] + nn_verts)
        nn_verts = nn_verts + cur_verts_nn
    tot_verts_list = torch.cat(tot_verts_list, dim=0)
    tot_faces_list = torch.cat(tot_faces_list, dim=0)
    return tot_verts_list, tot_faces_list


### get init s
class RobotAgent: # robot and the robot #
    def __init__(self, xml_fn, args=None) -> None:
        global urdf_fn
        urdf_fn = xml_fn
        self.xml_fn = xml_fn
        # self.args = args
        
        ## 
        active_robot =  parse_data_from_urdf(xml_fn)
        
        self.time_constant = nn.Embedding(
            num_embeddings=3, embedding_dim=1
        )
        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,
        )
        torch.nn.init.zeros_(self.optimizable_actions.weight) #
        
        self.learning_rate = 5e-4
    
        self.active_robot = active_robot
        
        
        self.set_init_states()
        init_visual_pts = self.get_init_state_visual_pts()
        self.init_visual_pts = init_visual_pts
        
        cur_verts, cur_faces = self.active_robot.get_init_visual_pts()
        self.robot_pts = cur_verts
        self.robot_faces = cur_faces
        
        
    def set_init_states_target_value(self, init_states):
        # 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)
        self.active_robot.set_initial_state(init_states)
    
    def set_init_states(self):
        # 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()
        # ### random rotation ###
        
        # # glb_rot, glb_trans #
        # init_states = {}
        # init_states['glb_rot'] = glb_rot;
        # init_states['glb_trans'] = glb_trans; 
        # self.active_robot.set_init_states(init_states)
        
        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)
    def get_init_state_visual_pts(self, expanded_pts=False, ret_joint_idxes=False):
        # visual_pts_list = [] # compute the transformation via current state #
        # visual_pts_list, visual_pts_mass_list = self.active_robot.compute_transformation_via_current_state( visual_pts_list)
        
        if ret_joint_idxes:
            joint_idxes = [] 
            cur_verts, cur_faces, joint_idxes = self.active_robot.get_init_visual_pts(expanded_pts=expanded_pts, joint_idxes=joint_idxes)
        else:
            cur_verts, cur_faces = self.active_robot.get_init_visual_pts(expanded_pts=expanded_pts, joint_idxes=None)
        self.faces = cur_faces
        # joint_idxes = torch.cat()
        # self.robot_pts = cur_verts
        # self.robot_faces = cur_faces
        # init_visual_pts = visual_pts_list
        if ret_joint_idxes:
            return cur_verts, joint_idxes
        else:
            return cur_verts
    
    def set_actions_and_update_states(self, actions, cur_timestep):
        # 
        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)) ### 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
    
    # state_vals = self.robot_agent.get_joint_state( cur_ts, state_vals, link_name_to_link_struct)
    def get_joint_state(self, cur_ts, state_vals):
        state_vals = self.active_robot.get_joint_state(cur_ts, state_vals)
        return state_vals
    
    def forward_stepping_test(self, ):
        # delta_glb_rot; delta_glb_trans #
        timestep_to_visual_pts = {}
        for i_step in range(50):
            actions = {}
            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
            self.set_actions_and_update_states(actions=actions, cur_timestep=i_step)
    
            cur_visual_pts = robot_agent.get_init_state_visual_pts()
            cur_visual_pts = cur_visual_pts.detach().cpu().numpy()
            timestep_to_visual_pts[i_step + 1] = cur_visual_pts
        return timestep_to_visual_pts
    
    def initialize_optimization(self, reference_pts_dict):
        self.n_timesteps = 50
        # self.n_timesteps = 19 # first 19-timesteps optimization #
        self.nn_tot_optimization_iters = 100
        # self.nn_tot_optimization_iters = 57
        # 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() for ts in self.ts_to_reference_pts
        }
    
    
    def forward_stepping_optimization(self, ):
        nn_tot_optimization_iters = self.nn_tot_optimization_iters
        params_to_train = []
        params_to_train += list(self.optimizable_actions.parameters())
        self.optimizer = torch.optim.Adam(params_to_train, lr=self.learning_rate)
        
        for i_iter in range(nn_tot_optimization_iters):
            
            tot_losses = []
            ts_to_robot_points = {}
            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)
                # 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
                # self.set_actions_and_update_states(actions=actions, cur_timestep=cur_ts) # update the interaction # 
                
                with torch.no_grad():
                    self.active_robot.calculate_inertia()
    
                self.active_robot.set_actions_and_update_states(actions_link_actions, cur_ts, 0.2)

                cur_visual_pts, cur_faces = self.active_robot.get_init_visual_pts()
                ts_to_robot_points[cur_ts + 1] = cur_visual_pts.clone()
                
                cur_reference_pts = self.ts_to_reference_pts[cur_ts + 1]
                diff = torch.sum((cur_visual_pts - cur_reference_pts) ** 2, dim=-1)
                diff = diff.mean()
                
                # diff.
                self.optimizer.zero_grad()
                diff.backward(retain_graph=True)
                # diff.backward(retain_graph=False)
                self.optimizer.step()
                
                tot_losses.append(diff.item())

                
            loss = sum(tot_losses) / float(len(tot_losses))
            print(f"Iter: {i_iter}, average loss: {loss}")
            # print(f"Iter: {i_iter}, average loss: {loss.item()}, start optimizing")
            # self.optimizer.zero_grad()
            # loss.backward()
            # self.optimizer.step()
        
        self.ts_to_robot_points = {
            ts: ts_to_robot_points[ts].detach().cpu().numpy() for ts in ts_to_robot_points
        }
        self.ts_to_ref_points = {
            ts: self.ts_to_reference_pts[ts].detach().cpu().numpy() for ts in ts_to_robot_points
        }
        return self.ts_to_robot_points, self.ts_to_ref_points




def rotation_matrix_from_axis_angle(axis, angle): # rotation_matrix_from_axis_angle -> 
        # sin_ = np.sin(angle) #  ti.math.sin(angle)
        # cos_ = np.cos(angle) #  ti.math.cos(angle)
        sin_ = torch.sin(angle) #  ti.math.sin(angle)
        cos_ = torch.cos(angle) #  ti.math.cos(angle)
        u_x, u_y, u_z = axis[0], axis[1], axis[2]
        u_xx = u_x * u_x
        u_yy = u_y * u_y
        u_zz = u_z * u_z
        u_xy = u_x * u_y
        u_xz = u_x * u_z
        u_yz = u_y * u_z ## 
        
        
        row_a = torch.stack(
            [cos_ + u_xx * (1 - cos_), u_xy * (1. - cos_) + u_z * sin_, u_xz * (1. - cos_) - u_y * sin_], dim=0
        )
        # print(f"row_a: {row_a.size()}")
        row_b = torch.stack(
            [u_xy * (1. - cos_) - u_z * sin_, cos_ + u_yy * (1. - cos_), u_yz * (1. - cos_) + u_x * sin_], dim=0
        )
        # print(f"row_b: {row_b.size()}")
        row_c = torch.stack(
            [u_xz * (1. - cos_) + u_y * sin_, u_yz * (1. - cos_) - u_x * sin_, cos_ + u_zz * (1. - cos_)], dim=0
        )
        # print(f"row_c: {row_c.size()}")
        
        ### rot_mtx for the rot_mtx ###
        rot_mtx = torch.stack(
            [row_a, row_b, row_c], dim=-1 ### rot_matrix of he matrix ##
        )
        
        return rot_mtx


def calibreate_urdf_files(urdf_fn):
    # active_robot =  parse_data_from_urdf(xml_fn)
    active_robot =  parse_data_from_urdf(urdf_fn)
    tot_joints = active_robot.tot_joints
    
    # class Joint_urdf: # 
    # def __init__(self, name, joint_type, parent_link, child_link, origin_xyz, axis_xyz, limit: Joint_Limit) -> None:
    #     self.name = name
    #     self.type = joint_type
    #     self.parent_link = parent_link
    #     self.child_link = child_link
    #     self.origin_xyz = origin_xyz
    #     self.axis_xyz = axis_xyz
    #     self.limit = limit
    
    with open(urdf_fn) as rf:
        urdf_string = rf.read()
    for cur_joint in tot_joints:
        print(f"type: {cur_joint.type}, origin: {cur_joint.origin_xyz}")
        cur_joint_origin = cur_joint.origin_xyz
        scaled_joint_origin = cur_joint_origin * 3.
        cur_joint_origin_string = cur_joint.origin_xyz_string
        if len(cur_joint_origin_string) == 0 or torch.sum(cur_joint_origin).item() == 0.:
            continue
            # <origin xyz="0.0 0.0 0.0"/>
        cur_joint_origin_string_wtag = "<origin xyz=" + "\"" + cur_joint_origin_string + "\"" + "/>"
        scaled_joint_origin_string_wtag = "<origin xyz=" + "\"" + f"{scaled_joint_origin[0].item()} {scaled_joint_origin[1].item()} {scaled_joint_origin[2].item()}" + "\"" + "/>"
        # scaled_joint_origin_string = f"{scaled_joint_origin[0].item()} {scaled_joint_origin[1].item()} {scaled_joint_origin[2].item()}"
        # urdf_string = urdf_string.replace(cur_joint_origin_string, scaled_joint_origin_string)
        urdf_string = urdf_string.replace(cur_joint_origin_string_wtag, scaled_joint_origin_string_wtag)
    changed_urdf_fn = urdf_fn.replace(".urdf", "_scaled.urdf")
    with open(changed_urdf_fn, "w") as wf:
        wf.write(urdf_string)
    print(f"changed_urdf_fn: {changed_urdf_fn}")
    # exit(0)
    
    
def get_GT_states_data_from_ckpt(ckpt_fn):
    mano_nn_substeps = 1
    num_steps = 60
    mano_robot_actions = nn.Embedding(
        num_embeddings=num_steps * mano_nn_substeps, embedding_dim=60,
    )
    torch.nn.init.zeros_(mano_robot_actions.weight)
    # params_to_train += list(self.robot_actions.parameters())
    
    mano_robot_delta_states = nn.Embedding(
        num_embeddings=num_steps * mano_nn_substeps, embedding_dim=60,
    )
    torch.nn.init.zeros_(mano_robot_delta_states.weight)
    # params_to_train += list(self.robot_delta_states.parameters())
    
    mano_robot_init_states = nn.Embedding(
        num_embeddings=1, embedding_dim=60,
    )
    torch.nn.init.zeros_(mano_robot_init_states.weight)
    # params_to_train += list(self.robot_init_states.parameters())
    
    mano_robot_glb_rotation = nn.Embedding(
        num_embeddings=num_steps * mano_nn_substeps, embedding_dim=4
    )
    mano_robot_glb_rotation.weight.data[:, 0] = 1.
    mano_robot_glb_rotation.weight.data[:, 1:] = 0.
    # params_to_train += list(self.robot_glb_rotation.parameters())
    
    
    mano_robot_glb_trans = nn.Embedding(
        num_embeddings=num_steps * mano_nn_substeps, embedding_dim=3
    )
    torch.nn.init.zeros_(mano_robot_glb_trans.weight)
    # params_to_train += list(self.robot_glb_trans.parameters())   
    
    mano_robot_states = nn.Embedding(
        num_embeddings=num_steps * mano_nn_substeps, embedding_dim=60,
    )
    torch.nn.init.zeros_(mano_robot_states.weight)
    mano_robot_states.weight.data[0, :] = mano_robot_init_states.weight.data[0, :].clone()
    
    
    ''' Load optimized MANO hand actions and states '''
    # ### laod optimized init actions #### # 
    # if 'model.load_optimized_init_actions' in self.conf and len(self.conf['model.load_optimized_init_actions']) > 0: 
    #     print(f"[MANO] Loading optimized init transformations from {self.conf['model.load_optimized_init_actions']}")
    cur_optimized_init_actions_fn = ckpt_fn
    optimized_init_actions_ckpt = torch.load(cur_optimized_init_actions_fn, map_location='cpu', )
    
    if 'mano_robot_states' in optimized_init_actions_ckpt:
        mano_robot_states.load_state_dict(optimized_init_actions_ckpt['mano_robot_states'])
    
    if 'mano_robot_init_states' in optimized_init_actions_ckpt:
        mano_robot_init_states.load_state_dict(optimized_init_actions_ckpt['mano_robot_init_states'])
        
    if 'mano_robot_glb_rotation' in optimized_init_actions_ckpt:
        mano_robot_glb_rotation.load_state_dict(optimized_init_actions_ckpt['mano_robot_glb_rotation'])
    
    if 'mano_robot_glb_trans' in optimized_init_actions_ckpt: # mano_robot_glb_trans
        mano_robot_glb_trans.load_state_dict(optimized_init_actions_ckpt['mano_robot_glb_trans'])
    
    mano_glb_trans_np_data = mano_robot_glb_trans.weight.data.detach().cpu().numpy()
    mano_glb_rotation_np_data = mano_robot_glb_rotation.weight.data.detach().cpu().numpy()
    mano_states_np_data = mano_robot_states.weight.data.detach().cpu().numpy()
    
    if optimized_init_actions_ckpt is not None and 'object_transl' in optimized_init_actions_ckpt:
        object_transl = optimized_init_actions_ckpt['object_transl'].detach().cpu().numpy()
        object_global_orient = optimized_init_actions_ckpt['object_global_orient'].detach().cpu().numpy()
    
    print(mano_robot_states.weight.data[1])
    
    #### TODO: add an arg to control where to save the gt-reference-data ####
    sv_gt_refereces = {
        'mano_glb_rot': mano_glb_rotation_np_data,
        'mano_glb_trans': mano_glb_trans_np_data,
        'mano_states': mano_states_np_data,
        'obj_rot': object_global_orient,
        'obj_trans': object_transl
    }
    # sv_gt_refereces_fn = "/home/xueyi/diffsim/Control-VAE/Data/ReferenceData/grab_train_split_20_cube_data.npy"
    # sv_gt_refereces_fn = "/home/xueyi/diffsim/Control-VAE/Data/ReferenceData/grab_train_split_25_ball_data.npy"
    # sv_gt_refereces_fn = "/home/xueyi/diffsim/Control-VAE/Data/ReferenceData/grab_train_split_54_cylinder_data.npy"
    sv_gt_refereces_fn = "/home/xueyi/diffsim/Control-VAE/Data/ReferenceData/grab_train_split_1_dingshuji_data.npy"
    np.save(sv_gt_refereces_fn, sv_gt_refereces)
    print(f'gt reference data saved to {sv_gt_refereces_fn}')
    #### TODO: add an arg to control where to save the gt-reference-data ####

def scale_and_save_meshes(meshes_folder):
    minn_robo_pts = -0.1
    maxx_robo_pts = 0.2
    extent_robo_pts = maxx_robo_pts - minn_robo_pts
    mult_const_after_cent = 0.437551664260203 ## should modify
    
    mult_const_after_cent = mult_const_after_cent / 3. * 0.9507
    
    meshes_fn = os.listdir(meshes_folder)
    meshes_fn = [fn for fn in meshes_fn if fn.endswith(".obj") and "scaled" not in fn]
    for cur_fn in meshes_fn:
        cur_mesh_name = cur_fn.split(".")[0]
        print(f"cur_mesh_name: {cur_mesh_name}")
        scaled_mesh_name = cur_mesh_name + "_scaled_bullet.obj"
        full_mesh_fn = os.path.join(meshes_folder, cur_fn)
        scaled_mesh_fn = os.path.join(meshes_folder, scaled_mesh_name)
        try:
            cur_mesh = trimesh.load_mesh(full_mesh_fn)
        except:
            continue
        cur_mesh.vertices = cur_mesh.vertices
        
        if 'palm' in cur_mesh_name:
            cur_mesh.vertices = (cur_mesh.vertices - minn_robo_pts) / extent_robo_pts
            cur_mesh.vertices = cur_mesh.vertices * 2. -1.
            cur_mesh.vertices = cur_mesh.vertices * mult_const_after_cent # mult_const #
        else:
            cur_mesh.vertices = (cur_mesh.vertices) / extent_robo_pts
            cur_mesh.vertices = cur_mesh.vertices * 2. # -1.
            cur_mesh.vertices = cur_mesh.vertices * mult_const_after_cent # mult_const #
        
        cur_mesh.export(scaled_mesh_fn)
        print(f"scaled_mesh_fn: {scaled_mesh_fn}")
    exit(0)
    
def scale_and_save_meshes_v2(meshes_folder):
    # /home/xueyi/diffsim/NeuS/rsc/redmax_hand/meshes/hand/body0_centered_scaled_v2.obj
    for body_idx in range(0, 18):
        cur_body_mesh_fn = f"body{body_idx}_centered_scaled_v2.obj"
        cur_body_mesh_fn = os.path.join(meshes_folder, cur_body_mesh_fn)
        cur_body_rescaled_mesh_fn = f"body{body_idx}_centered_scaled_v2_rescaled_grab.obj"
        cur_body_rescaled_mesh_fn = os.path.join(meshes_folder, cur_body_rescaled_mesh_fn)
        cur_mesh = trimesh.load_mesh(cur_body_mesh_fn)
        
        cur_mesh.vertices = cur_mesh.vertices / 4.0
        cur_mesh.export(cur_body_rescaled_mesh_fn)
        
    # minn_robo_pts = -0.1
    # maxx_robo_pts = 0.2
    # extent_robo_pts = maxx_robo_pts - minn_robo_pts
    # mult_const_after_cent = 0.437551664260203 ## should modify
    
    # mult_const_after_cent = mult_const_after_cent / 3. * 0.9507
    
    # meshes_fn = os.listdir(meshes_folder)
    # meshes_fn = [fn for fn in meshes_fn if fn.endswith(".obj") and "scaled" not in fn]
    # for cur_fn in meshes_fn:
    #     cur_mesh_name = cur_fn.split(".")[0]
    #     print(f"cur_mesh_name: {cur_mesh_name}")
    #     scaled_mesh_name = cur_mesh_name + "_scaled_bullet.obj"
    #     full_mesh_fn = os.path.join(meshes_folder, cur_fn)
    #     scaled_mesh_fn = os.path.join(meshes_folder, scaled_mesh_name)
    #     try:
    #         cur_mesh = trimesh.load_mesh(full_mesh_fn)
    #     except:
    #         continue
    #     cur_mesh.vertices = cur_mesh.vertices
        
    #     if 'palm' in cur_mesh_name:
    #         cur_mesh.vertices = (cur_mesh.vertices - minn_robo_pts) / extent_robo_pts
    #         cur_mesh.vertices = cur_mesh.vertices * 2. -1.
    #         cur_mesh.vertices = cur_mesh.vertices * mult_const_after_cent # mult_const #
    #     else:
    #         cur_mesh.vertices = (cur_mesh.vertices) / extent_robo_pts
    #         cur_mesh.vertices = cur_mesh.vertices * 2. # -1.
    #         cur_mesh.vertices = cur_mesh.vertices * mult_const_after_cent # mult_const #
        
    #     cur_mesh.export(scaled_mesh_fn)
    #     print(f"scaled_mesh_fn: {scaled_mesh_fn}")
    exit(0)
            
            
def calibreate_urdf_files_v2(urdf_fn):
    # active_robot =  parse_data_from_urdf(xml_fn)
    active_robot =  parse_data_from_urdf(urdf_fn)
    tot_joints = active_robot.tot_joints
    tot_links = active_robot.link_name_to_link_struct

    minn_robo_pts = -0.1
    maxx_robo_pts = 0.2
    extent_robo_pts = maxx_robo_pts - minn_robo_pts
    mult_const_after_cent = 0.437551664260203 ## should modify
    
    mult_const_after_cent = mult_const_after_cent / 3. * 0.9507
    
    with open(urdf_fn) as rf:
        urdf_string = rf.read()
    for cur_joint in tot_joints:
        # print(f"type: {cur_joint.type}, origin: {cur_joint.origin_xyz}")
        cur_joint_origin = cur_joint.origin_xyz
        
        # cur_joint_origin = (cur_joint_origin / extent_robo_pts) * 2.0 * mult_const_after_cent
        
        # cur_joint_origin = (cur_joint_origin / extent_robo_pts) * 2.0 * mult_const_after_cent
        
        if cur_joint.name in ['FFJ4' , 'MFJ4' ,'RFJ4' ,'LFJ5' ,'THJ5']:
            cur_joint_origin = (cur_joint_origin - minn_robo_pts) / extent_robo_pts
            cur_joint_origin = cur_joint_origin * 2.0 - 1.0
            cur_joint_origin = cur_joint_origin * mult_const_after_cent
        else:
            cur_joint_origin = (cur_joint_origin) / extent_robo_pts
            cur_joint_origin = cur_joint_origin * 2.0 # - 1.0
            cur_joint_origin = cur_joint_origin * mult_const_after_cent
            
        
        origin_list = cur_joint_origin.detach().cpu().tolist()
        origin_list = [str(cur_val) for cur_val in origin_list]
        origin_str = " ".join(origin_list)
        print(f"name: {cur_joint.name}, cur_joint_origin: {origin_str}")
        
        # scaled_joint_origin = cur_joint_origin * 3.
        # cur_joint_origin_string = cur_joint.origin_xyz_string
        # if len(cur_joint_origin_string) == 0 or torch.sum(cur_joint_origin).item() == 0.:
        #     continue
        #     # <origin xyz="0.0 0.0 0.0"/>
        # cur_joint_origin_string_wtag = "<origin xyz=" + "\"" + cur_joint_origin_string + "\"" + "/>"
        # scaled_joint_origin_string_wtag = "<origin xyz=" + "\"" + f"{scaled_joint_origin[0].item()} {scaled_joint_origin[1].item()} {scaled_joint_origin[2].item()}" + "\"" + "/>"
        # # scaled_joint_origin_string = f"{scaled_joint_origin[0].item()} {scaled_joint_origin[1].item()} {scaled_joint_origin[2].item()}"
        # # urdf_string = urdf_string.replace(cur_joint_origin_string, scaled_joint_origin_string)
        # urdf_string = urdf_string.replace(cur_joint_origin_string_wtag, scaled_joint_origin_string_wtag)
    # changed_urdf_fn = urdf_fn.replace(".urdf", "_scaled.urdf")
    # with open(changed_urdf_fn, "w") as wf:
    #     wf.write(urdf_string)
    # print(f"changed_urdf_fn: {changed_urdf_fn}")
    # # exit(0)
    
    # for cur_link_nm in tot_links:
    #     cur_link = tot_links[cur_link_nm]
    #     if cur_link.visual is None:
    #         continue
    #     xyz_visual = cur_link.visual.visual_xyz
    #     xyz_visual = (xyz_visual / extent_robo_pts) * 2.0 * mult_const_after_cent
    #     xyz_visual_list = xyz_visual.detach().cpu().tolist()
    #     xyz_visual_list = [str(cur_val) for cur_val in xyz_visual_list]
    #     xyz_visual_str = " ".join(xyz_visual_list)
    #     print(f"name: {cur_link.name}, xyz_visual: {xyz_visual_str}")
    
def get_shadow_GT_states_data_from_ckpt(ckpt_fn):
    mano_nn_substeps = 1
    num_steps = 60
    
    
    # robot actions # # 
    # robot_actions = nn.Embedding(
    #     num_embeddings=num_steps, embedding_dim=22,
    # )
    # 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,
    # )
    # 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,
    )
    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,
    # )
    # 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
    )
    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
    )
    torch.nn.init.zeros_(robot_glb_trans.weight)
        
    ''' Load optimized MANO hand actions and states '''
    cur_optimized_init_actions_fn = ckpt_fn
    optimized_init_actions_ckpt = torch.load(cur_optimized_init_actions_fn, map_location='cpu', )
    
    print(f"optimized_init_actions_ckpt: {optimized_init_actions_ckpt.keys()}")
    
    if 'robot_glb_rotation' in optimized_init_actions_ckpt:
        robot_glb_rotation.load_state_dict(optimized_init_actions_ckpt['robot_glb_rotation'])
    
    if 'robot_states' in optimized_init_actions_ckpt:
        robot_states.load_state_dict(optimized_init_actions_ckpt['robot_states'])
        
    if 'robot_glb_trans' in optimized_init_actions_ckpt:
        robot_glb_trans.load_state_dict(optimized_init_actions_ckpt['robot_glb_trans'])
    
    # if 'mano_robot_glb_trans' in optimized_init_actions_ckpt: # mano_robot_glb_trans
    #     mano_robot_glb_trans.load_state_dict(optimized_init_actions_ckpt['mano_robot_glb_trans'])
    
    robot_glb_trans_np_data = robot_glb_trans.weight.data.detach().cpu().numpy()
    robot_glb_rotation_np_data = robot_glb_rotation.weight.data.detach().cpu().numpy()
    robot_states_np_data = robot_states.weight.data.detach().cpu().numpy()
    
    if optimized_init_actions_ckpt is not None and 'object_transl' in optimized_init_actions_ckpt:
        object_transl = optimized_init_actions_ckpt['object_transl'].detach().cpu().numpy()
        object_global_orient = optimized_init_actions_ckpt['object_global_orient'].detach().cpu().numpy()
    
    # print(mano_robot_states.weight.data[1])
    
    #### TODO: add an arg to control where to save the gt-reference-data ####
    sv_gt_refereces = {
        'mano_glb_rot': robot_glb_rotation_np_data,
        'mano_glb_trans': robot_glb_trans_np_data,
        'mano_states': robot_states_np_data,
        'obj_rot': object_global_orient,
        'obj_trans': object_transl
    }
    # sv_gt_refereces_fn = "/home/xueyi/diffsim/Control-VAE/Data/ReferenceData/grab_train_split_20_cube_data.npy"
    # sv_gt_refereces_fn = "/home/xueyi/diffsim/Control-VAE/Data/ReferenceData/grab_train_split_25_ball_data.npy"
    # sv_gt_refereces_fn = "/home/xueyi/diffsim/Control-VAE/Data/ReferenceData/grab_train_split_54_cylinder_data.npy"
    # sv_gt_refereces_fn = "/home/xueyi/diffsim/Control-VAE/Data/ReferenceData/shadow_grab_train_split_224_tiantianquan_data.npy"
    sv_gt_refereces_fn = "/home/xueyi/diffsim/Control-VAE/Data/ReferenceData/shadow_grab_train_split_54_cylinder_data.npy"
    np.save(sv_gt_refereces_fn, sv_gt_refereces)
    print(f'gt reference data saved to {sv_gt_refereces_fn}')
    #### TODO: add an arg to control where to save the gt-reference-data ####

## saved the robot file ##


def calibrate_left_shadow_hand():
    rgt_shadow_hand_des_folder = "/home/xueyi/diffsim/NeuS/rsc/shadow_hand_description"
    lft_shadow_hand_des_folder = "/home/xueyi/diffsim/NeuS/rsc/shadow_hand_description_left"
    os.makedirs(lft_shadow_hand_des_folder, exist_ok=True)
    lft_shadow_hand_mesh_folder = os.path.join(lft_shadow_hand_des_folder, "meshes")
    os.makedirs(lft_shadow_hand_mesh_folder, exist_ok=True)
    rgt_shadow_hand_mesh_folder = os.path.join(rgt_shadow_hand_des_folder, "meshes")
    tot_rgt_hand_meshes = os.listdir(rgt_shadow_hand_mesh_folder)
    tot_rgt_hand_meshes = [fn for fn in tot_rgt_hand_meshes if fn.endswith(".obj")]
    for cur_hand_mesh_fn in tot_rgt_hand_meshes:
        full_rgt_mesh_fn = os.path.join(rgt_shadow_hand_mesh_folder, cur_hand_mesh_fn)
        try:
            full_rgt_mesh = trimesh.load(full_rgt_mesh_fn, force='mesh')
        except:
            continue
        full_rgt_mesh_verts = full_rgt_mesh.vertices
        full_rgt_mesh_faces = full_rgt_mesh.faces
        full_rgt_mesh_verts[:, 1] = -1. * full_rgt_mesh_verts[:, 1] ## flip the y-axis 
        lft_mesh = trimesh.Trimesh(vertices=full_rgt_mesh_verts, faces=full_rgt_mesh_faces)
        lft_mesh_fn = os.path.join(lft_shadow_hand_mesh_folder, cur_hand_mesh_fn)
        lft_mesh.export(lft_mesh_fn)
        print(f"lft_mesh_fn: {lft_mesh_fn}")
    exit(0)


## urd for the left hand 
def calibreate_urdf_files_left_hand(urdf_fn):
    # active_robot =  parse_data_from_urdf(xml_fn)
    active_robot =  parse_data_from_urdf(urdf_fn)
    tot_joints = active_robot.tot_joints
    tot_links = active_robot.link_name_to_link_struct

    minn_robo_pts = -0.1
    maxx_robo_pts = 0.2
    extent_robo_pts = maxx_robo_pts - minn_robo_pts
    mult_const_after_cent = 0.437551664260203 ## should modify
    
    mult_const_after_cent = mult_const_after_cent / 3. * 0.9507
    
    with open(urdf_fn) as rf:
        urdf_string = rf.read()
    for cur_joint in tot_joints:
        # print(f"type: {cur_joint.type}, origin: {cur_joint.origin_xyz}")
        cur_joint_origin = cur_joint.origin_xyz
        
        cur_joint_axis = cur_joint.axis_xyz
        
        cur_joint_origin = cur_joint_origin.detach()
        cur_joint_axis = cur_joint_axis.detach()
        
        cur_joint_origin[1] = -1.0 * cur_joint_origin[1]
        cur_joint_axis[1] = -1.0 * cur_joint_axis[1]
        
        
        origin_list = cur_joint_origin.detach().cpu().tolist()
        origin_list = [str(cur_val) for cur_val in origin_list]
        origin_str = " ".join(origin_list)
        
        axis_list = cur_joint_axis.detach().cpu().tolist()
        axis_list = [str(cur_val) for cur_val in axis_list]
        axis_str = " ".join(axis_list)
        print(f"name: {cur_joint.name}, cur_joint_origin: {origin_str}, axis_str: {axis_str}")
        
        # cur_joint_origin = (cur_joint_origin / extent_robo_pts) * 2.0 * mult_const_after_cent
        
        # cur_joint_origin = (cur_joint_origin / extent_robo_pts) * 2.0 * mult_const_after_cent
        
        # if cur_joint.name in ['FFJ4' , 'MFJ4' ,'RFJ4' ,'LFJ5' ,'THJ5']:
        #     cur_joint_origin = (cur_joint_origin - minn_robo_pts) / extent_robo_pts
        #     cur_joint_origin = cur_joint_origin * 2.0 - 1.0
        #     cur_joint_origin = cur_joint_origin * mult_const_after_cent
        # else:
        #     cur_joint_origin = (cur_joint_origin) / extent_robo_pts
        #     cur_joint_origin = cur_joint_origin * 2.0 # - 1.0
        #     cur_joint_origin = cur_joint_origin * mult_const_after_cent
            
        
        # origin_list = cur_joint_origin.detach().cpu().tolist()
        # origin_list = [str(cur_val) for cur_val in origin_list]
        # origin_str = " ".join(origin_list)
        # print(f"name: {cur_joint.name}, cur_joint_origin: {origin_str}")

def calibreate_urdf_files_v4(urdf_fn, dst_urdf_fn):
    # active_robot =  parse_data_from_urdf(xml_fn)
    active_robot =  parse_data_from_urdf(urdf_fn)
    tot_joints = active_robot.tot_joints
    tot_links = active_robot.link_name_to_link_struct

    # minn_robo_pts = -0.1
    # maxx_robo_pts = 0.2
    # extent_robo_pts = maxx_robo_pts - minn_robo_pts
    # mult_const_after_cent = 0.437551664260203 ## should modify
    
    # mult_const_after_cent = mult_const_after_cent / 3. * 0.9507
    
    with open(urdf_fn) as rf:
        urdf_string = rf.read()
    for cur_joint in tot_joints:
        # print(f"type: {cur_joint.type}, origin: {cur_joint.origin_xyz}")
        cur_joint_origin = cur_joint.origin_xyz
        modified_joint_origin = cur_joint_origin / 4.
        
        origin_list = cur_joint_origin.detach().cpu().tolist()
        origin_list = [str(cur_val) for cur_val in origin_list]
        origin_str = " ".join(origin_list)
        
        dst_list = modified_joint_origin.detach().cpu().tolist()
        dst_list = [str(cur_val) for cur_val in dst_list]
        dst_str = " ".join(dst_list)
        
        urdf_string = urdf_string.replace(origin_str, dst_str)
    
    with open(dst_urdf_fn, "w") as wf:
        wf.write(urdf_string)
        wf.close()
        # # cur_joint_origin = (cur_joint_origin / extent_robo_pts) * 2.0 * mult_const_after_cent
        
        # # cur_joint_origin = (cur_joint_origin / extent_robo_pts) * 2.0 * mult_const_after_cent
        
        # if cur_joint.name in ['FFJ4' , 'MFJ4' ,'RFJ4' ,'LFJ5' ,'THJ5']:
        #     cur_joint_origin = (cur_joint_origin - minn_robo_pts) / extent_robo_pts
        #     cur_joint_origin = cur_joint_origin * 2.0 - 1.0
        #     cur_joint_origin = cur_joint_origin * mult_const_after_cent
        # else:
        #     cur_joint_origin = (cur_joint_origin) / extent_robo_pts
        #     cur_joint_origin = cur_joint_origin * 2.0 # - 1.0
        #     cur_joint_origin = cur_joint_origin * mult_const_after_cent
            
        
        # origin_list = cur_joint_origin.detach().cpu().tolist()
        # origin_list = [str(cur_val) for cur_val in origin_list]
        # origin_str = " ".join(origin_list)
        # print(f"name: {cur_joint.name}, cur_joint_origin: {origin_str}")
        

def test_gt_ref_data(gt_ref_data_fn):
    cur_gt_ref_data = np.load(gt_ref_data_fn, allow_pickle=True).item()
    print(cur_gt_ref_data.keys())
    
    mano_glb_rot, glb_trans, states = cur_gt_ref_data['mano_glb_rot'], cur_gt_ref_data['mano_glb_trans'], cur_gt_ref_data['mano_states']
    return mano_glb_rot, glb_trans, states

    
def get_states(gt_ref_data_fn):
    states = np.load(gt_ref_data_fn, allow_pickle=True).item()
    return states['target']

#### Big TODO: the external contact forces from the manipulated object to the robot ####
if __name__=='__main__': # # #
    
    gt_ref_data_fn = "/home/xueyi/diffsim/Control-VAE/Data/ReferenceData/shadow_grab_train_split_85_bunny_wact_data.npy"
    # mano_glb_rot, glb_trans, states = test_gt_ref_data(gt_ref_data_fn)
    # eixt(0)
    mano_states_fn = '/home/xueyi/diffsim/NeuS/raw_data/evalulated_traj_sm_l512_wana_v3_subiters1024_optim_params_shadow_85_bunny_std0d01_netv1_mass10000_new_dp1d0_wtable_gn9d8__step_2.npy'
    mano_states_fn = '/home/xueyi/diffsim/NeuS/raw_data/evalulated_traj_sm_l512_wana_v3_subiters1024_optim_params_shadow_102_mouse_wact_std0d01_netv1_mass10000_new_dp1d0_dtv2tsv2ctlv2_netv3optt_lstd_langdamp_wcmase_ni4_wtable_gn_adp1d0_trwmonly_cs0d6_predtarmano_wambient__step_9.npy'
    mano_states = get_states(mano_states_fn)
    
    blended_ratio = 0.5
    
    blended_states = []
    
    tot_rot_mtxes = []
    tot_trans = []
    for i_state in range(len(mano_states)):
        cur_trans = mano_states[i_state][:3]
        cur_rot = mano_states[i_state][3:6]
        cur_states = mano_states[i_state][6:]
        
        cur_rot_struct = R.from_euler('zyx', cur_rot[[2, 1, 0]], degrees=False)
        cur_rot_mtx = cur_rot_struct.as_matrix()
        
        tot_rot_mtxes.append(cur_rot_mtx)
        tot_trans.append(cur_trans)
        
        
        cur_state = cur_states #  states[i_state]
        cur_modified_state = mano_states[0][6:] + (cur_state - mano_states[0][6:] ) * blended_ratio
        
        cur_modified_state = np.concatenate([np.zeros((2,), dtype=np.float32), cur_modified_state], axis=-1)
        blended_states.append(cur_modified_state)
        # return blended_states
    
    tot_rot_mtxes = np.stack(tot_rot_mtxes, axis=0)
    tot_trans = np.stack(tot_trans, axis=0)
    blended_states = np.stack(blended_states, axis=0)
    
    # urdf_fn = "/home/xueyi/diffsim/NeuS/rsc/redmax_hand/redmax_hand_test_3_wcollision.urdf"
    # dst_urdf_fn = "/home/xueyi/diffsim/NeuS/rsc/redmax_hand/redmax_hand_test_3_wcollision_rescaled_grab.urdf"
    # calibreate_urdf_files_v4(urdf_fn, dst_urdf_fn)
    # exit(0)
    
    # meshes_folder = "/home/xueyi/diffsim/NeuS/rsc/redmax_hand/meshes/hand"
    # scale_and_save_meshes_v2(meshes_folder)
    # exit(0)
    
    urdf_fn = "/home/xueyi/diffsim/NeuS/rsc/shadow_hand_description/shadowhand_new_scaled.urdf"
    urdf_fn = "/home/xueyi/diffsim/NeuS/rsc/shadow_hand_description/shadowhand_new_scaled_nroot_new.urdf"
    # urdf_fn = "/home/xueyi/diffsim/NeuS/rsc/shadow_hand_description/shadowhand_new_scaled_nroot.urdf"
    # urdf_fn = "/home/xueyi/diffsim/NeuS/rsc/redmax_hand/redmax_hand_test_3_wcollision_rescaled_grab.urdf"
    robot_agent = RobotAgent(urdf_fn)
    
    init_vertices, init_faces = robot_agent.active_robot.init_vertices, robot_agent.active_robot.init_faces
    init_vertices = init_vertices.detach().cpu().numpy()
    init_faces = init_faces.detach().cpu().numpy()
    
    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()
        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()
        
        cur_pts_transformed = np.matmul(
            tot_rot_mtxes[i_ts], cur_pts.T
        ).T + tot_trans[i_ts][None]
        tot_transformed_pts.append(cur_pts_transformed)
    tot_transformed_pts = np.stack(tot_transformed_pts, axis=0)
    np.save("/home/xueyi/diffsim/NeuS/raw_data/transformed_pts.npy", {'tot_transformed_pts': tot_transformed_pts, 'init_faces': init_faces})
    exit(0)
        
    
    robot_agent.active_robot.set_delta_state_and_update_v2()
    
    init_vertices, init_faces = robot_agent.active_robot.init_vertices, robot_agent.active_robot.init_faces
    init_vertices = init_vertices.detach().cpu().numpy()
    init_faces = init_faces.detach().cpu().numpy()
    print(f"init_vertices: {init_vertices.shape}, init_faces: {init_faces.shape}")
    shadow_hand_mesh = trimesh.Trimesh(vertices=init_vertices, faces=init_faces)
    # shadow_hand_sv_fn = "/home/xueyi/diffsim/NeuS/raw_data/shadow_hand_lft.obj"
    shadow_hand_sv_fn = "/home/xueyi/diffsim/NeuS/raw_data/shadow_hand_new.ply"
    shadow_hand_mesh.export(shadow_hand_sv_fn)
    np.save("/home/xueyi/diffsim/NeuS/raw_data/faces.npy", init_faces)
    
    exit(0)
    
    init_vertices, init_faces = robot_agent.active_robot.init_vertices, robot_agent.active_robot.init_faces
    init_vertices = init_vertices.detach().cpu().numpy()
    init_faces = init_faces.detach().cpu().numpy()
    print(f"init_vertices: {init_vertices.shape}, init_faces: {init_faces.shape}")
    shadow_hand_mesh = trimesh.Trimesh(vertices=init_vertices, faces=init_faces)
    # shadow_hand_sv_fn = "/home/xueyi/diffsim/NeuS/raw_data/shadow_hand_lft.obj"
    shadow_hand_sv_fn = "/home/xueyi/diffsim/NeuS/raw_data/scaled_shadow_hand.obj"
    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)
    robot_agent.set_initial_state(init_joint_states)
    
    
    cur_verts, cur_faces = robot_agent.get_init_visual_pts()
    cur_mesh = trimesh.Trimesh(vertices=cur_verts.detach().cpu().numpy(), faces=cur_faces.detach().cpu().numpy())
    shadow_hand_sv_fn = "/home/xueyi/diffsim/NeuS/raw_data/scaled_redmax_hand_rescaled_grab_wstate.obj"
    cur_mesh.export(shadow_hand_sv_fn)
    exit(0)
    
    
    
    urdf_fn = "/home/xueyi/diffsim/NeuS/rsc/shadow_hand_description/shadowhand_new_scaled.urdf"
    
    ## 
    lft_urdf_fn = "/home/xueyi/diffsim/NeuS/rsc/shadow_hand_description_left/shadowhand_left_new_scaled.urdf"
    
    
    urdf_fn = "/home/xueyi/diffsim/NeuS/rsc/redmax_hand/redmax_hand_test_3_wcollision.urdf"
    
    ## 
    lft_urdf_fn = "/home/xueyi/diffsim/NeuS/rsc/redmax_hand/redmax_hand_test_3_wcollision.urdf"
    
    robot_agent = RobotAgent(lft_urdf_fn)
    init_vertices, init_faces = robot_agent.active_robot.init_vertices, robot_agent.active_robot.init_faces
    init_vertices = init_vertices.detach().cpu().numpy()
    init_faces = init_faces.detach().cpu().numpy()
    print(f"init_vertices: {init_vertices.shape}, init_faces: {init_faces.shape}")
    shadow_hand_mesh = trimesh.Trimesh(vertices=init_vertices, faces=init_faces)
    shadow_hand_sv_fn = "/home/xueyi/diffsim/NeuS/raw_data/shadow_hand_lft.obj"
    shadow_hand_sv_fn = "/home/xueyi/diffsim/NeuS/raw_data/redmax_hand.obj"
    shadow_hand_mesh.export(shadow_hand_sv_fn)
    exit(0)
    
    
    rgt_urdf_fn = "/home/xueyi/diffsim/NeuS/rsc/shadow_hand_description/shadowhand_new_scaled.urdf"
    # rgt_urdf_fn
    calibreate_urdf_files_left_hand(rgt_urdf_fn)
    exit(0)
    
    calibrate_left_shadow_hand()
    exit(0)
    
    # ckpt_fn = "/data3/datasets/xueyi/neus/exp/hand_test_routine_2_light_color_wtime_active_passive/wmask_reverse_value_totviews_tag_train_retargeted_shadow_hand_states_/checkpoints/ckpt_320000.pth"
    # ckpt_fn = "/data3/datasets/xueyi/neus/exp/hand_test_routine_2_light_color_wtime_active_passive/wmask_reverse_value_totviews_tag_train_retargeted_shadow_hand_states_optrobot__seq_54_optrules_/checkpoints/ckpt_030000.pth"
    # get_shadow_GT_states_data_from_ckpt(ckpt_fn)
    # exit(0)
    
    # urdf_fn = "/home/xueyi/diffsim/NeuS/rsc/shadow_hand_description/shadowhand_new_scaled.urdf"
    # calibreate_urdf_files_v2(urdf_fn)
    # exit(0)
    
    meshes_folder = "/home/xueyi/diffsim/NeuS/rsc/shadow_hand_description/meshes"
    # scale_and_save_meshes(meshes_folder)
    # exit(0)
    
    # sv_ckpt_fn = "/data3/datasets/xueyi/neus/exp/hand_test_routine_2_light_color_wtime_active_passive/wmask_reverse_value_totviews_tag_train_mano_states_grab_train_54_cylinder_tst_/checkpoints/ckpt_070000.pth"
    # sv_ckpt_fn = "/data3/datasets/xueyi/neus/exp/hand_test_routine_2_light_color_wtime_active_passive/wmask_reverse_value_totviews_tag_train_mano_states_grab_train_1_dingshuji_tst_/checkpoints/ckpt_070000.pth"
    # get_GT_states_data_from_ckpt(sv_ckpt_fn)
    # exit(0)
    
    
    # urdf_fn = "/home/xueyi/diffsim/NeuS/rsc/mano/mano_mean_wcollision_scaled_scaled_0_9507_nroot.urdf"
    # robot_agent = RobotAgent(urdf_fn)
    # exit(0)
    
    # urdf_fn = "/home/xueyi/diffsim/NeuS/rsc/mano/mano_mean_nocoll_simplified.urdf"
    # urdf_fn = "/home/xueyi/diffsim/NeuS/rsc/mano/mano_mean_wcollision_scaled.urdf"
    # calibreate_urdf_files(urdf_fn)
    # exit(0)
    
    # urdf_fn = "/home/xueyi/diffsim/NeuS/rsc/mano/mano_mean_nocoll_simplified.urdf"
    urdf_fn = "/home/xueyi/diffsim/NeuS/rsc/shadow_hand_description/shadowhand_new.urdf"
    robot_agent = RobotAgent(urdf_fn)
    
    init_vertices, init_faces = robot_agent.active_robot.init_vertices, robot_agent.active_robot.init_faces
    init_vertices = init_vertices.detach().cpu().numpy()
    init_faces = init_faces.detach().cpu().numpy()
    
    shadow_hand_mesh = trimesh.Trimesh(vertices=init_vertices, faces=init_faces)
    shadow_hand_sv_fn = "/home/xueyi/diffsim/NeuS/raw_data/shadow_hand.obj"
    shadow_hand_mesh.export(shadow_hand_sv_fn)
    exit(0)
    
    
    ref_dict_npy = "reference_verts.npy"
    robot_agent.initialize_optimization(ref_dict_npy)
    ts_to_robot_points, ts_to_ref_points = robot_agent.forward_stepping_optimization()
    np.save(f"ts_to_robot_points.npy", ts_to_robot_points)
    np.save(f"ts_to_ref_points.npy", ts_to_ref_points)
    exit(0)
    
    urdf_fn = "/home/xueyi/diffsim/NeuS/rsc/mano/mano_mean_nocoll_simplified.urdf"
    cur_robot = parse_data_from_urdf(urdf_fn)
    # self.init_vertices, self.init_faces
    init_vertices, init_faces = cur_robot.init_vertices, cur_robot.init_faces
    
    
    
    init_vertices = init_vertices.detach().cpu().numpy()
    init_faces = init_faces.detach().cpu().numpy()
    
    
    ## initial  states ehre ##3
    # mesh_obj = trimesh.Trimesh(vertices=init_vertices, faces=init_faces)
    # mesh_obj.export(f"hand_urdf.ply")
    
    ##### Test the set initial state function #####
    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)
    cur_ones_actions = torch.ones((60, ), dtype=torch.float32) #  * 100
    
    ts_to_mesh_verts = {}
    for i_ts in range(50):
        cur_robot.calculate_inertia()
    
        cur_robot.set_actions_and_update_states(cur_ones_actions, i_ts, 0.2) ###
    
    
        cur_verts, cur_faces = cur_robot.get_init_visual_pts()
        cur_mesh = trimesh.Trimesh(vertices=cur_verts.detach().cpu().numpy(), faces=cur_faces.detach().cpu().numpy())
        
        ts_to_mesh_verts[i_ts + i_ts] = cur_verts.detach().cpu().numpy()
        # cur_mesh.export(f"stated_mano_mesh.ply")
        # cur_mesh.export(f"zero_actioned_mano_mesh.ply")
        cur_mesh.export(f"ones_actioned_mano_mesh_ts_{i_ts}.ply")
    
    np.save(f"reference_verts.npy", ts_to_mesh_verts)
    
    exit(0)
    
    xml_fn = "/home/xueyi/diffsim/DiffHand/assets/hand_sphere.xml"
    robot_agent = RobotAgent(xml_fn=xml_fn, args=None)
    init_visual_pts = robot_agent.init_visual_pts.detach().cpu().numpy()
    exit(0)