a

.gitignore

@@ -0,0 +1,193 @@
+# Created by .ignore support plugin (hsz.mobi)
+### Python template
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+*.npy
+
+*.zip
+
+*.ply
+
+
+#### old files ####
+exp_runner_arti_*
+
+exp_runner.py
+exp_runner_arti.py
+exp_runner_dyn_model.py
+exp_runner_sim.py
+get-pip.py
+
+test.py
+
+
+#### old scripts and data&exp folders ###
+scripts/
+confs/
+data/
+ckpts/
+exp/
+uni_rep/
+
+*/*_local.sh
+*/*_local.conf
+
+# Distribution / packaging
+.Python
+env/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+*.egg-info/
+.installed.cfg
+*.egg
+
+rsc/
+raw_data/
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*,cover
+.hypothesis/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# IPython Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# dotenv
+.env
+
+# virtualenv
+venv/
+ENV/
+
+# Spyder project settings
+.spyderproject
+
+# Rope project settings
+.ropeproject
+### VirtualEnv template
+# Virtualenv
+# http://iamzed.com/2009/05/07/a-primer-on-virtualenv/
+.Python
+[Bb]in
+[Ii]nclude
+[Ll]ib
+[Ll]ib64
+[Ll]ocal
+# [Ss]cripts
+pyvenv.cfg
+.venv
+pip-selfcheck.json
+### JetBrains template
+# Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio and Webstorm
+# Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839
+
+# User-specific stuff:
+.idea/workspace.xml
+.idea/tasks.xml
+.idea/dictionaries
+.idea/vcs.xml
+.idea/jsLibraryMappings.xml
+
+# Sensitive or high-churn files:
+.idea/dataSources.ids
+.idea/dataSources.xml
+.idea/dataSources.local.xml
+.idea/sqlDataSources.xml
+.idea/dynamic.xml
+.idea/uiDesigner.xml
+
+# Gradle:
+.idea/gradle.xml
+.idea/libraries
+
+# Mongo Explorer plugin:
+.idea/mongoSettings.xml
+
+.idea/
+
+## File-based project format:
+*.iws
+
+
+## Plugin-specific files:
+
+# IntelliJ
+/out/
+
+# mpeltonen/sbt-idea plugin
+.idea_modules/
+
+# JIRA plugin
+atlassian-ide-plugin.xml
+
+# Crashlytics plugin (for Android Studio and IntelliJ)
+com_crashlytics_export_strings.xml
+crashlytics.properties
+crashlytics-build.properties
+fabric.properties
+
+data
+public_data
+exp
+tmp
+
+.models/*.npy
+.models/*.ply

README.md

@@ -1,8 +1,8 @@
 ---
 title: Quasi Physical Sims
-emoji: 📚
-colorFrom: red
-colorTo: gray
+emoji: 🏃
+colorFrom: gray
+colorTo: indigo
 sdk: gradio
 sdk_version: 4.25.0
 app_file: app.py

confs_new/dyn_grab_arti_shadow_dm.conf

@@ -0,0 +1,288 @@
+general {
+
+    base_exp_dir = exp/CASE_NAME/wmask
+
+    tag = "train_retargeted_shadow_hand_seq_102_diffhand_model_"
+
+    recording = [
+        ./,
+        ./models
+    ]
+}
+
+
+dataset {
+    data_dir = public_data/CASE_NAME/
+    render_cameras_name = cameras_sphere.npz
+    object_cameras_name = cameras_sphere.npz
+
+    obj_idx = 102
+}
+
+train {
+    learning_rate = 5e-4
+    # learning_rate = 5e-6
+    learning_rate_actions = 5e-6
+    learning_rate_alpha = 0.05
+    end_iter = 300000
+
+    # batch_size = 128 #  64
+    # batch_size = 4000
+    # batch_size = 3096 #  64
+    batch_size = 1024
+    validate_resolution_level = 4
+    warm_up_end = 5000
+    anneal_end = 0
+    use_white_bkgd = False
+
+    # save_freq = 10000
+    save_freq = 10000
+    val_freq = 20 #  2500
+    val_mesh_freq = 20 #  5000
+    report_freq = 10
+    ### igr weight ###
+    igr_weight = 0.1
+    mask_weight = 0.1
+}
+ 
+model {
+    
+
+    load_redmax_robot_actions = ""
+    penetrating_depth_penalty = 0.0
+    train_states = True
+
+    minn_dist_threshold = 0.000
+    obj_mass = 100.0
+    obj_mass = 30.0
+
+    # use_mano_hand_for_test = False
+    use_mano_hand_for_test = True
+
+    # train_residual_friction = False
+    train_residual_friction = True
+
+    # use_LBFGS = True
+    use_LBFGS = False
+
+
+    use_mano_hand_for_test = False
+    train_residual_friction = True 
+    
+    extract_delta_mesh = False
+    freeze_weights = True
+    # gt_act_xs_def = True
+    gt_act_xs_def = False
+    use_bending_network =  True
+    use_delta_bending = True
+
+    use_passive_nets = True
+    use_split_network = True
+
+    n_timesteps = 60
+
+
+    # using_delta_glb_trans = True
+    using_delta_glb_trans = False
+    # train_multi_seqs = True
+
+
+    # optimize_with_intermediates = False
+    optimize_with_intermediates = True
+
+
+    loss_tangential_diff_coef = 1000
+    loss_tangential_diff_coef = 0
+
+    
+
+    optimize_active_object = True
+
+    no_friction_constraint = False
+
+    optimize_glb_transformations = True
+
+
+    sim_model_path = "DiffHand/assets/hand_sphere_only_hand_testt.xml"
+    mano_sim_model_path = "rsc/mano/mano_mean_wcollision_scaled_scaled_0_9507_nroot.urdf"
+    mano_mult_const_after_cent = 1.0
+    sim_num_steps = 1000000
+
+    bending_net_type = "active_force_field_v18"    
+
+
+    ### try to train the residual friction ? ###
+    train_residual_friction = True
+    optimize_rules = True
+
+    load_optimized_init_actions = ""
+
+
+    use_optimizable_params = True
+
+
+    ### grab train seq 224 ###
+    penetration_determining = "sdf_of_canon"
+    train_with_forces_to_active = False
+    loss_scale_coef = 1000.0
+    # penetration_proj_k_to_robot_friction = 40000000.0
+    # penetration_proj_k_to_robot_friction = 100000000.0 # as friction coefs here #
+    use_same_contact_spring_k = False
+
+
+    # sim_model_path = "/home/xueyi/diffsim/DiffHand/assets/hand_sphere_only_hand_testt.xml"
+    sim_model_path = "rsc/shadow_hand_description/shadowhand_new.urdf"
+
+
+    optimize_rules = False
+
+    penetration_determining = "sdf_of_canon"
+
+    optimize_rules = False
+
+    optim_sim_model_params_from_mano = False
+    optimize_rules = False    
+
+
+    penetration_proj_k_to_robot_friction = 100000000.0 # as friction coefs here # ## confs ##
+    penetration_proj_k_to_robot = 40000000.0 #
+
+
+    penetrating_depth_penalty = 1
+
+    minn_dist_threshold_robot_to_obj = 0.0
+
+
+    minn_dist_threshold_robot_to_obj = 0.1
+
+    optim_sim_model_params_from_mano = False
+    optimize_rules = False    
+
+    optim_sim_model_params_from_mano = True
+    optimize_rules = True    
+    minn_dist_threshold_robot_to_obj = 0.0
+
+    optim_sim_model_params_from_mano = False
+    optimize_rules = False    
+    minn_dist_threshold_robot_to_obj = 0.1
+
+
+    ### kinematics confgs ###
+    obj_sdf_fn = "data/grab/102/102_obj.npy" 
+    kinematic_mano_gt_sv_fn = "data/grab/102/102_sv_dict.npy"
+    scaled_obj_mesh_fn = "data/grab/102/102_obj.obj"
+    # ckpt_fn = ""
+    # load_optimized_init_transformations = ""
+    optim_sim_model_params_from_mano = True
+    optimize_rules = True    
+    minn_dist_threshold_robot_to_obj = 0.0
+
+    optim_sim_model_params_from_mano = False
+    optimize_rules = False    
+    optimize_rules = True
+
+    ckpt_fn = "ckpts/grab/102/retargeted_shadow.pth"
+    load_optimized_init_transformations = "ckpts/grab/102/retargeted_shadow.pth"
+    optimize_rules = False
+
+    optimize_rules = True
+    
+    ## opt roboto ##
+    opt_robo_glb_trans = True
+    opt_robo_glb_rot = False # opt rot # ## opt rot ##
+    opt_robo_states = True
+    
+
+    use_multi_stages = False
+
+    minn_dist_threshold_robot_to_obj = 0.1
+
+    penetration_proj_k_to_robot = 40000
+    penetration_proj_k_to_robot_friction = 100000
+
+    drive_robot = "actions"
+    opt_robo_glb_trans = False
+    opt_robo_states = True
+    opt_robo_glb_rot = False 
+
+    train_with_forces_to_active = True
+
+
+
+    load_redmax_robot_actions_fn  = ""
+
+
+    optimize_rules = False
+    
+
+
+    loss_scale_coef = 1.0
+
+
+
+    use_opt_rigid_translations=True
+
+    train_def = True
+
+
+    optimizable_rigid_translations=True
+    
+    nerf {
+        D = 8,
+        d_in = 4,
+        d_in_view = 3,
+        W = 256,
+        multires = 10,
+        multires_view = 4,
+        output_ch = 4,
+        skips=[4],
+        use_viewdirs=True
+    }
+
+    sdf_network {
+        d_out = 257,
+        d_in = 3,
+        d_hidden = 256,
+        n_layers = 8,
+        skip_in = [4],
+        multires = 6,
+        bias = 0.5,
+        scale = 1.0,
+        geometric_init = True,
+        weight_norm = True,
+    } 
+
+    variance_network {
+        init_val = 0.3
+    }
+
+    rendering_network {
+        d_feature = 256,
+        mode = idr,
+        d_in = 9,
+        d_out = 3,
+        d_hidden = 256,
+        n_layers = 4,
+        weight_norm = True,
+        multires_view = 4,
+        squeeze_out = True,
+    }
+
+    neus_renderer {
+        n_samples = 64,
+        n_importance = 64,
+        n_outside = 0,
+        up_sample_steps = 4 ,
+        perturb = 1.0,
+    }
+
+    bending_network {
+        multires = 6,
+        bending_latent_size = 32,
+        d_in = 3,
+        rigidity_hidden_dimensions = 64,
+        rigidity_network_depth = 5,
+        use_rigidity_network = False,
+        bending_n_timesteps = 10,
+    }
+}

confs_new/dyn_grab_arti_shadow_dm_curriculum.conf

@@ -0,0 +1,326 @@
+general {
+    # base_exp_dir = exp/CASE_NAME/wmask
+    base_exp_dir = /data2/datasets/xueyi/neus/exp/CASE_NAME/wmask
+
+    tag = "train_retargeted_shadow_hand_seq_102_diffhand_model_curriculum_"
+
+    recording = [
+        ./,
+        ./models
+    ]
+}
+
+dataset {
+    data_dir = public_data/CASE_NAME/
+    render_cameras_name = cameras_sphere.npz
+    object_cameras_name = cameras_sphere.npz
+
+    obj_idx = 102
+}
+
+train {
+    learning_rate = 5e-4
+    learning_rate_actions = 5e-6
+    # learning_rate = 5e-6
+    # learning_rate = 5e-5
+    learning_rate_alpha = 0.05
+    end_iter = 300000
+
+    # batch_size = 128 #  64
+    # batch_size = 4000
+    # batch_size = 3096 #  64
+    batch_size = 1024 
+    validate_resolution_level = 4
+    warm_up_end = 5000
+    anneal_end = 0
+    use_white_bkgd = False
+
+    # save_freq = 10000
+    save_freq = 10000
+    val_freq = 20 #  2500
+    val_mesh_freq = 20 #  5000
+    report_freq = 10
+    ### igr weight ###
+    igr_weight = 0.1
+    mask_weight = 0.1
+}
+ 
+model {
+    
+    
+    penetration_proj_k_to_robot = 40
+
+    penetrating_depth_penalty = 1.0
+    penetrating_depth_penalty = 0.0
+    train_states = True
+    penetration_proj_k_to_robot = 4000000000.0
+
+
+    minn_dist_threshold = 0.000
+    # minn_dist_threshold = 0.01
+    obj_mass = 100.0
+    obj_mass = 30.0
+
+    optimize_rules = True
+
+    use_mano_hand_for_test = False
+    use_mano_hand_for_test = True
+
+    train_residual_friction = False
+    train_residual_friction = True
+
+    use_LBFGS = True
+    use_LBFGS = False
+
+    use_mano_hand_for_test = False
+    train_residual_friction = True 
+    
+    extract_delta_mesh = False
+    freeze_weights = True
+    # gt_act_xs_def = True
+    gt_act_xs_def = False
+    use_bending_network =  True
+    ### for ts = 3 ###
+    # use_delta_bending = False
+    ### for ts = 3 ###
+    use_delta_bending = True
+    use_passive_nets = True
+    # use_passive_nets = False # sv mesh root #
+
+    use_split_network = True
+
+    penetration_determining = "plane_primitives"
+
+
+    n_timesteps = 3 #  
+    # n_timesteps = 5 # 
+    n_timesteps = 7
+    n_timesteps = 60
+
+    
+
+
+    using_delta_glb_trans = True
+    using_delta_glb_trans = False
+
+    optimize_with_intermediates = False
+    optimize_with_intermediates = True
+
+
+    loss_tangential_diff_coef = 1000
+    loss_tangential_diff_coef = 0
+
+    
+
+    optimize_active_object = False
+    optimize_active_object = True
+
+    # optimize_expanded_pts = False
+    # optimize_expanded_pts = True
+
+    no_friction_constraint = False
+
+    optimize_glb_transformations = True
+    sim_model_path = "DiffHand/assets/hand_sphere_only_hand_testt.xml"
+    mano_sim_model_path = "rsc/mano/mano_mean_wcollision_scaled_scaled_0_9507_nroot.urdf"
+    mano_mult_const_after_cent = 1.0
+    sim_num_steps = 1000000
+
+    bending_net_type = "active_force_field_v18"    
+
+
+    ### try to train the residual friction ? ###
+    train_residual_friction = True
+    optimize_rules = True
+    ### cube ###
+    load_optimized_init_actions = ""
+
+    optimize_rules = False
+
+
+    ## optimize rules ## penetration proj k to robot ##
+    optimize_rules = True
+    penetration_proj_k_to_robot = 4000000.0
+    use_optimizable_params = True
+
+    penetration_determining = "ball_primitives" # uing ball primitives 
+    optimize_rules = True # 
+    penetration_proj_k_to_robot = 4000000.0 # 
+    use_optimizable_params = True
+    train_with_forces_to_active = False
+
+    # penetration_determining = "ball_primitives"
+    ### obj sdf and normals for colllision eteftion and responses ##
+    ### grab train seq 54; cylinder ###
+    penetration_determining = "sdf_of_canon"
+    optimize_rules = True
+    train_with_forces_to_active = False
+
+    ### grab train seq 1 ###
+    penetration_determining = "sdf_of_canon"
+    train_with_forces_to_active = False
+
+    ### grab train seq 224 ###
+    penetration_determining = "sdf_of_canon"
+    train_with_forces_to_active = False
+    loss_scale_coef = 1000.0
+    penetration_proj_k_to_robot_friction = 40000000.0
+    penetration_proj_k_to_robot_friction = 100000000.0
+    use_same_contact_spring_k = False
+    sim_model_path = "DiffHand/assets/hand_sphere_only_hand_testt.xml"
+    sim_model_path = "rsc/shadow_hand_description/shadowhand_new.urdf"
+
+
+    penetration_determining = "sdf_of_canon"
+    optimize_rules = True
+    # optimize_rules = True
+
+    optimize_rules = False
+
+    optimize_rules = True
+
+
+    optimize_rules = False
+
+    optim_sim_model_params_from_mano = True
+    optimize_rules = True
+    optim_sim_model_params_from_mano = False
+    optimize_rules = False    
+
+    penetration_proj_k_to_robot_friction = 100000000.0
+    penetration_proj_k_to_robot = 40000000.0
+
+
+    penetrating_depth_penalty = 1
+
+    minn_dist_threshold_robot_to_obj = 0.0
+
+
+    minn_dist_threshold_robot_to_obj = 0.1
+
+    optim_sim_model_params_from_mano = True
+    optimize_rules = True
+    optim_sim_model_params_from_mano = False
+    optimize_rules = False    
+    optim_sim_model_params_from_mano = False
+    optimize_rules = False    
+
+    load_optimized_init_transformations = ""
+    optim_sim_model_params_from_mano = True
+    optimize_rules = True    
+    minn_dist_threshold_robot_to_obj = 0.0
+
+
+    optim_sim_model_params_from_mano = False
+
+    minn_dist_threshold_robot_to_obj = 0.1
+
+
+    ### kinematics confgs ###
+    obj_sdf_fn = "data/grab/102/102_obj.npy" 
+    kinematic_mano_gt_sv_fn = "data/grab/102/102_sv_dict.npy"
+    scaled_obj_mesh_fn = "data/grab/102/102_obj.obj"
+    # ckpt_fn = ""
+    load_optimized_init_transformations = ""
+    optim_sim_model_params_from_mano = True
+    optimize_rules = True    
+    minn_dist_threshold_robot_to_obj = 0.0
+
+    optim_sim_model_params_from_mano = False
+
+    optimize_rules = True
+
+    ckpt_fn = "ckpts/grab/102/retargeted_shadow.pth"
+    ckpt_fn = "/data2/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_102_optactswreacts_redmaxacts_rules_/checkpoints/ckpt_035459.pth"
+    load_optimized_init_transformations = "ckpts/grab/102/retargeted_shadow.pth"
+
+
+    optimize_rules = True
+    
+    ## opt roboto ##
+    opt_robo_glb_trans = True
+    opt_robo_glb_rot = False # opt rot # ## opt rot ##
+    opt_robo_states = True
+
+
+    load_redmax_robot_actions_fn  = "ckpts/grab/102/diffhand_act.npy"
+
+
+
+    ckpt_fn = ""
+
+    use_multi_stages = True
+    train_with_forces_to_active = True
+
+
+    # optimize_rules = False
+    loss_scale_coef = 1.0 ## loss scale coef ## loss scale coef #### 
+
+
+
+    use_opt_rigid_translations=True
+
+    train_def = True
+
+    # optimizable_rigid_translations = False #
+    optimizable_rigid_translations=True
+    
+    nerf {
+        D = 8,
+        d_in = 4,
+        d_in_view = 3,
+        W = 256,
+        multires = 10,
+        multires_view = 4,
+        output_ch = 4,
+        skips=[4],
+        use_viewdirs=True
+    }
+
+    sdf_network {
+        d_out = 257,
+        d_in = 3,
+        d_hidden = 256,
+        n_layers = 8,
+        skip_in = [4],
+        multires = 6,
+        bias = 0.5,
+        scale = 1.0,
+        geometric_init = True,
+        weight_norm = True,
+    } 
+
+    variance_network {
+        init_val = 0.3
+    }
+
+    rendering_network {
+        d_feature = 256,
+        mode = idr,
+        d_in = 9,
+        d_out = 3,
+        d_hidden = 256,
+        n_layers = 4,
+        weight_norm = True,
+        multires_view = 4,
+        squeeze_out = True,
+    }
+
+    neus_renderer {
+        n_samples = 64,
+        n_importance = 64,
+        n_outside = 0,
+        up_sample_steps = 4 ,
+        perturb = 1.0,
+    }
+
+    bending_network {
+        multires = 6,
+        bending_latent_size = 32,
+        d_in = 3,
+        rigidity_hidden_dimensions = 64,
+        rigidity_network_depth = 5,
+        use_rigidity_network = False,
+        bending_n_timesteps = 10,
+    }
+}

confs_new/dyn_grab_arti_shadow_dm_singlestage.conf

@@ -0,0 +1,318 @@
+general {
+    base_exp_dir = exp/CASE_NAME/wmask
+
+    tag = "train_retargeted_shadow_hand_seq_102_diffhand_model_curriculum_"
+
+    recording = [
+        ./,
+        ./models
+    ]
+}
+
+dataset {
+    data_dir = public_data/CASE_NAME/
+    render_cameras_name = cameras_sphere.npz
+    object_cameras_name = cameras_sphere.npz
+
+    obj_idx = 102
+}
+
+train {
+    learning_rate = 5e-4
+    learning_rate_actions = 5e-6
+    learning_rate_alpha = 0.05
+    end_iter = 300000
+
+    batch_size = 1024 
+    validate_resolution_level = 4
+    warm_up_end = 5000
+    anneal_end = 0
+    use_white_bkgd = False
+
+
+    save_freq = 10000
+    val_freq = 20
+    val_mesh_freq = 20
+    report_freq = 10
+    igr_weight = 0.1
+    mask_weight = 0.1
+}
+ 
+model {
+    
+    
+    penetration_proj_k_to_robot = 40
+
+    penetrating_depth_penalty = 1.0
+    penetrating_depth_penalty = 0.0
+    train_states = True
+    penetration_proj_k_to_robot = 4000000000.0
+
+
+    minn_dist_threshold = 0.000
+    # minn_dist_threshold = 0.01
+    obj_mass = 100.0
+    obj_mass = 30.0
+
+    optimize_rules = True
+
+    use_mano_hand_for_test = False
+    use_mano_hand_for_test = True
+
+    train_residual_friction = False
+    train_residual_friction = True
+
+    use_LBFGS = True
+    use_LBFGS = False
+
+    use_mano_hand_for_test = False
+    train_residual_friction = True 
+    
+    extract_delta_mesh = False
+    freeze_weights = True
+    # gt_act_xs_def = True
+    gt_act_xs_def = False
+    use_bending_network =  True
+    ### for ts = 3 ###
+    # use_delta_bending = False
+    ### for ts = 3 ###
+    use_delta_bending = True
+    use_passive_nets = True
+    # use_passive_nets = False # sv mesh root #
+
+    use_split_network = True
+
+    penetration_determining = "plane_primitives"
+
+
+    n_timesteps = 3 #  
+    # n_timesteps = 5 # 
+    n_timesteps = 7
+    n_timesteps = 60
+
+    
+
+
+    using_delta_glb_trans = True
+    using_delta_glb_trans = False
+
+    optimize_with_intermediates = False
+    optimize_with_intermediates = True
+
+
+    loss_tangential_diff_coef = 1000
+    loss_tangential_diff_coef = 0
+
+    
+
+    optimize_active_object = False
+    optimize_active_object = True
+
+    # optimize_expanded_pts = False
+    # optimize_expanded_pts = True
+
+    no_friction_constraint = False
+
+    optimize_glb_transformations = True
+    sim_model_path = "DiffHand/assets/hand_sphere_only_hand_testt.xml"
+    mano_sim_model_path = "rsc/mano/mano_mean_wcollision_scaled_scaled_0_9507_nroot.urdf"
+    mano_mult_const_after_cent = 1.0
+    sim_num_steps = 1000000
+
+    bending_net_type = "active_force_field_v18"    
+
+
+    ### try to train the residual friction ? ###
+    train_residual_friction = True
+    optimize_rules = True
+    ### cube ###
+    load_optimized_init_actions = ""
+
+    optimize_rules = False
+
+
+    ## optimize rules ## penetration proj k to robot ##
+    optimize_rules = True
+    penetration_proj_k_to_robot = 4000000.0
+    use_optimizable_params = True
+
+    penetration_determining = "ball_primitives" # uing ball primitives 
+    optimize_rules = True # 
+    penetration_proj_k_to_robot = 4000000.0 # 
+    use_optimizable_params = True
+    train_with_forces_to_active = False
+
+    # penetration_determining = "ball_primitives"
+    ### obj sdf and normals for colllision eteftion and responses ##
+    ### grab train seq 54; cylinder ###
+    penetration_determining = "sdf_of_canon"
+    optimize_rules = True
+    train_with_forces_to_active = False
+
+    ### grab train seq 1 ###
+    penetration_determining = "sdf_of_canon"
+    train_with_forces_to_active = False
+
+    ### grab train seq 224 ###
+    penetration_determining = "sdf_of_canon"
+    train_with_forces_to_active = False
+    loss_scale_coef = 1000.0
+    penetration_proj_k_to_robot_friction = 40000000.0
+    penetration_proj_k_to_robot_friction = 100000000.0
+    use_same_contact_spring_k = False
+    sim_model_path = "DiffHand/assets/hand_sphere_only_hand_testt.xml"
+    sim_model_path = "rsc/shadow_hand_description/shadowhand_new.urdf"
+
+
+    penetration_determining = "sdf_of_canon"
+    optimize_rules = True
+    # optimize_rules = True
+
+    optimize_rules = False
+
+    optimize_rules = True
+
+
+    optimize_rules = False
+
+    optim_sim_model_params_from_mano = True
+    optimize_rules = True
+    optim_sim_model_params_from_mano = False
+    optimize_rules = False    
+
+    penetration_proj_k_to_robot_friction = 100000000.0
+    penetration_proj_k_to_robot = 40000000.0
+
+
+    penetrating_depth_penalty = 1
+
+    minn_dist_threshold_robot_to_obj = 0.0
+
+
+    minn_dist_threshold_robot_to_obj = 0.1
+
+    optim_sim_model_params_from_mano = True
+    optimize_rules = True
+    optim_sim_model_params_from_mano = False
+    optimize_rules = False    
+    optim_sim_model_params_from_mano = False
+    optimize_rules = False    
+
+    load_optimized_init_transformations = ""
+    optim_sim_model_params_from_mano = True
+    optimize_rules = True    
+    minn_dist_threshold_robot_to_obj = 0.0
+
+
+    optim_sim_model_params_from_mano = False
+
+    minn_dist_threshold_robot_to_obj = 0.1
+
+
+    ### kinematics confgs ###
+    obj_sdf_fn = "data/grab/102/102_obj.npy" 
+    kinematic_mano_gt_sv_fn = "data/grab/102/102_sv_dict.npy"
+    scaled_obj_mesh_fn = "data/grab/102/102_obj.obj"
+    # ckpt_fn = ""
+    load_optimized_init_transformations = ""
+    optim_sim_model_params_from_mano = True
+    optimize_rules = True    
+    minn_dist_threshold_robot_to_obj = 0.0
+
+    optim_sim_model_params_from_mano = False
+
+    optimize_rules = True
+
+    ckpt_fn = "ckpts/grab/102/retargeted_shadow.pth"
+    load_optimized_init_transformations = "ckpts/grab/102/retargeted_shadow.pth"
+
+
+    optimize_rules = True
+    
+    ## opt roboto ##
+    opt_robo_glb_trans = True
+    opt_robo_glb_rot = False
+    opt_robo_states = True
+
+
+    load_redmax_robot_actions_fn  = "ckpts/grab/102/diffhand_act.npy"
+
+
+
+    ckpt_fn = ""
+
+    use_multi_stages = False
+    train_with_forces_to_active = True
+
+
+    # optimize_rules = False
+    loss_scale_coef = 1.0 ## loss scale coef ## loss scale coef #### 
+
+
+
+    use_opt_rigid_translations=True
+
+    train_def = True
+
+    # optimizable_rigid_translations = False #
+    optimizable_rigid_translations=True
+    
+    nerf {
+        D = 8,
+        d_in = 4,
+        d_in_view = 3,
+        W = 256,
+        multires = 10,
+        multires_view = 4,
+        output_ch = 4,
+        skips=[4],
+        use_viewdirs=True
+    }
+
+    sdf_network {
+        d_out = 257,
+        d_in = 3,
+        d_hidden = 256,
+        n_layers = 8,
+        skip_in = [4],
+        multires = 6,
+        bias = 0.5,
+        scale = 1.0,
+        geometric_init = True,
+        weight_norm = True,
+    } 
+
+    variance_network {
+        init_val = 0.3
+    }
+
+    rendering_network {
+        d_feature = 256,
+        mode = idr,
+        d_in = 9,
+        d_out = 3,
+        d_hidden = 256,
+        n_layers = 4,
+        weight_norm = True,
+        multires_view = 4,
+        squeeze_out = True,
+    }
+
+    neus_renderer {
+        n_samples = 64,
+        n_importance = 64,
+        n_outside = 0,
+        up_sample_steps = 4 ,
+        perturb = 1.0,
+    }
+
+    bending_network {
+        multires = 6,
+        bending_latent_size = 32,
+        d_in = 3,
+        rigidity_hidden_dimensions = 64,
+        rigidity_network_depth = 5,
+        use_rigidity_network = False,
+        bending_n_timesteps = 10,
+    }
+}

confs_new/dyn_grab_pointset_mano.conf

@@ -0,0 +1,215 @@
+general {
+
+
+    base_exp_dir = exp/CASE_NAME/wmask
+
+
+    # tag = "train_retargeted_shadow_hand_seq_102_mano_sparse_retargeting_"
+    tag = "train_dyn_mano_acts_"
+
+    recording = [
+        ./,
+        ./models
+    ]
+}
+
+dataset {
+    data_dir = public_data/CASE_NAME/
+    render_cameras_name = cameras_sphere.npz
+    object_cameras_name = cameras_sphere.npz
+    obj_idx = 102
+}
+
+train {
+    learning_rate = 5e-4
+    learning_rate_alpha = 0.05
+    end_iter = 300000
+
+    batch_size = 1024
+    validate_resolution_level = 4
+    warm_up_end = 5000
+    anneal_end = 0
+    use_white_bkgd = False
+
+    # save_freq = 10000
+    save_freq = 10000
+    val_freq = 20
+    val_mesh_freq = 20
+    report_freq = 10
+    igr_weight = 0.1
+    mask_weight = 0.1
+}
+ 
+model {
+
+    optimize_dyn_actions = True
+
+
+    optimize_robot = True
+
+    use_penalty_based_friction = True
+
+    use_split_params = False
+
+    use_sqr_spring_stiffness = True
+
+    use_pre_proj_frictions = True
+
+
+
+    use_sqrt_dist = True
+    contact_maintaining_dist_thres = 0.2
+
+    robot_actions_diff_coef = 0.001
+
+
+    use_sdf_as_contact_dist = True
+
+
+    # 
+    use_contact_dist_as_sdf = False
+
+    use_glb_proj_delta = True
+
+
+
+# penetration_proj_k_to_robot = 30
+    penetrating_depth_penalty = 1.0
+    train_states = True
+
+
+
+    minn_dist_threshold = 0.000
+    obj_mass = 30.0
+
+
+    use_LBFGS = True
+    use_LBFGS = False
+
+    use_mano_hand_for_test = False # use the dynamic mano model here #
+
+    extract_delta_mesh = False
+    freeze_weights = True
+    gt_act_xs_def = False
+    use_bending_network =  True
+    ### for ts = 3 ###
+    # use_delta_bending = False
+    ### for ts = 3 ###
+   
+
+
+
+    sim_model_path = "rsc/shadow_hand_description/shadowhand_new.urdf"
+    mano_sim_model_path = "rsc/mano/mano_mean_wcollision_scaled_scaled_0_9507_nroot.urdf"
+
+    obj_sdf_fn = "data/grab/102/102_obj.npy" 
+    kinematic_mano_gt_sv_fn = "data/grab/102/102_sv_dict.npy"
+    scaled_obj_mesh_fn = "data/grab/102/102_obj.obj"
+
+    bending_net_type = "active_force_field_v18"    
+    sim_num_steps = 1000000
+    n_timesteps = 60
+    optim_sim_model_params_from_mano = False
+    penetration_determining = "sdf_of_canon"
+    train_with_forces_to_active = False
+    loss_scale_coef = 1000.0
+    use_same_contact_spring_k = False
+    use_optimizable_params = True # 
+    train_residual_friction = True
+    mano_mult_const_after_cent = 1.0
+    optimize_glb_transformations = True
+    no_friction_constraint = False
+    optimize_active_object = True
+    loss_tangential_diff_coef = 0
+    optimize_with_intermediates = True
+    using_delta_glb_trans = False
+    train_multi_seqs = False
+    use_split_network = True
+    use_delta_bending = True
+
+
+
+
+    ##### contact spring model settings ####
+    minn_dist_threshold_robot_to_obj = 0.1
+    penetration_proj_k_to_robot_friction = 10000000.0
+    penetration_proj_k_to_robot = 4000000.0
+    ##### contact spring model settings ####
+ 
+
+    ######  ######
+    # drive_pointset = "states"
+    fix_obj = True # to track the hand only 
+    optimize_rules = False
+    train_pointset_acts_via_deltas = False
+    load_optimized_init_actions = ""
+    load_optimized_init_transformations = ""
+    ckpt_fn = ""
+    retar_only_glb = True
+    # use_multi_stages = True
+    ###### Stage 1: threshold, ks settings 1, optimize offsets ######
+
+    use_opt_rigid_translations=True
+
+    train_def = True
+    optimizable_rigid_translations=True
+    
+    nerf {
+        D = 8,
+        d_in = 4,
+        d_in_view = 3,
+        W = 256,
+        multires = 10,
+        multires_view = 4,
+        output_ch = 4,
+        skips=[4],
+        use_viewdirs=True
+    }
+
+    sdf_network {
+        d_out = 257,
+        d_in = 3,
+        d_hidden = 256,
+        n_layers = 8,
+        skip_in = [4],
+        multires = 6,
+        bias = 0.5,
+        scale = 1.0,
+        geometric_init = True,
+        weight_norm = True,
+    } 
+
+    variance_network {
+        init_val = 0.3
+    }
+
+    rendering_network {
+        d_feature = 256,
+        mode = idr,
+        d_in = 9,
+        d_out = 3,
+        d_hidden = 256,
+        n_layers = 4,
+        weight_norm = True,
+        multires_view = 4,
+        squeeze_out = True,
+    }
+
+    neus_renderer {
+        n_samples = 64,
+        n_importance = 64,
+        n_outside = 0,
+        up_sample_steps = 4 ,
+        perturb = 1.0,
+    }
+
+    bending_network {
+        multires = 6,
+        bending_latent_size = 32,
+        d_in = 3,
+        rigidity_hidden_dimensions = 64,
+        rigidity_network_depth = 5,
+        use_rigidity_network = False,
+        bending_n_timesteps = 10,
+    }
+}

confs_new/dyn_grab_pointset_mano_dyn.conf

@@ -0,0 +1,218 @@
+general {
+
+
+    base_exp_dir = exp/CASE_NAME/wmask
+
+
+    # tag = "train_retargeted_shadow_hand_seq_102_mano_sparse_retargeting_"
+    # tag = "train_dyn_mano_acts_"
+    tag = "train_dyn_mano_acts_wreact_optps_"
+
+    recording = [
+        ./,
+        ./models
+    ]
+}
+
+dataset {
+    data_dir = public_data/CASE_NAME/
+    render_cameras_name = cameras_sphere.npz
+    object_cameras_name = cameras_sphere.npz
+    obj_idx = 102
+}
+
+train {
+    learning_rate = 5e-4
+    learning_rate_alpha = 0.05
+    end_iter = 300000
+
+    batch_size = 1024
+    validate_resolution_level = 4
+    warm_up_end = 5000
+    anneal_end = 0
+    use_white_bkgd = False
+
+    # save_freq = 10000
+    save_freq = 10000
+    val_freq = 20
+    val_mesh_freq = 20
+    report_freq = 10
+    igr_weight = 0.1
+    mask_weight = 0.1
+}
+ 
+model {
+
+    optimize_dyn_actions = True
+
+
+    optimize_robot = True
+
+    use_penalty_based_friction = True
+
+    use_split_params = False
+
+    use_sqr_spring_stiffness = True
+
+    use_pre_proj_frictions = True
+
+
+
+    use_sqrt_dist = True
+    contact_maintaining_dist_thres = 0.2
+
+    robot_actions_diff_coef = 0.001
+
+
+    use_sdf_as_contact_dist = True
+
+
+    # 
+    use_contact_dist_as_sdf = False
+
+    use_glb_proj_delta = True
+
+
+
+# penetration_proj_k_to_robot = 30
+    penetrating_depth_penalty = 1.0
+    train_states = True
+
+
+
+    minn_dist_threshold = 0.000
+    obj_mass = 30.0
+
+
+    use_LBFGS = True
+    use_LBFGS = False
+
+    use_mano_hand_for_test = False # use the dynamic mano model here #
+
+    extract_delta_mesh = False
+    freeze_weights = True
+    gt_act_xs_def = False
+    use_bending_network =  True
+    ### for ts = 3 ###
+    # use_delta_bending = False
+    ### for ts = 3 ###
+   
+
+
+
+    sim_model_path = "rsc/shadow_hand_description/shadowhand_new.urdf"
+    mano_sim_model_path = "rsc/mano/mano_mean_wcollision_scaled_scaled_0_9507_nroot.urdf"
+
+    obj_sdf_fn = "data/grab/102/102_obj.npy" 
+    kinematic_mano_gt_sv_fn = "data/grab/102/102_sv_dict.npy"
+    scaled_obj_mesh_fn = "data/grab/102/102_obj.obj"
+
+    bending_net_type = "active_force_field_v18"    
+    sim_num_steps = 1000000
+    n_timesteps = 60
+    optim_sim_model_params_from_mano = False
+    penetration_determining = "sdf_of_canon"
+    train_with_forces_to_active = False
+    loss_scale_coef = 1000.0
+    use_same_contact_spring_k = False
+    use_optimizable_params = True # 
+    train_residual_friction = True
+    mano_mult_const_after_cent = 1.0
+    optimize_glb_transformations = True
+    no_friction_constraint = False
+    optimize_active_object = True
+    loss_tangential_diff_coef = 0
+    optimize_with_intermediates = True
+    using_delta_glb_trans = False
+    train_multi_seqs = False
+    use_split_network = True
+    use_delta_bending = True
+
+
+
+
+    ##### contact spring model settings ####
+    minn_dist_threshold_robot_to_obj = 0.1
+    penetration_proj_k_to_robot_friction = 10000000.0
+    penetration_proj_k_to_robot = 4000000.0
+    ##### contact spring model settings ####
+ 
+
+    ###### Stage 1: optimize for the parametes ######
+    # drive_pointset = "states"
+    fix_obj = False
+    optimize_rules = True
+    train_pointset_acts_via_deltas = False
+    load_optimized_init_actions = "ckpts/grab/102/dyn_mano_arti.pth"
+    load_optimized_init_transformations = ""
+    ckpt_fn = "ckpts/grab/102/dyn_mano_arti.pth"
+    # retar_only_glb = True
+    # use_multi_stages = True
+    ###### Stage 1: optimize for the parametes ######
+
+    use_opt_rigid_translations=True
+
+    train_def = True
+    optimizable_rigid_translations=True
+    
+    nerf {
+        D = 8,
+        d_in = 4,
+        d_in_view = 3,
+        W = 256,
+        multires = 10,
+        multires_view = 4,
+        output_ch = 4,
+        skips=[4],
+        use_viewdirs=True
+    }
+
+    sdf_network {
+        d_out = 257,
+        d_in = 3,
+        d_hidden = 256,
+        n_layers = 8,
+        skip_in = [4],
+        multires = 6,
+        bias = 0.5,
+        scale = 1.0,
+        geometric_init = True,
+        weight_norm = True,
+    } 
+
+    variance_network {
+        init_val = 0.3
+    }
+
+    rendering_network {
+        d_feature = 256,
+        mode = idr,
+        d_in = 9,
+        d_out = 3,
+        d_hidden = 256,
+        n_layers = 4,
+        weight_norm = True,
+        multires_view = 4,
+        squeeze_out = True,
+    }
+
+    neus_renderer {
+        n_samples = 64,
+        n_importance = 64,
+        n_outside = 0,
+        up_sample_steps = 4 ,
+        perturb = 1.0,
+    }
+
+    bending_network {
+        multires = 6,
+        bending_latent_size = 32,
+        d_in = 3,
+        rigidity_hidden_dimensions = 64,
+        rigidity_network_depth = 5,
+        use_rigidity_network = False,
+        bending_n_timesteps = 10,
+    }
+}
+
+

confs_new/dyn_grab_pointset_mano_dyn_optacts.conf

@@ -0,0 +1,218 @@
+general {
+
+
+    base_exp_dir = exp/CASE_NAME/wmask
+
+
+    # tag = "train_retargeted_shadow_hand_seq_102_mano_sparse_retargeting_"
+    # tag = "train_dyn_mano_acts_"
+    tag = "train_dyn_mano_acts_wreact_optps_optacts_"
+
+    recording = [
+        ./,
+        ./models
+    ]
+}
+
+dataset {
+    data_dir = public_data/CASE_NAME/
+    render_cameras_name = cameras_sphere.npz
+    object_cameras_name = cameras_sphere.npz
+    obj_idx = 102
+}
+
+train {
+    learning_rate = 5e-4
+    learning_rate_alpha = 0.05
+    end_iter = 300000
+
+    batch_size = 1024
+    validate_resolution_level = 4
+    warm_up_end = 5000
+    anneal_end = 0
+    use_white_bkgd = False
+
+    # save_freq = 10000
+    save_freq = 10000
+    val_freq = 20
+    val_mesh_freq = 20
+    report_freq = 10
+    igr_weight = 0.1
+    mask_weight = 0.1
+}
+ 
+model {
+
+    optimize_dyn_actions = True
+
+
+    optimize_robot = True
+
+    use_penalty_based_friction = True
+
+    use_split_params = False
+
+    use_sqr_spring_stiffness = True
+
+    use_pre_proj_frictions = True
+
+
+
+    use_sqrt_dist = True
+    contact_maintaining_dist_thres = 0.2
+
+    robot_actions_diff_coef = 0.001
+
+
+    use_sdf_as_contact_dist = True
+
+
+    # 
+    use_contact_dist_as_sdf = False
+
+    use_glb_proj_delta = True
+
+
+
+# penetration_proj_k_to_robot = 30
+    penetrating_depth_penalty = 1.0
+    train_states = True
+
+
+
+    minn_dist_threshold = 0.000
+    obj_mass = 30.0
+
+
+    use_LBFGS = True
+    use_LBFGS = False
+
+    use_mano_hand_for_test = False # use the dynamic mano model here #
+
+    extract_delta_mesh = False
+    freeze_weights = True
+    gt_act_xs_def = False
+    use_bending_network =  True
+    ### for ts = 3 ###
+    # use_delta_bending = False
+    ### for ts = 3 ###
+   
+
+
+
+    sim_model_path = "rsc/shadow_hand_description/shadowhand_new.urdf"
+    mano_sim_model_path = "rsc/mano/mano_mean_wcollision_scaled_scaled_0_9507_nroot.urdf"
+
+    obj_sdf_fn = "data/grab/102/102_obj.npy" 
+    kinematic_mano_gt_sv_fn = "data/grab/102/102_sv_dict.npy"
+    scaled_obj_mesh_fn = "data/grab/102/102_obj.obj"
+
+    bending_net_type = "active_force_field_v18"    
+    sim_num_steps = 1000000
+    n_timesteps = 60
+    optim_sim_model_params_from_mano = False
+    penetration_determining = "sdf_of_canon"
+    train_with_forces_to_active = False
+    loss_scale_coef = 1000.0
+    use_same_contact_spring_k = False
+    use_optimizable_params = True # 
+    train_residual_friction = True
+    mano_mult_const_after_cent = 1.0
+    optimize_glb_transformations = True
+    no_friction_constraint = False
+    optimize_active_object = True
+    loss_tangential_diff_coef = 0
+    optimize_with_intermediates = True
+    using_delta_glb_trans = False
+    train_multi_seqs = False
+    use_split_network = True
+    use_delta_bending = True
+
+
+
+
+    ##### contact spring model settings ####
+    minn_dist_threshold_robot_to_obj = 0.1
+    penetration_proj_k_to_robot_friction = 10000000.0
+    penetration_proj_k_to_robot = 4000000.0
+    ##### contact spring model settings ####
+ 
+
+    ###### Stage 1: optimize for the parametes ######
+    # drive_pointset = "states"
+    fix_obj = False
+    optimize_rules = False
+    train_pointset_acts_via_deltas = False
+    load_optimized_init_actions = "ckpts/grab/102/dyn_mano_arti.pth"
+    load_optimized_init_transformations = ""
+    ckpt_fn = "ckpts/grab/102/dyn_mano_arti.pth"
+    # retar_only_glb = True
+    # use_multi_stages = True
+    ###### Stage 1: optimize for the parametes ######
+
+    use_opt_rigid_translations=True
+
+    train_def = True
+    optimizable_rigid_translations=True
+    
+    nerf {
+        D = 8,
+        d_in = 4,
+        d_in_view = 3,
+        W = 256,
+        multires = 10,
+        multires_view = 4,
+        output_ch = 4,
+        skips=[4],
+        use_viewdirs=True
+    }
+
+    sdf_network {
+        d_out = 257,
+        d_in = 3,
+        d_hidden = 256,
+        n_layers = 8,
+        skip_in = [4],
+        multires = 6,
+        bias = 0.5,
+        scale = 1.0,
+        geometric_init = True,
+        weight_norm = True,
+    } 
+
+    variance_network {
+        init_val = 0.3
+    }
+
+    rendering_network {
+        d_feature = 256,
+        mode = idr,
+        d_in = 9,
+        d_out = 3,
+        d_hidden = 256,
+        n_layers = 4,
+        weight_norm = True,
+        multires_view = 4,
+        squeeze_out = True,
+    }
+
+    neus_renderer {
+        n_samples = 64,
+        n_importance = 64,
+        n_outside = 0,
+        up_sample_steps = 4 ,
+        perturb = 1.0,
+    }
+
+    bending_network {
+        multires = 6,
+        bending_latent_size = 32,
+        d_in = 3,
+        rigidity_hidden_dimensions = 64,
+        rigidity_network_depth = 5,
+        use_rigidity_network = False,
+        bending_n_timesteps = 10,
+    }
+}
+
+

confs_new/dyn_grab_pointset_points_dyn.conf

@@ -0,0 +1,257 @@
+general {
+
+
+    base_exp_dir = exp/CASE_NAME/wmask
+    
+    tag = "train_retargeted_shadow_hand_seq_102_mano_pointset_acts_"
+
+    recording = [
+        ./,
+        ./models
+    ]
+}
+
+dataset {
+    data_dir = public_data/CASE_NAME/
+    render_cameras_name = cameras_sphere.npz
+    object_cameras_name = cameras_sphere.npz
+    obj_idx = 102
+}
+
+train {
+    learning_rate = 5e-4
+    learning_rate_alpha = 0.05
+    end_iter = 300000
+
+    batch_size = 1024 #  64
+    validate_resolution_level = 4
+    warm_up_end = 5000
+    anneal_end = 0
+    use_white_bkgd = False
+
+    # save_freq = 10000
+    save_freq = 10000
+    val_freq = 20 #  2500
+    val_mesh_freq = 20 #  5000
+    report_freq = 10
+    ### igr weight ###
+    igr_weight = 0.1
+    mask_weight = 0.1
+}
+ 
+model {
+
+    optimize_dyn_actions = True
+
+
+    optimize_robot = True
+
+    use_penalty_based_friction = True
+
+    use_split_params = False
+
+    use_sqr_spring_stiffness = True
+
+    use_pre_proj_frictions = True
+
+
+
+    use_sqrt_dist = True
+    contact_maintaining_dist_thres = 0.2
+
+    robot_actions_diff_coef = 0.001
+
+
+    use_sdf_as_contact_dist = True
+
+
+    # 
+    use_contact_dist_as_sdf = False
+
+    use_glb_proj_delta = True
+
+
+
+# penetration_proj_k_to_robot = 30
+    penetrating_depth_penalty = 1.0
+    train_states = True
+
+
+
+    minn_dist_threshold = 0.000
+    obj_mass = 30.0
+
+
+    use_LBFGS = True
+    use_LBFGS = False
+
+    use_mano_hand_for_test = False # use the dynamic mano model here #
+
+    extract_delta_mesh = False
+    freeze_weights = True
+    gt_act_xs_def = False
+    use_bending_network =  True
+    ### for ts = 3 ###
+    # use_delta_bending = False
+    ### for ts = 3 ###
+   
+
+
+
+    sim_model_path = "rsc/shadow_hand_description/shadowhand_new.urdf"
+    mano_sim_model_path = "rsc/mano/mano_mean_wcollision_scaled_scaled_0_9507_nroot.urdf"
+
+    obj_sdf_fn = "data/grab/102/102_obj.npy" 
+    kinematic_mano_gt_sv_fn = "data/grab/102/102_sv_dict.npy"
+    scaled_obj_mesh_fn = "data/grab/102/102_obj.obj"
+
+    bending_net_type = "active_force_field_v18"    
+    sim_num_steps = 1000000
+    n_timesteps = 60
+    optim_sim_model_params_from_mano = False
+    penetration_determining = "sdf_of_canon"
+    train_with_forces_to_active = False
+    loss_scale_coef = 1000.0
+    use_same_contact_spring_k = False
+    use_optimizable_params = True # 
+    train_residual_friction = True
+    mano_mult_const_after_cent = 1.0
+    optimize_glb_transformations = True
+    no_friction_constraint = False
+    optimize_active_object = True
+    loss_tangential_diff_coef = 0
+    optimize_with_intermediates = True
+    using_delta_glb_trans = False
+    train_multi_seqs = False
+    use_split_network = True
+    use_delta_bending = True
+
+
+
+
+
+
+
+
+
+    ###### threshold, ks settings 1, optimize acts ######
+    # drive_pointset = "actions"
+    # fix_obj = True
+    # optimize_rules = False
+    # train_pointset_acts_via_deltas = True
+    # load_optimized_init_actions = "/data/xueyi/NeuS/exp/hand_test_routine_2_light_color_wtime_active_passive/wmask_reverse_value_totviews_tag_train_dyn_mano_hand_seq_102_mouse_optdynactions_points_optrobo_offsetdriven_optrules_multk100_wfixobj_optdelta_radius0d4_/checkpoints/ckpt_002000.pth"
+    # load_optimized_init_actions = "/data/xueyi/NeuS/exp/hand_test_routine_2_light_color_wtime_active_passive/wmask_reverse_value_totviews_tag_train_dyn_mano_hand_seq_102_mouse_optdynactions_points_optrobo_offsetdriven_optrules_multk100_wfixobj_optdelta_radius0d2_/checkpoints/ckpt_008000.pth"
+    ###### threshold, ks settings 1, optimize acts ######
+
+
+    ##### contact spring model settings ####
+    minn_dist_threshold_robot_to_obj = 0.1
+    penetration_proj_k_to_robot_friction = 10000000.0
+    penetration_proj_k_to_robot = 4000000.0
+    ##### contact spring model settings ####
+ 
+
+    ###### Stage 1: threshold, ks settings 1, optimize offsets ######
+    drive_pointset = "states"
+    fix_obj = True
+    optimize_rules = False
+    train_pointset_acts_via_deltas = False
+    load_optimized_init_actions = "ckpts/grab/102/dyn_mano_arti.pth"
+    ###### Stage 1: threshold, ks settings 1, optimize offsets ######
+
+
+    ###### Stage 2: threshold, ks settings 1, optimize acts ######
+    drive_pointset = "actions"
+    fix_obj = True
+    optimize_rules = False
+    train_pointset_acts_via_deltas = True
+    load_optimized_init_actions = "ckpts/grab/102/dyn_mano_pointset_states.pt"
+    ###### Stage 2: threshold, ks settings 1, optimize acts ######
+
+
+    ###### Stage 3: threshold, ks settings 1, optimize params from acts ######
+    drive_pointset = "actions"
+    fix_obj = False
+    optimize_rules = True
+    train_pointset_acts_via_deltas = True
+    load_optimized_init_actions = "ckpts/grab/102/dyn_mano_pointset_acts.pt"
+    ##### model parameters optimized from the MANO hand trajectory #####
+    ckpt_fn = "ckpts/grab/102/dyn_mano_opts.pt"
+    ###### Stage 3: threshold, ks settings 1, optimize params from acts ######
+
+
+    ###### Stage 4: threshold, ks settings 1, optimize acts from optimized params ######
+    drive_pointset = "actions"
+    fix_obj = False
+    optimize_rules = False
+    train_pointset_acts_via_deltas = True ## pointset acts via deltas ###
+    ##### model parameters optimized from the MANO hand expanded set trajectory #####
+    ckpt_fn = "ckpts/grab/102/dyn_mano_pointset_optimized_acts_optimized_ps.pth"
+    load_optimized_init_actions = "ckpts/grab/102/dyn_mano_pointset_optimized_acts.pth"
+    ###### Stage 4: threshold, ks settings 1, optimize acts from optimized params ######
+
+
+    use_opt_rigid_translations=True
+
+    train_def = True
+    optimizable_rigid_translations=True
+    
+    nerf {
+        D = 8,
+        d_in = 4,
+        d_in_view = 3,
+        W = 256,
+        multires = 10,
+        multires_view = 4,
+        output_ch = 4,
+        skips=[4],
+        use_viewdirs=True
+    }
+
+    sdf_network {
+        d_out = 257,
+        d_in = 3,
+        d_hidden = 256,
+        n_layers = 8,
+        skip_in = [4],
+        multires = 6,
+        bias = 0.5,
+        scale = 1.0,
+        geometric_init = True,
+        weight_norm = True,
+    } 
+
+    variance_network {
+        init_val = 0.3
+    }
+
+    rendering_network {
+        d_feature = 256,
+        mode = idr,
+        d_in = 9,
+        d_out = 3,
+        d_hidden = 256,
+        n_layers = 4,
+        weight_norm = True,
+        multires_view = 4,
+        squeeze_out = True,
+    }
+
+    neus_renderer {
+        n_samples = 64,
+        n_importance = 64,
+        n_outside = 0,
+        up_sample_steps = 4 ,
+        perturb = 1.0,
+    }
+
+    bending_network {
+        multires = 6,
+        bending_latent_size = 32,
+        d_in = 3,
+        rigidity_hidden_dimensions = 64,
+        rigidity_network_depth = 5,
+        use_rigidity_network = False,
+        bending_n_timesteps = 10,
+    }
+}

confs_new/dyn_grab_pointset_points_dyn_retar.conf

@@ -0,0 +1,274 @@
+general {
+
+
+    base_exp_dir = exp/CASE_NAME/wmask
+    
+    tag = "train_retargeted_shadow_hand_seq_102_mano_pointset_acts_"
+    tag = "train_retargeted_shadow_hand_seq_102_mano_pointset_acts_retar_to_shadow_"
+
+    recording = [
+        ./,
+        ./models
+    ]
+}
+
+dataset {
+    data_dir = public_data/CASE_NAME/
+    render_cameras_name = cameras_sphere.npz
+    object_cameras_name = cameras_sphere.npz
+    obj_idx = 102
+}
+
+train {
+    learning_rate = 5e-4
+    learning_rate_alpha = 0.05
+    end_iter = 300000
+
+    batch_size = 1024 #  64
+    validate_resolution_level = 4
+    warm_up_end = 5000
+    anneal_end = 0
+    use_white_bkgd = False
+
+    # save_freq = 10000
+    save_freq = 10000
+    val_freq = 20 #  2500
+    val_mesh_freq = 20 #  5000
+    report_freq = 10
+    ### igr weight ###
+    igr_weight = 0.1
+    mask_weight = 0.1
+}
+ 
+model {
+
+    optimize_dyn_actions = True
+
+
+    optimize_robot = True
+
+    use_penalty_based_friction = True
+
+    use_split_params = False
+
+    use_sqr_spring_stiffness = True
+
+    use_pre_proj_frictions = True
+
+
+
+    use_sqrt_dist = True
+    contact_maintaining_dist_thres = 0.2
+
+    robot_actions_diff_coef = 0.001
+
+
+    use_sdf_as_contact_dist = True
+
+
+    # 
+    use_contact_dist_as_sdf = False
+
+    use_glb_proj_delta = True
+
+
+
+# penetration_proj_k_to_robot = 30
+    penetrating_depth_penalty = 1.0
+    train_states = True
+
+
+
+    minn_dist_threshold = 0.000
+    obj_mass = 30.0
+
+
+    use_LBFGS = True
+    use_LBFGS = False
+
+    use_mano_hand_for_test = False # use the dynamic mano model here #
+
+    extract_delta_mesh = False
+    freeze_weights = True
+    gt_act_xs_def = False
+    use_bending_network =  True
+    ### for ts = 3 ###
+    # use_delta_bending = False
+    ### for ts = 3 ###
+   
+
+
+
+    sim_model_path = "rsc/shadow_hand_description/shadowhand_new.urdf"
+    mano_sim_model_path = "rsc/mano/mano_mean_wcollision_scaled_scaled_0_9507_nroot.urdf"
+
+    obj_sdf_fn = "data/grab/102/102_obj.npy" 
+    kinematic_mano_gt_sv_fn = "data/grab/102/102_sv_dict.npy"
+    scaled_obj_mesh_fn = "data/grab/102/102_obj.obj"
+
+    bending_net_type = "active_force_field_v18"    
+    sim_num_steps = 1000000
+    n_timesteps = 60
+    optim_sim_model_params_from_mano = False
+    penetration_determining = "sdf_of_canon"
+    train_with_forces_to_active = False
+    loss_scale_coef = 1000.0
+    use_same_contact_spring_k = False
+    use_optimizable_params = True # 
+    train_residual_friction = True
+    mano_mult_const_after_cent = 1.0
+    optimize_glb_transformations = True
+    no_friction_constraint = False
+    optimize_active_object = True
+    loss_tangential_diff_coef = 0
+    optimize_with_intermediates = True
+    using_delta_glb_trans = False
+    train_multi_seqs = False
+    use_split_network = True
+    use_delta_bending = True
+
+
+
+
+
+
+
+
+
+    ###### threshold, ks settings 1, optimize acts ######
+    # drive_pointset = "actions"
+    # fix_obj = True
+    # optimize_rules = False
+    # train_pointset_acts_via_deltas = True
+    # load_optimized_init_actions = "/data/xueyi/NeuS/exp/hand_test_routine_2_light_color_wtime_active_passive/wmask_reverse_value_totviews_tag_train_dyn_mano_hand_seq_102_mouse_optdynactions_points_optrobo_offsetdriven_optrules_multk100_wfixobj_optdelta_radius0d4_/checkpoints/ckpt_002000.pth"
+    # load_optimized_init_actions = "/data/xueyi/NeuS/exp/hand_test_routine_2_light_color_wtime_active_passive/wmask_reverse_value_totviews_tag_train_dyn_mano_hand_seq_102_mouse_optdynactions_points_optrobo_offsetdriven_optrules_multk100_wfixobj_optdelta_radius0d2_/checkpoints/ckpt_008000.pth"
+    ###### threshold, ks settings 1, optimize acts ######
+
+
+    ##### contact spring model settings ####
+    minn_dist_threshold_robot_to_obj = 0.1
+    penetration_proj_k_to_robot_friction = 10000000.0
+    penetration_proj_k_to_robot = 4000000.0
+    ##### contact spring model settings ####
+ 
+
+    # ###### Stage 1: threshold, ks settings 1, optimize offsets ######
+    # drive_pointset = "states"
+    # fix_obj = True
+    # optimize_rules = False
+    # train_pointset_acts_via_deltas = False
+    # load_optimized_init_actions = "ckpts/grab/102/dyn_mano_arti.pth"
+    # ###### Stage 1: threshold, ks settings 1, optimize offsets ######
+
+
+    # ###### Stage 2: threshold, ks settings 1, optimize acts ######
+    # drive_pointset = "actions"
+    # fix_obj = True
+    # optimize_rules = False
+    # train_pointset_acts_via_deltas = True
+    # load_optimized_init_actions = "ckpts/grab/102/dyn_mano_pointset_states.pt"
+    # ###### Stage 2: threshold, ks settings 1, optimize acts ######
+
+
+    # ###### Stage 3: threshold, ks settings 1, optimize params from acts ######
+    # drive_pointset = "actions"
+    # fix_obj = False
+    # optimize_rules = True
+    # train_pointset_acts_via_deltas = True
+    # load_optimized_init_actions = "ckpts/grab/102/dyn_mano_pointset_acts.pt"
+    # ##### model parameters optimized from the MANO hand trajectory #####
+    # ckpt_fn = "ckpts/grab/102/dyn_mano_opts.pt"
+    # ###### Stage 3: threshold, ks settings 1, optimize params from acts ######
+
+
+    # ###### Stage 4: threshold, ks settings 1, optimize acts from optimized params ######
+    # drive_pointset = "actions"
+    # fix_obj = False
+    # optimize_rules = False
+    # train_pointset_acts_via_deltas = True ## pointset acts via deltas ###
+    # ##### model parameters optimized from the MANO hand expanded set trajectory #####
+    # ckpt_fn = "ckpts/grab/102/dyn_mano_pointset_optimized_acts_optimized_ps.pth"
+    # load_optimized_init_actions = "ckpts/grab/102/dyn_mano_pointset_optimized_acts.pth"
+    # ###### Stage 4: threshold, ks settings 1, optimize acts from optimized params ######
+
+
+
+    ###### Retargeting Stage 1 ######
+    drive_pointset = "actions"
+    fix_obj = False
+    optimize_rules = False
+    train_pointset_acts_via_deltas = True
+    ##### model parameters optimized from the MANO hand expanded set trajectory #####
+    ckpt_fn = "ckpts/grab/102/dyn_mano_pointset_optimized_acts_optimized_ps_optimized_acts.pth"
+    load_optimized_init_actions = "ckpts/grab/102/dyn_mano_pointset_optimized_acts_optimized_ps_optimized_acts.pth"
+    load_optimized_init_transformations = "ckpts/grab/102/dyn_mano_shadow_arti.pth"
+    finger_cd_loss = 1.0
+    optimize_pointset_motion_only = True
+    ###### Retargeting Stage 1 ######
+
+
+
+    use_opt_rigid_translations=True
+
+    train_def = True
+    optimizable_rigid_translations=True
+    
+    nerf {
+        D = 8,
+        d_in = 4,
+        d_in_view = 3,
+        W = 256,
+        multires = 10,
+        multires_view = 4,
+        output_ch = 4,
+        skips=[4],
+        use_viewdirs=True
+    }
+
+    sdf_network {
+        d_out = 257,
+        d_in = 3,
+        d_hidden = 256,
+        n_layers = 8,
+        skip_in = [4],
+        multires = 6,
+        bias = 0.5,
+        scale = 1.0,
+        geometric_init = True,
+        weight_norm = True,
+    } 
+
+    variance_network {
+        init_val = 0.3
+    }
+
+    rendering_network {
+        d_feature = 256,
+        mode = idr,
+        d_in = 9,
+        d_out = 3,
+        d_hidden = 256,
+        n_layers = 4,
+        weight_norm = True,
+        multires_view = 4,
+        squeeze_out = True,
+    }
+
+    neus_renderer {
+        n_samples = 64,
+        n_importance = 64,
+        n_outside = 0,
+        up_sample_steps = 4 ,
+        perturb = 1.0,
+    }
+
+    bending_network {
+        multires = 6,
+        bending_latent_size = 32,
+        d_in = 3,
+        rigidity_hidden_dimensions = 64,
+        rigidity_network_depth = 5,
+        use_rigidity_network = False,
+        bending_n_timesteps = 10,
+    }
+}

confs_new/dyn_grab_pointset_points_dyn_retar_pts.conf

@@ -0,0 +1,281 @@
+general {
+
+
+    base_exp_dir = exp/CASE_NAME/wmask
+    
+    tag = "train_retargeted_shadow_hand_seq_102_mano_pointset_acts_"
+    tag = "train_retargeted_shadow_hand_seq_102_mano_pointset_acts_retar_to_shadow_pointset_"
+
+    recording = [
+        ./,
+        ./models
+    ]
+}
+
+dataset {
+    data_dir = public_data/CASE_NAME/
+    render_cameras_name = cameras_sphere.npz
+    object_cameras_name = cameras_sphere.npz
+    obj_idx = 102
+}
+
+train {
+    learning_rate = 5e-4
+    learning_rate_alpha = 0.05
+    end_iter = 300000
+
+    batch_size = 1024 #  64
+    validate_resolution_level = 4
+    warm_up_end = 5000
+    anneal_end = 0
+    use_white_bkgd = False
+
+    # save_freq = 10000
+    save_freq = 10000
+    val_freq = 20 #  2500
+    val_mesh_freq = 20 #  5000
+    report_freq = 10
+    ### igr weight ###
+    igr_weight = 0.1
+    mask_weight = 0.1
+}
+ 
+model {
+
+    optimize_dyn_actions = True
+
+
+    optimize_robot = True
+
+    use_penalty_based_friction = True
+
+    use_split_params = False
+
+    use_sqr_spring_stiffness = True
+
+    use_pre_proj_frictions = True
+
+
+
+    use_sqrt_dist = True
+    contact_maintaining_dist_thres = 0.2
+
+    robot_actions_diff_coef = 0.001
+
+
+    use_sdf_as_contact_dist = True
+
+
+    # 
+    use_contact_dist_as_sdf = False
+
+    use_glb_proj_delta = True
+
+
+
+# penetration_proj_k_to_robot = 30
+    penetrating_depth_penalty = 1.0
+    train_states = True
+
+
+
+    minn_dist_threshold = 0.000
+    obj_mass = 30.0
+
+
+    use_LBFGS = True
+    use_LBFGS = False
+
+    use_mano_hand_for_test = False # use the dynamic mano model here #
+
+    extract_delta_mesh = False
+    freeze_weights = True
+    gt_act_xs_def = False
+    use_bending_network =  True
+    ### for ts = 3 ###
+    # use_delta_bending = False
+    ### for ts = 3 ###
+   
+
+
+
+    sim_model_path = "rsc/shadow_hand_description/shadowhand_new.urdf"
+    mano_sim_model_path = "rsc/mano/mano_mean_wcollision_scaled_scaled_0_9507_nroot.urdf"
+
+    obj_sdf_fn = "data/grab/102/102_obj.npy" 
+    kinematic_mano_gt_sv_fn = "data/grab/102/102_sv_dict.npy"
+    scaled_obj_mesh_fn = "data/grab/102/102_obj.obj"
+
+    bending_net_type = "active_force_field_v18"    
+    sim_num_steps = 1000000
+    n_timesteps = 60
+    optim_sim_model_params_from_mano = False
+    penetration_determining = "sdf_of_canon"
+    train_with_forces_to_active = False
+    loss_scale_coef = 1000.0
+    use_same_contact_spring_k = False
+    use_optimizable_params = True # 
+    train_residual_friction = True
+    mano_mult_const_after_cent = 1.0
+    optimize_glb_transformations = True
+    no_friction_constraint = False
+    optimize_active_object = True
+    loss_tangential_diff_coef = 0
+    optimize_with_intermediates = True
+    using_delta_glb_trans = False
+    train_multi_seqs = False
+    use_split_network = True
+    use_delta_bending = True
+
+
+
+
+
+
+
+
+
+    ###### threshold, ks settings 1, optimize acts ######
+    # drive_pointset = "actions"
+    # fix_obj = True
+    # optimize_rules = False
+    # train_pointset_acts_via_deltas = True
+    # load_optimized_init_actions = "/data/xueyi/NeuS/exp/hand_test_routine_2_light_color_wtime_active_passive/wmask_reverse_value_totviews_tag_train_dyn_mano_hand_seq_102_mouse_optdynactions_points_optrobo_offsetdriven_optrules_multk100_wfixobj_optdelta_radius0d4_/checkpoints/ckpt_002000.pth"
+    # load_optimized_init_actions = "/data/xueyi/NeuS/exp/hand_test_routine_2_light_color_wtime_active_passive/wmask_reverse_value_totviews_tag_train_dyn_mano_hand_seq_102_mouse_optdynactions_points_optrobo_offsetdriven_optrules_multk100_wfixobj_optdelta_radius0d2_/checkpoints/ckpt_008000.pth"
+    ###### threshold, ks settings 1, optimize acts ######
+
+
+    ##### contact spring model settings ####
+    minn_dist_threshold_robot_to_obj = 0.1
+    penetration_proj_k_to_robot_friction = 10000000.0
+    penetration_proj_k_to_robot = 4000000.0
+    ##### contact spring model settings ####
+ 
+
+    # ###### Stage 1: threshold, ks settings 1, optimize offsets ######
+    # drive_pointset = "states"
+    # fix_obj = True
+    # optimize_rules = False
+    # train_pointset_acts_via_deltas = False
+    # load_optimized_init_actions = "ckpts/grab/102/dyn_mano_arti.pth"
+    # ###### Stage 1: threshold, ks settings 1, optimize offsets ######
+
+
+    # ###### Stage 2: threshold, ks settings 1, optimize acts ######
+    # drive_pointset = "actions"
+    # fix_obj = True
+    # optimize_rules = False
+    # train_pointset_acts_via_deltas = True
+    # load_optimized_init_actions = "ckpts/grab/102/dyn_mano_pointset_states.pt"
+    # ###### Stage 2: threshold, ks settings 1, optimize acts ######
+
+
+    # ###### Stage 3: threshold, ks settings 1, optimize params from acts ######
+    # drive_pointset = "actions"
+    # fix_obj = False
+    # optimize_rules = True
+    # train_pointset_acts_via_deltas = True
+    # load_optimized_init_actions = "ckpts/grab/102/dyn_mano_pointset_acts.pt"
+    # ##### model parameters optimized from the MANO hand trajectory #####
+    # ckpt_fn = "ckpts/grab/102/dyn_mano_opts.pt"
+    # ###### Stage 3: threshold, ks settings 1, optimize params from acts ######
+
+
+    # ###### Stage 4: threshold, ks settings 1, optimize acts from optimized params ######
+    # drive_pointset = "actions"
+    # fix_obj = False
+    # optimize_rules = False
+    # train_pointset_acts_via_deltas = True ## pointset acts via deltas ###
+    # ##### model parameters optimized from the MANO hand expanded set trajectory #####
+    # ckpt_fn = "ckpts/grab/102/dyn_mano_pointset_optimized_acts_optimized_ps.pth"
+    # load_optimized_init_actions = "ckpts/grab/102/dyn_mano_pointset_optimized_acts.pth"
+    # ###### Stage 4: threshold, ks settings 1, optimize acts from optimized params ######
+
+
+
+    ###### Retargeting Stage 1 ######
+    drive_pointset = "actions"
+    fix_obj = False
+    optimize_rules = False
+    train_pointset_acts_via_deltas = True
+    ##### model parameters optimized from the MANO hand expanded set trajectory #####
+    ckpt_fn = "ckpts/grab/102/dyn_mano_pointset_optimized_acts_optimized_ps_optimized_acts.pth"
+    load_optimized_init_actions = "ckpts/grab/102/dyn_mano_pointset_optimized_acts_optimized_ps_optimized_acts.pth"
+    load_optimized_init_transformations = "ckpts/grab/102/dyn_mano_shadow_arti_retar.pth"
+    finger_cd_loss = 1.0
+    optimize_pointset_motion_only = True
+    ###### Retargeting Stage 1 ######
+
+
+
+    ###### Retargeting Stage 2 ######
+    load_optimized_init_transformations = "ckpts/grab/102/dyn_mano_shadow_arti_retar_optimized_arti.pth"
+    ###### Retargeting Stage 2 ######
+
+
+
+
+    use_opt_rigid_translations=True
+
+    train_def = True
+    optimizable_rigid_translations=True
+    
+    nerf {
+        D = 8,
+        d_in = 4,
+        d_in_view = 3,
+        W = 256,
+        multires = 10,
+        multires_view = 4,
+        output_ch = 4,
+        skips=[4],
+        use_viewdirs=True
+    }
+
+    sdf_network {
+        d_out = 257,
+        d_in = 3,
+        d_hidden = 256,
+        n_layers = 8,
+        skip_in = [4],
+        multires = 6,
+        bias = 0.5,
+        scale = 1.0,
+        geometric_init = True,
+        weight_norm = True,
+    } 
+
+    variance_network {
+        init_val = 0.3
+    }
+
+    rendering_network {
+        d_feature = 256,
+        mode = idr,
+        d_in = 9,
+        d_out = 3,
+        d_hidden = 256,
+        n_layers = 4,
+        weight_norm = True,
+        multires_view = 4,
+        squeeze_out = True,
+    }
+
+    neus_renderer {
+        n_samples = 64,
+        n_importance = 64,
+        n_outside = 0,
+        up_sample_steps = 4 ,
+        perturb = 1.0,
+    }
+
+    bending_network {
+        multires = 6,
+        bending_latent_size = 32,
+        d_in = 3,
+        rigidity_hidden_dimensions = 64,
+        rigidity_network_depth = 5,
+        use_rigidity_network = False,
+        bending_n_timesteps = 10,
+    }
+}

confs_new/dyn_grab_pointset_points_dyn_retar_pts_opts.conf

@@ -0,0 +1,287 @@
+general {
+
+
+    base_exp_dir = exp/CASE_NAME/wmask
+    
+    tag = "train_retargeted_shadow_hand_seq_102_mano_pointset_acts_"
+    tag = "train_retargeted_shadow_hand_seq_102_mano_pointset_acts_retar_to_shadow_pointset_"
+    tag = "train_retargeted_shadow_hand_seq_102_mano_pointset_acts_retar_to_shadow_pointset_optrules_"
+
+    recording = [
+        ./,
+        ./models
+    ]
+}
+
+dataset {
+    data_dir = public_data/CASE_NAME/
+    render_cameras_name = cameras_sphere.npz
+    object_cameras_name = cameras_sphere.npz
+    obj_idx = 102
+}
+
+train {
+    learning_rate = 5e-4
+    learning_rate_alpha = 0.05
+    end_iter = 300000
+
+    batch_size = 1024 #  64
+    validate_resolution_level = 4
+    warm_up_end = 5000
+    anneal_end = 0
+    use_white_bkgd = False
+
+    save_freq = 10000
+    val_freq = 20
+    val_mesh_freq = 20
+    report_freq = 10
+    igr_weight = 0.1
+    mask_weight = 0.1
+}
+ 
+model {
+
+    optimize_dyn_actions = True
+
+
+    optimize_robot = True
+
+    use_penalty_based_friction = True
+
+    use_split_params = False
+
+    use_sqr_spring_stiffness = True
+
+    use_pre_proj_frictions = True
+
+
+
+    use_sqrt_dist = True
+    contact_maintaining_dist_thres = 0.2
+
+    robot_actions_diff_coef = 0.001
+
+
+    use_sdf_as_contact_dist = True
+
+
+    # 
+    use_contact_dist_as_sdf = False
+
+    use_glb_proj_delta = True
+
+
+
+# penetration_proj_k_to_robot = 30
+    penetrating_depth_penalty = 1.0
+    train_states = True
+
+
+
+    minn_dist_threshold = 0.000
+    obj_mass = 30.0
+
+
+    use_LBFGS = True
+    use_LBFGS = False
+
+    use_mano_hand_for_test = False # use the dynamic mano model here #
+
+    extract_delta_mesh = False
+    freeze_weights = True
+    gt_act_xs_def = False
+    use_bending_network =  True
+    ### for ts = 3 ###
+    # use_delta_bending = False
+    ### for ts = 3 ###
+   
+
+
+
+    sim_model_path = "rsc/shadow_hand_description/shadowhand_new.urdf"
+    mano_sim_model_path = "rsc/mano/mano_mean_wcollision_scaled_scaled_0_9507_nroot.urdf"
+
+    obj_sdf_fn = "data/grab/102/102_obj.npy" 
+    kinematic_mano_gt_sv_fn = "data/grab/102/102_sv_dict.npy"
+    scaled_obj_mesh_fn = "data/grab/102/102_obj.obj"
+
+    bending_net_type = "active_force_field_v18"    
+    sim_num_steps = 1000000
+    n_timesteps = 60
+    optim_sim_model_params_from_mano = False
+    penetration_determining = "sdf_of_canon"
+    train_with_forces_to_active = False
+    loss_scale_coef = 1000.0
+    use_same_contact_spring_k = False
+    use_optimizable_params = True # 
+    train_residual_friction = True
+    mano_mult_const_after_cent = 1.0
+    optimize_glb_transformations = True
+    no_friction_constraint = False
+    optimize_active_object = True
+    loss_tangential_diff_coef = 0
+    optimize_with_intermediates = True
+    using_delta_glb_trans = False
+    train_multi_seqs = False
+    use_split_network = True
+    use_delta_bending = True
+
+
+
+
+
+
+
+
+
+    ###### threshold, ks settings 1, optimize acts ######
+    # drive_pointset = "actions"
+    # fix_obj = True
+    # optimize_rules = False
+    # train_pointset_acts_via_deltas = True
+    # load_optimized_init_actions = "/data/xueyi/NeuS/exp/hand_test_routine_2_light_color_wtime_active_passive/wmask_reverse_value_totviews_tag_train_dyn_mano_hand_seq_102_mouse_optdynactions_points_optrobo_offsetdriven_optrules_multk100_wfixobj_optdelta_radius0d4_/checkpoints/ckpt_002000.pth"
+    # load_optimized_init_actions = "/data/xueyi/NeuS/exp/hand_test_routine_2_light_color_wtime_active_passive/wmask_reverse_value_totviews_tag_train_dyn_mano_hand_seq_102_mouse_optdynactions_points_optrobo_offsetdriven_optrules_multk100_wfixobj_optdelta_radius0d2_/checkpoints/ckpt_008000.pth"
+    ###### threshold, ks settings 1, optimize acts ######
+
+
+    ##### contact spring model settings ####
+    minn_dist_threshold_robot_to_obj = 0.1
+    penetration_proj_k_to_robot_friction = 10000000.0
+    penetration_proj_k_to_robot = 4000000.0
+    ##### contact spring model settings ####
+ 
+
+    # ###### Stage 1: threshold, ks settings 1, optimize offsets ######
+    # drive_pointset = "states"
+    # fix_obj = True
+    # optimize_rules = False
+    # train_pointset_acts_via_deltas = False
+    # load_optimized_init_actions = "ckpts/grab/102/dyn_mano_arti.pth"
+    # ###### Stage 1: threshold, ks settings 1, optimize offsets ######
+
+
+    # ###### Stage 2: threshold, ks settings 1, optimize acts ######
+    # drive_pointset = "actions"
+    # fix_obj = True
+    # optimize_rules = False
+    # train_pointset_acts_via_deltas = True
+    # load_optimized_init_actions = "ckpts/grab/102/dyn_mano_pointset_states.pt"
+    # ###### Stage 2: threshold, ks settings 1, optimize acts ######
+
+
+    # ###### Stage 3: threshold, ks settings 1, optimize params from acts ######
+    # drive_pointset = "actions"
+    # fix_obj = False
+    # optimize_rules = True
+    # train_pointset_acts_via_deltas = True
+    # load_optimized_init_actions = "ckpts/grab/102/dyn_mano_pointset_acts.pt"
+    # ##### model parameters optimized from the MANO hand trajectory #####
+    # ckpt_fn = "ckpts/grab/102/dyn_mano_opts.pt"
+    # ###### Stage 3: threshold, ks settings 1, optimize params from acts ######
+
+
+    # ###### Stage 4: threshold, ks settings 1, optimize acts from optimized params ######
+    # drive_pointset = "actions"
+    # fix_obj = False
+    # optimize_rules = False
+    # train_pointset_acts_via_deltas = True ## pointset acts via deltas ###
+    # ##### model parameters optimized from the MANO hand expanded set trajectory #####
+    # ckpt_fn = "ckpts/grab/102/dyn_mano_pointset_optimized_acts_optimized_ps.pth"
+    # load_optimized_init_actions = "ckpts/grab/102/dyn_mano_pointset_optimized_acts.pth"
+    # ###### Stage 4: threshold, ks settings 1, optimize acts from optimized params ######
+
+
+
+    ###### Retargeting Stage 1 ######
+    drive_pointset = "actions"
+    fix_obj = False
+    optimize_rules = False
+    train_pointset_acts_via_deltas = True
+    ##### model parameters optimized from the MANO hand expanded set trajectory #####
+    ckpt_fn = "ckpts/grab/102/dyn_mano_pointset_optimized_acts_optimized_ps_optimized_acts.pth"
+    load_optimized_init_actions = "ckpts/grab/102/dyn_mano_pointset_optimized_acts_optimized_ps_optimized_acts.pth"
+    load_optimized_init_transformations = "ckpts/grab/102/dyn_mano_shadow_arti_retar.pth"
+    finger_cd_loss = 1.0
+    optimize_pointset_motion_only = True
+    ###### Retargeting Stage 1 ######
+
+
+    ###### Retargeting Stage 2 ######
+    load_optimized_init_transformations = "ckpts/grab/102/dyn_mano_shadow_arti_retar_optimized_arti.pth"
+    ###### Retargeting Stage 2 ######
+
+
+    ###### Retargeting Stage 3 ######
+    load_optimized_init_transformations = "ckpts/grab/102/dyn_mano_shadow_arti_retar_optimized_arti_optimized_pts.pth"
+    optimize_anchored_pts = False
+    optimize_rules = True
+    optimize_pointset_motion_only = False
+    ###### Retargeting Stage 3 ######
+    
+
+
+
+    use_opt_rigid_translations=True
+
+    train_def = True
+    optimizable_rigid_translations=True
+    
+    nerf {
+        D = 8,
+        d_in = 4,
+        d_in_view = 3,
+        W = 256,
+        multires = 10,
+        multires_view = 4,
+        output_ch = 4,
+        skips=[4],
+        use_viewdirs=True
+    }
+
+    sdf_network {
+        d_out = 257,
+        d_in = 3,
+        d_hidden = 256,
+        n_layers = 8,
+        skip_in = [4],
+        multires = 6,
+        bias = 0.5,
+        scale = 1.0,
+        geometric_init = True,
+        weight_norm = True,
+    } 
+
+    variance_network {
+        init_val = 0.3
+    }
+
+    rendering_network {
+        d_feature = 256,
+        mode = idr,
+        d_in = 9,
+        d_out = 3,
+        d_hidden = 256,
+        n_layers = 4,
+        weight_norm = True,
+        multires_view = 4,
+        squeeze_out = True,
+    }
+
+    neus_renderer {
+        n_samples = 64,
+        n_importance = 64,
+        n_outside = 0,
+        up_sample_steps = 4 ,
+        perturb = 1.0,
+    }
+
+    bending_network {
+        multires = 6,
+        bending_latent_size = 32,
+        d_in = 3,
+        rigidity_hidden_dimensions = 64,
+        rigidity_network_depth = 5,
+        use_rigidity_network = False,
+        bending_n_timesteps = 10,
+    }
+}

confs_new/dyn_grab_pointset_points_dyn_s1.conf

@@ -0,0 +1,256 @@
+general {
+
+
+    base_exp_dir = exp/CASE_NAME/wmask
+    
+    tag = "train_retargeted_shadow_hand_seq_102_mano_pointset_acts_optstates_"
+
+    recording = [
+        ./,
+        ./models
+    ]
+}
+
+dataset {
+    data_dir = public_data/CASE_NAME/
+    render_cameras_name = cameras_sphere.npz
+    object_cameras_name = cameras_sphere.npz
+    obj_idx = 102
+}
+
+train {
+    learning_rate = 5e-4
+    learning_rate_alpha = 0.05
+    end_iter = 300000
+
+    batch_size = 1024
+    validate_resolution_level = 4
+    warm_up_end = 5000
+    anneal_end = 0
+    use_white_bkgd = False
+
+    # save_freq = 10000
+    save_freq = 10000
+    val_freq = 20
+    val_mesh_freq = 20
+    report_freq = 10
+    igr_weight = 0.1
+    mask_weight = 0.1
+}
+ 
+model {
+
+    optimize_dyn_actions = True
+
+
+    optimize_robot = True
+
+    use_penalty_based_friction = True
+
+    use_split_params = False
+
+    use_sqr_spring_stiffness = True
+
+    use_pre_proj_frictions = True
+
+
+
+    use_sqrt_dist = True
+    contact_maintaining_dist_thres = 0.2
+
+    robot_actions_diff_coef = 0.001
+
+
+    use_sdf_as_contact_dist = True
+
+
+    # 
+    use_contact_dist_as_sdf = False
+
+    use_glb_proj_delta = True
+
+
+
+# penetration_proj_k_to_robot = 30
+    penetrating_depth_penalty = 1.0
+    train_states = True
+
+
+
+    minn_dist_threshold = 0.000
+    obj_mass = 30.0
+
+
+    use_LBFGS = True
+    use_LBFGS = False
+
+    use_mano_hand_for_test = False # use the dynamic mano model here #
+
+    extract_delta_mesh = False
+    freeze_weights = True
+    gt_act_xs_def = False
+    use_bending_network =  True
+    ### for ts = 3 ###
+    # use_delta_bending = False
+    ### for ts = 3 ###
+   
+
+
+
+    sim_model_path = "rsc/shadow_hand_description/shadowhand_new.urdf"
+    mano_sim_model_path = "rsc/mano/mano_mean_wcollision_scaled_scaled_0_9507_nroot.urdf"
+
+    obj_sdf_fn = "data/grab/102/102_obj.npy" 
+    kinematic_mano_gt_sv_fn = "data/grab/102/102_sv_dict.npy"
+    scaled_obj_mesh_fn = "data/grab/102/102_obj.obj"
+
+    bending_net_type = "active_force_field_v18"    
+    sim_num_steps = 1000000
+    n_timesteps = 60
+    optim_sim_model_params_from_mano = False
+    penetration_determining = "sdf_of_canon"
+    train_with_forces_to_active = False
+    loss_scale_coef = 1000.0
+    use_same_contact_spring_k = False
+    use_optimizable_params = True # 
+    train_residual_friction = True
+    mano_mult_const_after_cent = 1.0
+    optimize_glb_transformations = True
+    no_friction_constraint = False
+    optimize_active_object = True
+    loss_tangential_diff_coef = 0
+    optimize_with_intermediates = True
+    using_delta_glb_trans = False
+    train_multi_seqs = False
+    use_split_network = True
+    use_delta_bending = True
+
+
+
+
+
+
+
+
+
+    ###### threshold, ks settings 1, optimize acts ######
+    # drive_pointset = "actions"
+    # fix_obj = True
+    # optimize_rules = False
+    # train_pointset_acts_via_deltas = True
+    # load_optimized_init_actions = "/data/xueyi/NeuS/exp/hand_test_routine_2_light_color_wtime_active_passive/wmask_reverse_value_totviews_tag_train_dyn_mano_hand_seq_102_mouse_optdynactions_points_optrobo_offsetdriven_optrules_multk100_wfixobj_optdelta_radius0d4_/checkpoints/ckpt_002000.pth"
+    # load_optimized_init_actions = "/data/xueyi/NeuS/exp/hand_test_routine_2_light_color_wtime_active_passive/wmask_reverse_value_totviews_tag_train_dyn_mano_hand_seq_102_mouse_optdynactions_points_optrobo_offsetdriven_optrules_multk100_wfixobj_optdelta_radius0d2_/checkpoints/ckpt_008000.pth"
+    ###### threshold, ks settings 1, optimize acts ######
+
+
+    ##### contact spring model settings ####
+    minn_dist_threshold_robot_to_obj = 0.1
+    penetration_proj_k_to_robot_friction = 10000000.0
+    penetration_proj_k_to_robot = 4000000.0
+    ##### contact spring model settings ####
+ 
+
+    ###### Stage 1: threshold, ks settings 1, optimize offsets ######
+    drive_pointset = "states"
+    fix_obj = True
+    optimize_rules = False
+    train_pointset_acts_via_deltas = False
+    load_optimized_init_actions = "ckpts/grab/102/dyn_mano_arti.pth"
+    ###### Stage 1: threshold, ks settings 1, optimize offsets ######
+
+
+    # ###### Stage 2: threshold, ks settings 1, optimize acts ######
+    # drive_pointset = "actions"
+    # fix_obj = True
+    # optimize_rules = False
+    # train_pointset_acts_via_deltas = True
+    # load_optimized_init_actions = "ckpts/grab/102/dyn_mano_pointset_states.pt"
+    # ###### Stage 2: threshold, ks settings 1, optimize acts ######
+
+
+    # ###### Stage 3: threshold, ks settings 1, optimize params from acts ######
+    # drive_pointset = "actions"
+    # fix_obj = False
+    # optimize_rules = True
+    # train_pointset_acts_via_deltas = True
+    # load_optimized_init_actions = "ckpts/grab/102/dyn_mano_pointset_acts.pt"
+    # ##### model parameters optimized from the MANO hand trajectory #####
+    # ckpt_fn = "ckpts/grab/102/dyn_mano_opts.pt"
+    # ###### Stage 3: threshold, ks settings 1, optimize params from acts ######
+
+
+    # ###### Stage 4: threshold, ks settings 1, optimize acts from optimized params ######
+    # drive_pointset = "actions"
+    # fix_obj = False
+    # optimize_rules = False
+    # train_pointset_acts_via_deltas = True ## pointset acts via deltas ###
+    # ##### model parameters optimized from the MANO hand expanded set trajectory #####
+    # ckpt_fn = "ckpts/grab/102/dyn_mano_pointset_optimized_acts_optimized_ps.pth"
+    # load_optimized_init_actions = "ckpts/grab/102/dyn_mano_pointset_optimized_acts.pth"
+    # ###### Stage 4: threshold, ks settings 1, optimize acts from optimized params ######
+
+
+    use_opt_rigid_translations=True
+
+    train_def = True
+    optimizable_rigid_translations=True
+    
+    nerf {
+        D = 8,
+        d_in = 4,
+        d_in_view = 3,
+        W = 256,
+        multires = 10,
+        multires_view = 4,
+        output_ch = 4,
+        skips=[4],
+        use_viewdirs=True
+    }
+
+    sdf_network {
+        d_out = 257,
+        d_in = 3,
+        d_hidden = 256,
+        n_layers = 8,
+        skip_in = [4],
+        multires = 6,
+        bias = 0.5,
+        scale = 1.0,
+        geometric_init = True,
+        weight_norm = True,
+    } 
+
+    variance_network {
+        init_val = 0.3
+    }
+
+    rendering_network {
+        d_feature = 256,
+        mode = idr,
+        d_in = 9,
+        d_out = 3,
+        d_hidden = 256,
+        n_layers = 4,
+        weight_norm = True,
+        multires_view = 4,
+        squeeze_out = True,
+    }
+
+    neus_renderer {
+        n_samples = 64,
+        n_importance = 64,
+        n_outside = 0,
+        up_sample_steps = 4 ,
+        perturb = 1.0,
+    }
+
+    bending_network {
+        multires = 6,
+        bending_latent_size = 32,
+        d_in = 3,
+        rigidity_hidden_dimensions = 64,
+        rigidity_network_depth = 5,
+        use_rigidity_network = False,
+        bending_n_timesteps = 10,
+    }
+}

confs_new/dyn_grab_pointset_points_dyn_s2.conf

@@ -0,0 +1,258 @@
+general {
+
+
+    base_exp_dir = exp/CASE_NAME/wmask
+    
+    tag = "train_retargeted_shadow_hand_seq_102_mano_pointset_acts_optstates_optacts_"
+
+    recording = [
+        ./,
+        ./models
+    ]
+}
+
+dataset {
+    data_dir = public_data/CASE_NAME/
+    render_cameras_name = cameras_sphere.npz
+    object_cameras_name = cameras_sphere.npz
+    obj_idx = 102
+}
+
+train {
+    learning_rate = 5e-4
+    learning_rate_alpha = 0.05
+    end_iter = 300000
+
+    batch_size = 1024 #  64
+    validate_resolution_level = 4
+    warm_up_end = 5000
+    anneal_end = 0
+    use_white_bkgd = False
+
+    # save_freq = 10000
+    save_freq = 10000
+    val_freq = 20 #  2500
+    val_mesh_freq = 20 #  5000
+    report_freq = 10
+    ### igr weight ###
+    igr_weight = 0.1
+    mask_weight = 0.1
+}
+ 
+model {
+
+    optimize_dyn_actions = True
+
+
+    optimize_robot = True
+
+    use_penalty_based_friction = True
+
+    use_split_params = False
+
+    use_sqr_spring_stiffness = True
+
+    use_pre_proj_frictions = True
+
+
+
+    use_sqrt_dist = True
+    contact_maintaining_dist_thres = 0.2
+
+    robot_actions_diff_coef = 0.001
+
+
+    use_sdf_as_contact_dist = True
+
+
+    # 
+    use_contact_dist_as_sdf = False
+
+    use_glb_proj_delta = True
+
+
+
+# penetration_proj_k_to_robot = 30
+    penetrating_depth_penalty = 1.0
+    train_states = True
+
+
+
+    minn_dist_threshold = 0.000
+    obj_mass = 30.0
+
+
+    use_LBFGS = True
+    use_LBFGS = False
+
+    use_mano_hand_for_test = False # use the dynamic mano model here #
+
+    extract_delta_mesh = False
+    freeze_weights = True
+    gt_act_xs_def = False
+    use_bending_network =  True
+    ### for ts = 3 ###
+    # use_delta_bending = False
+    ### for ts = 3 ###
+   
+
+
+
+    sim_model_path = "rsc/shadow_hand_description/shadowhand_new.urdf"
+    mano_sim_model_path = "rsc/mano/mano_mean_wcollision_scaled_scaled_0_9507_nroot.urdf"
+
+    obj_sdf_fn = "data/grab/102/102_obj.npy" 
+    kinematic_mano_gt_sv_fn = "data/grab/102/102_sv_dict.npy"
+    scaled_obj_mesh_fn = "data/grab/102/102_obj.obj"
+
+    bending_net_type = "active_force_field_v18"    
+    sim_num_steps = 1000000
+    n_timesteps = 60
+    optim_sim_model_params_from_mano = False
+    penetration_determining = "sdf_of_canon"
+    train_with_forces_to_active = False
+    loss_scale_coef = 1000.0
+    use_same_contact_spring_k = False
+    use_optimizable_params = True # 
+    train_residual_friction = True
+    mano_mult_const_after_cent = 1.0
+    optimize_glb_transformations = True
+    no_friction_constraint = False
+    optimize_active_object = True
+    loss_tangential_diff_coef = 0
+    optimize_with_intermediates = True
+    using_delta_glb_trans = False
+    train_multi_seqs = False
+    use_split_network = True
+    use_delta_bending = True
+
+
+
+
+
+
+
+
+
+    ###### threshold, ks settings 1, optimize acts ######
+    # drive_pointset = "actions"
+    # fix_obj = True
+    # optimize_rules = False
+    # train_pointset_acts_via_deltas = True
+    # load_optimized_init_actions = "/data/xueyi/NeuS/exp/hand_test_routine_2_light_color_wtime_active_passive/wmask_reverse_value_totviews_tag_train_dyn_mano_hand_seq_102_mouse_optdynactions_points_optrobo_offsetdriven_optrules_multk100_wfixobj_optdelta_radius0d4_/checkpoints/ckpt_002000.pth"
+    # load_optimized_init_actions = "/data/xueyi/NeuS/exp/hand_test_routine_2_light_color_wtime_active_passive/wmask_reverse_value_totviews_tag_train_dyn_mano_hand_seq_102_mouse_optdynactions_points_optrobo_offsetdriven_optrules_multk100_wfixobj_optdelta_radius0d2_/checkpoints/ckpt_008000.pth"
+    ###### threshold, ks settings 1, optimize acts ######
+
+
+    ##### contact spring model settings ####
+    minn_dist_threshold_robot_to_obj = 0.1
+    penetration_proj_k_to_robot_friction = 10000000.0
+    penetration_proj_k_to_robot = 4000000.0
+    ##### contact spring model settings ####
+ 
+
+    ###### Stage 1: threshold, ks settings 1, optimize offsets ######
+    drive_pointset = "states"
+    fix_obj = True
+    optimize_rules = False
+    train_pointset_acts_via_deltas = False
+    load_optimized_init_actions = "ckpts/grab/102/dyn_mano_arti.pth"
+    ###### Stage 1: threshold, ks settings 1, optimize offsets ######
+
+
+    ###### Stage 2: threshold, ks settings 1, optimize acts ######
+    drive_pointset = "actions"
+    fix_obj = True
+    optimize_rules = False
+    train_pointset_acts_via_deltas = True
+    load_optimized_init_actions = "ckpts/grab/102/dyn_mano_pointset_states.pt" ## pre-optimized ckpts
+    load_optimized_init_actions = "exp/hand_test_routine_2_light_color_wtime_active_passive/wmask_reverse_value_totviews_tag_train_retargeted_shadow_hand_seq_102_mano_pointset_acts_optstates_/checkpoints/ckpt_004000.pth"
+    ###### Stage 2: threshold, ks settings 1, optimize acts ######
+
+
+    # ###### Stage 3: threshold, ks settings 1, optimize params from acts ######
+    # drive_pointset = "actions"
+    # fix_obj = False
+    # optimize_rules = True
+    # train_pointset_acts_via_deltas = True
+    # load_optimized_init_actions = "ckpts/grab/102/dyn_mano_pointset_acts.pt"
+    # ##### model parameters optimized from the MANO hand trajectory #####
+    # ckpt_fn = "ckpts/grab/102/dyn_mano_opts.pt"
+    # ###### Stage 3: threshold, ks settings 1, optimize params from acts ######
+
+
+    # ###### Stage 4: threshold, ks settings 1, optimize acts from optimized params ######
+    # drive_pointset = "actions"
+    # fix_obj = False
+    # optimize_rules = False
+    # train_pointset_acts_via_deltas = True ## pointset acts via deltas ###
+    # ##### model parameters optimized from the MANO hand expanded set trajectory #####
+    # ckpt_fn = "ckpts/grab/102/dyn_mano_pointset_optimized_acts_optimized_ps.pth"
+    # load_optimized_init_actions = "ckpts/grab/102/dyn_mano_pointset_optimized_acts.pth"
+    # ###### Stage 4: threshold, ks settings 1, optimize acts from optimized params ######
+
+
+    use_opt_rigid_translations=True
+
+    train_def = True
+    optimizable_rigid_translations=True
+    
+    nerf {
+        D = 8,
+        d_in = 4,
+        d_in_view = 3,
+        W = 256,
+        multires = 10,
+        multires_view = 4,
+        output_ch = 4,
+        skips=[4],
+        use_viewdirs=True
+    }
+
+    sdf_network {
+        d_out = 257,
+        d_in = 3,
+        d_hidden = 256,
+        n_layers = 8,
+        skip_in = [4],
+        multires = 6,
+        bias = 0.5,
+        scale = 1.0,
+        geometric_init = True,
+        weight_norm = True,
+    } 
+
+    variance_network {
+        init_val = 0.3
+    }
+
+    rendering_network {
+        d_feature = 256,
+        mode = idr,
+        d_in = 9,
+        d_out = 3,
+        d_hidden = 256,
+        n_layers = 4,
+        weight_norm = True,
+        multires_view = 4,
+        squeeze_out = True,
+    }
+
+    neus_renderer {
+        n_samples = 64,
+        n_importance = 64,
+        n_outside = 0,
+        up_sample_steps = 4 ,
+        perturb = 1.0,
+    }
+
+    bending_network {
+        multires = 6,
+        bending_latent_size = 32,
+        d_in = 3,
+        rigidity_hidden_dimensions = 64,
+        rigidity_network_depth = 5,
+        use_rigidity_network = False,
+        bending_n_timesteps = 10,
+    }
+}

confs_new/dyn_grab_pointset_points_dyn_s3.conf

@@ -0,0 +1,259 @@
+general {
+
+
+    base_exp_dir = exp/CASE_NAME/wmask
+    
+    tag = "train_retargeted_shadow_hand_seq_102_mano_pointset_acts_optstates_optacts_optsysps_"
+
+    recording = [
+        ./,
+        ./models
+    ]
+}
+
+dataset {
+    data_dir = public_data/CASE_NAME/
+    render_cameras_name = cameras_sphere.npz
+    object_cameras_name = cameras_sphere.npz
+    obj_idx = 102
+}
+
+train {
+    learning_rate = 5e-4
+    learning_rate_alpha = 0.05
+    end_iter = 300000
+
+    batch_size = 1024 #  64
+    validate_resolution_level = 4
+    warm_up_end = 5000
+    anneal_end = 0
+    use_white_bkgd = False
+
+    # save_freq = 10000
+    save_freq = 10000
+    val_freq = 20 #  2500
+    val_mesh_freq = 20 #  5000
+    report_freq = 10
+    ### igr weight ###
+    igr_weight = 0.1
+    mask_weight = 0.1
+}
+ 
+model {
+
+    optimize_dyn_actions = True
+
+
+    optimize_robot = True
+
+    use_penalty_based_friction = True
+
+    use_split_params = False
+
+    use_sqr_spring_stiffness = True
+
+    use_pre_proj_frictions = True
+
+
+
+    use_sqrt_dist = True
+    contact_maintaining_dist_thres = 0.2
+
+    robot_actions_diff_coef = 0.001
+
+
+    use_sdf_as_contact_dist = True
+
+
+    # 
+    use_contact_dist_as_sdf = False
+
+    use_glb_proj_delta = True
+
+
+
+# penetration_proj_k_to_robot = 30
+    penetrating_depth_penalty = 1.0
+    train_states = True
+
+
+
+    minn_dist_threshold = 0.000
+    obj_mass = 30.0
+
+
+    use_LBFGS = True
+    use_LBFGS = False
+
+    use_mano_hand_for_test = False # use the dynamic mano model here #
+
+    extract_delta_mesh = False
+    freeze_weights = True
+    gt_act_xs_def = False
+    use_bending_network =  True
+    ### for ts = 3 ###
+    # use_delta_bending = False
+    ### for ts = 3 ###
+   
+
+
+
+    sim_model_path = "rsc/shadow_hand_description/shadowhand_new.urdf"
+    mano_sim_model_path = "rsc/mano/mano_mean_wcollision_scaled_scaled_0_9507_nroot.urdf"
+
+    obj_sdf_fn = "data/grab/102/102_obj.npy" 
+    kinematic_mano_gt_sv_fn = "data/grab/102/102_sv_dict.npy"
+    scaled_obj_mesh_fn = "data/grab/102/102_obj.obj"
+
+    bending_net_type = "active_force_field_v18"    
+    sim_num_steps = 1000000
+    n_timesteps = 60
+    optim_sim_model_params_from_mano = False
+    penetration_determining = "sdf_of_canon"
+    train_with_forces_to_active = False
+    loss_scale_coef = 1000.0
+    use_same_contact_spring_k = False
+    use_optimizable_params = True # 
+    train_residual_friction = True
+    mano_mult_const_after_cent = 1.0
+    optimize_glb_transformations = True
+    no_friction_constraint = False
+    optimize_active_object = True
+    loss_tangential_diff_coef = 0
+    optimize_with_intermediates = True
+    using_delta_glb_trans = False
+    train_multi_seqs = False
+    use_split_network = True
+    use_delta_bending = True
+
+
+
+
+
+
+
+
+
+    ###### threshold, ks settings 1, optimize acts ######
+    # drive_pointset = "actions"
+    # fix_obj = True
+    # optimize_rules = False
+    # train_pointset_acts_via_deltas = True
+    # load_optimized_init_actions = "/data/xueyi/NeuS/exp/hand_test_routine_2_light_color_wtime_active_passive/wmask_reverse_value_totviews_tag_train_dyn_mano_hand_seq_102_mouse_optdynactions_points_optrobo_offsetdriven_optrules_multk100_wfixobj_optdelta_radius0d4_/checkpoints/ckpt_002000.pth"
+    # load_optimized_init_actions = "/data/xueyi/NeuS/exp/hand_test_routine_2_light_color_wtime_active_passive/wmask_reverse_value_totviews_tag_train_dyn_mano_hand_seq_102_mouse_optdynactions_points_optrobo_offsetdriven_optrules_multk100_wfixobj_optdelta_radius0d2_/checkpoints/ckpt_008000.pth"
+    ###### threshold, ks settings 1, optimize acts ######
+
+
+    ##### contact spring model settings ####
+    minn_dist_threshold_robot_to_obj = 0.1
+    penetration_proj_k_to_robot_friction = 10000000.0
+    penetration_proj_k_to_robot = 4000000.0
+    ##### contact spring model settings ####
+ 
+
+    ###### Stage 1: threshold, ks settings 1, optimize offsets ######
+    drive_pointset = "states"
+    fix_obj = True
+    optimize_rules = False
+    train_pointset_acts_via_deltas = False
+    load_optimized_init_actions = "ckpts/grab/102/dyn_mano_arti.pth"
+    ###### Stage 1: threshold, ks settings 1, optimize offsets ######
+
+
+    ###### Stage 2: threshold, ks settings 1, optimize acts ######
+    drive_pointset = "actions"
+    fix_obj = True
+    optimize_rules = False
+    train_pointset_acts_via_deltas = True
+    load_optimized_init_actions = "ckpts/grab/102/dyn_mano_pointset_states.pt"
+    ###### Stage 2: threshold, ks settings 1, optimize acts ######
+
+
+    ###### Stage 3: threshold, ks settings 1, optimize params from acts ######
+    drive_pointset = "actions"
+    fix_obj = False
+    optimize_rules = True
+    train_pointset_acts_via_deltas = True
+    load_optimized_init_actions = "ckpts/grab/102/dyn_mano_pointset_acts.pt"
+    load_optimized_init_actions = "exp/hand_test_routine_2_light_color_wtime_active_passive/wmask_reverse_value_totviews_tag_train_retargeted_shadow_hand_seq_102_mano_pointset_acts_optstates_optacts_/checkpoints/ckpt_004000.pth"
+    ##### model parameters optimized from the MANO hand trajectory #####
+    ckpt_fn = "ckpts/grab/102/dyn_mano_opts.pt"
+    ckpt_fn = "exp/hand_test_routine_2_light_color_wtime_active_passive/wmask_reverse_value_totviews_tag_train_retargeted_shadow_hand_seq_102_mano_pointset_acts_optstates_optacts_/checkpoints/ckpt_004000.pth"
+    ###### Stage 3: threshold, ks settings 1, optimize params from acts ######
+
+
+    # ###### Stage 4: threshold, ks settings 1, optimize acts from optimized params ######
+    # drive_pointset = "actions"
+    # fix_obj = False
+    # optimize_rules = False
+    # train_pointset_acts_via_deltas = True ## pointset acts via deltas ###
+    # ##### model parameters optimized from the MANO hand expanded set trajectory #####
+    # ckpt_fn = "ckpts/grab/102/dyn_mano_pointset_optimized_acts_optimized_ps.pth"
+    # load_optimized_init_actions = "ckpts/grab/102/dyn_mano_pointset_optimized_acts.pth"
+    # ###### Stage 4: threshold, ks settings 1, optimize acts from optimized params ######
+
+
+    use_opt_rigid_translations=True
+
+    train_def = True
+    optimizable_rigid_translations=True
+    
+    nerf {
+        D = 8,
+        d_in = 4,
+        d_in_view = 3,
+        W = 256,
+        multires = 10,
+        multires_view = 4,
+        output_ch = 4,
+        skips=[4],
+        use_viewdirs=True
+    }
+
+    sdf_network {
+        d_out = 257,
+        d_in = 3,
+        d_hidden = 256,
+        n_layers = 8,
+        skip_in = [4],
+        multires = 6,
+        bias = 0.5,
+        scale = 1.0,
+        geometric_init = True,
+        weight_norm = True,
+    } 
+
+    variance_network {
+        init_val = 0.3
+    }
+
+    rendering_network {
+        d_feature = 256,
+        mode = idr,
+        d_in = 9,
+        d_out = 3,
+        d_hidden = 256,
+        n_layers = 4,
+        weight_norm = True,
+        multires_view = 4,
+        squeeze_out = True,
+    }
+
+    neus_renderer {
+        n_samples = 64,
+        n_importance = 64,
+        n_outside = 0,
+        up_sample_steps = 4 ,
+        perturb = 1.0,
+    }
+
+    bending_network {
+        multires = 6,
+        bending_latent_size = 32,
+        d_in = 3,
+        rigidity_hidden_dimensions = 64,
+        rigidity_network_depth = 5,
+        use_rigidity_network = False,
+        bending_n_timesteps = 10,
+    }
+}

confs_new/dyn_grab_pointset_points_dyn_s4.conf

@@ -0,0 +1,259 @@
+general {
+
+
+    base_exp_dir = exp/CASE_NAME/wmask
+    
+    tag = "train_retargeted_shadow_hand_seq_102_mano_pointset_acts_optstates_optacts_optsysps_optacts_"
+
+    recording = [
+        ./,
+        ./models
+    ]
+}
+
+dataset {
+    data_dir = public_data/CASE_NAME/
+    render_cameras_name = cameras_sphere.npz
+    object_cameras_name = cameras_sphere.npz
+    obj_idx = 102
+}
+
+train {
+    learning_rate = 5e-4
+    learning_rate_alpha = 0.05
+    end_iter = 300000
+
+    batch_size = 1024 #  64
+    validate_resolution_level = 4
+    warm_up_end = 5000
+    anneal_end = 0
+    use_white_bkgd = False
+
+    # save_freq = 10000
+    save_freq = 10000
+    val_freq = 20 #  2500
+    val_mesh_freq = 20 #  5000
+    report_freq = 10
+    ### igr weight ###
+    igr_weight = 0.1
+    mask_weight = 0.1
+}
+ 
+model {
+
+    optimize_dyn_actions = True
+
+
+    optimize_robot = True
+
+    use_penalty_based_friction = True
+
+    use_split_params = False
+
+    use_sqr_spring_stiffness = True
+
+    use_pre_proj_frictions = True
+
+
+
+    use_sqrt_dist = True
+    contact_maintaining_dist_thres = 0.2
+
+    robot_actions_diff_coef = 0.001
+
+
+    use_sdf_as_contact_dist = True
+
+
+    # 
+    use_contact_dist_as_sdf = False
+
+    use_glb_proj_delta = True
+
+
+
+# penetration_proj_k_to_robot = 30
+    penetrating_depth_penalty = 1.0
+    train_states = True
+
+
+
+    minn_dist_threshold = 0.000
+    obj_mass = 30.0
+
+
+    use_LBFGS = True
+    use_LBFGS = False
+
+    use_mano_hand_for_test = False # use the dynamic mano model here #
+
+    extract_delta_mesh = False
+    freeze_weights = True
+    gt_act_xs_def = False
+    use_bending_network =  True
+    ### for ts = 3 ###
+    # use_delta_bending = False
+    ### for ts = 3 ###
+   
+
+
+
+    sim_model_path = "rsc/shadow_hand_description/shadowhand_new.urdf"
+    mano_sim_model_path = "rsc/mano/mano_mean_wcollision_scaled_scaled_0_9507_nroot.urdf"
+
+    obj_sdf_fn = "data/grab/102/102_obj.npy" 
+    kinematic_mano_gt_sv_fn = "data/grab/102/102_sv_dict.npy"
+    scaled_obj_mesh_fn = "data/grab/102/102_obj.obj"
+
+    bending_net_type = "active_force_field_v18"    
+    sim_num_steps = 1000000
+    n_timesteps = 60
+    optim_sim_model_params_from_mano = False
+    penetration_determining = "sdf_of_canon"
+    train_with_forces_to_active = False
+    loss_scale_coef = 1000.0
+    use_same_contact_spring_k = False
+    use_optimizable_params = True # 
+    train_residual_friction = True
+    mano_mult_const_after_cent = 1.0
+    optimize_glb_transformations = True
+    no_friction_constraint = False
+    optimize_active_object = True
+    loss_tangential_diff_coef = 0
+    optimize_with_intermediates = True
+    using_delta_glb_trans = False
+    train_multi_seqs = False
+    use_split_network = True
+    use_delta_bending = True
+
+
+
+
+
+
+
+
+
+    ###### threshold, ks settings 1, optimize acts ######
+    # drive_pointset = "actions"
+    # fix_obj = True
+    # optimize_rules = False
+    # train_pointset_acts_via_deltas = True
+    # load_optimized_init_actions = "/data/xueyi/NeuS/exp/hand_test_routine_2_light_color_wtime_active_passive/wmask_reverse_value_totviews_tag_train_dyn_mano_hand_seq_102_mouse_optdynactions_points_optrobo_offsetdriven_optrules_multk100_wfixobj_optdelta_radius0d4_/checkpoints/ckpt_002000.pth"
+    # load_optimized_init_actions = "/data/xueyi/NeuS/exp/hand_test_routine_2_light_color_wtime_active_passive/wmask_reverse_value_totviews_tag_train_dyn_mano_hand_seq_102_mouse_optdynactions_points_optrobo_offsetdriven_optrules_multk100_wfixobj_optdelta_radius0d2_/checkpoints/ckpt_008000.pth"
+    ###### threshold, ks settings 1, optimize acts ######
+
+
+    ##### contact spring model settings ####
+    minn_dist_threshold_robot_to_obj = 0.1
+    penetration_proj_k_to_robot_friction = 10000000.0
+    penetration_proj_k_to_robot = 4000000.0
+    ##### contact spring model settings ####
+ 
+
+    ###### Stage 1: threshold, ks settings 1, optimize offsets ######
+    drive_pointset = "states"
+    fix_obj = True
+    optimize_rules = False
+    train_pointset_acts_via_deltas = False
+    load_optimized_init_actions = "ckpts/grab/102/dyn_mano_arti.pth"
+    ###### Stage 1: threshold, ks settings 1, optimize offsets ######
+
+
+    ###### Stage 2: threshold, ks settings 1, optimize acts ######
+    drive_pointset = "actions"
+    fix_obj = True
+    optimize_rules = False
+    train_pointset_acts_via_deltas = True
+    load_optimized_init_actions = "ckpts/grab/102/dyn_mano_pointset_states.pt"
+    ###### Stage 2: threshold, ks settings 1, optimize acts ######
+
+
+    ###### Stage 3: threshold, ks settings 1, optimize params from acts ######
+    drive_pointset = "actions"
+    fix_obj = False
+    optimize_rules = True
+    train_pointset_acts_via_deltas = True
+    load_optimized_init_actions = "ckpts/grab/102/dyn_mano_pointset_acts.pt"
+    ##### model parameters optimized from the MANO hand trajectory #####
+    ckpt_fn = "ckpts/grab/102/dyn_mano_opts.pt"
+    ###### Stage 3: threshold, ks settings 1, optimize params from acts ######
+
+
+    ###### Stage 4: threshold, ks settings 1, optimize acts from optimized params ######
+    drive_pointset = "actions"
+    fix_obj = False
+    optimize_rules = False
+    train_pointset_acts_via_deltas = True ## pointset acts via deltas ###
+    ##### model parameters optimized from the MANO hand expanded set trajectory #####
+    ckpt_fn = "ckpts/grab/102/dyn_mano_pointset_optimized_acts_optimized_ps.pth"
+    load_optimized_init_actions = "ckpts/grab/102/dyn_mano_pointset_optimized_acts.pth"
+    ckpt_fn = "exp/hand_test_routine_2_light_color_wtime_active_passive/wmask_reverse_value_totviews_tag_train_retargeted_shadow_hand_seq_102_mano_pointset_acts_optstates_optacts_optsysps_/checkpoints/ckpt_044000.pth"
+    load_optimized_init_actions = "exp/hand_test_routine_2_light_color_wtime_active_passive/wmask_reverse_value_totviews_tag_train_retargeted_shadow_hand_seq_102_mano_pointset_acts_optstates_optacts_optsysps_/checkpoints/ckpt_044000.pth"
+    ###### Stage 4: threshold, ks settings 1, optimize acts from optimized params ######
+
+
+    use_opt_rigid_translations=True
+
+    train_def = True
+    optimizable_rigid_translations=True
+    
+    nerf {
+        D = 8,
+        d_in = 4,
+        d_in_view = 3,
+        W = 256,
+        multires = 10,
+        multires_view = 4,
+        output_ch = 4,
+        skips=[4],
+        use_viewdirs=True
+    }
+
+    sdf_network {
+        d_out = 257,
+        d_in = 3,
+        d_hidden = 256,
+        n_layers = 8,
+        skip_in = [4],
+        multires = 6,
+        bias = 0.5,
+        scale = 1.0,
+        geometric_init = True,
+        weight_norm = True,
+    } 
+
+    variance_network {
+        init_val = 0.3
+    }
+
+    rendering_network {
+        d_feature = 256,
+        mode = idr,
+        d_in = 9,
+        d_out = 3,
+        d_hidden = 256,
+        n_layers = 4,
+        weight_norm = True,
+        multires_view = 4,
+        squeeze_out = True,
+    }
+
+    neus_renderer {
+        n_samples = 64,
+        n_importance = 64,
+        n_outside = 0,
+        up_sample_steps = 4 ,
+        perturb = 1.0,
+    }
+
+    bending_network {
+        multires = 6,
+        bending_latent_size = 32,
+        d_in = 3,
+        rigidity_hidden_dimensions = 64,
+        rigidity_network_depth = 5,
+        use_rigidity_network = False,
+        bending_n_timesteps = 10,
+    }
+}

confs_new/dyn_grab_sparse_retar.conf

@@ -0,0 +1,214 @@
+general {
+
+
+    base_exp_dir = exp/CASE_NAME/wmask
+
+
+    tag = "train_retargeted_shadow_hand_seq_102_mano_sparse_retargeting_"
+
+    recording = [
+        ./,
+        ./models
+    ]
+}
+
+dataset {
+    data_dir = public_data/CASE_NAME/
+    render_cameras_name = cameras_sphere.npz
+    object_cameras_name = cameras_sphere.npz
+    obj_idx = 102
+}
+
+train {
+    learning_rate = 5e-4
+    learning_rate_alpha = 0.05
+    end_iter = 300000
+
+    batch_size = 1024
+    validate_resolution_level = 4
+    warm_up_end = 5000
+    anneal_end = 0
+    use_white_bkgd = False
+
+    # save_freq = 10000
+    save_freq = 10000
+    val_freq = 20
+    val_mesh_freq = 20
+    report_freq = 10
+    igr_weight = 0.1
+    mask_weight = 0.1
+}
+ 
+model {
+
+    optimize_dyn_actions = True
+
+
+    optimize_robot = True
+
+    use_penalty_based_friction = True
+
+    use_split_params = False
+
+    use_sqr_spring_stiffness = True
+
+    use_pre_proj_frictions = True
+
+
+
+    use_sqrt_dist = True
+    contact_maintaining_dist_thres = 0.2
+
+    robot_actions_diff_coef = 0.001
+
+
+    use_sdf_as_contact_dist = True
+
+
+    # 
+    use_contact_dist_as_sdf = False
+
+    use_glb_proj_delta = True
+
+
+
+# penetration_proj_k_to_robot = 30
+    penetrating_depth_penalty = 1.0
+    train_states = True
+
+
+
+    minn_dist_threshold = 0.000
+    obj_mass = 30.0
+
+
+    use_LBFGS = True
+    use_LBFGS = False
+
+    use_mano_hand_for_test = False # use the dynamic mano model here #
+
+    extract_delta_mesh = False
+    freeze_weights = True
+    gt_act_xs_def = False
+    use_bending_network =  True
+    ### for ts = 3 ###
+    # use_delta_bending = False
+    ### for ts = 3 ###
+   
+
+
+
+    sim_model_path = "rsc/shadow_hand_description/shadowhand_new.urdf"
+    mano_sim_model_path = "rsc/mano/mano_mean_wcollision_scaled_scaled_0_9507_nroot.urdf"
+
+    obj_sdf_fn = "data/grab/102/102_obj.npy" 
+    kinematic_mano_gt_sv_fn = "data/grab/102/102_sv_dict.npy"
+    scaled_obj_mesh_fn = "data/grab/102/102_obj.obj"
+
+    bending_net_type = "active_force_field_v18"    
+    sim_num_steps = 1000000
+    n_timesteps = 60
+    optim_sim_model_params_from_mano = False
+    penetration_determining = "sdf_of_canon"
+    train_with_forces_to_active = False
+    loss_scale_coef = 1000.0
+    use_same_contact_spring_k = False
+    use_optimizable_params = True # 
+    train_residual_friction = True
+    mano_mult_const_after_cent = 1.0
+    optimize_glb_transformations = True
+    no_friction_constraint = False
+    optimize_active_object = True
+    loss_tangential_diff_coef = 0
+    optimize_with_intermediates = True
+    using_delta_glb_trans = False
+    train_multi_seqs = False
+    use_split_network = True
+    use_delta_bending = True
+
+
+
+
+    ##### contact spring model settings ####
+    minn_dist_threshold_robot_to_obj = 0.1
+    penetration_proj_k_to_robot_friction = 10000000.0
+    penetration_proj_k_to_robot = 4000000.0
+    ##### contact spring model settings ####
+ 
+
+    ######  ######
+    drive_pointset = "states"
+    fix_obj = True
+    optimize_rules = False
+    train_pointset_acts_via_deltas = False
+    load_optimized_init_actions = ""
+    load_optimized_init_transformations = ""
+    ckpt_fn = ""
+    retar_only_glb = True
+    # use_multi_stages = True
+    ###### Stage 1: threshold, ks settings 1, optimize offsets ######
+
+    use_opt_rigid_translations=True
+
+    train_def = True
+    optimizable_rigid_translations=True
+    
+    nerf {
+        D = 8,
+        d_in = 4,
+        d_in_view = 3,
+        W = 256,
+        multires = 10,
+        multires_view = 4,
+        output_ch = 4,
+        skips=[4],
+        use_viewdirs=True
+    }
+
+    sdf_network {
+        d_out = 257,
+        d_in = 3,
+        d_hidden = 256,
+        n_layers = 8,
+        skip_in = [4],
+        multires = 6,
+        bias = 0.5,
+        scale = 1.0,
+        geometric_init = True,
+        weight_norm = True,
+    } 
+
+    variance_network {
+        init_val = 0.3
+    }
+
+    rendering_network {
+        d_feature = 256,
+        mode = idr,
+        d_in = 9,
+        d_out = 3,
+        d_hidden = 256,
+        n_layers = 4,
+        weight_norm = True,
+        multires_view = 4,
+        squeeze_out = True,
+    }
+
+    neus_renderer {
+        n_samples = 64,
+        n_importance = 64,
+        n_outside = 0,
+        up_sample_steps = 4 ,
+        perturb = 1.0,
+    }
+
+    bending_network {
+        multires = 6,
+        bending_latent_size = 32,
+        d_in = 3,
+        rigidity_hidden_dimensions = 64,
+        rigidity_network_depth = 5,
+        use_rigidity_network = False,
+        bending_n_timesteps = 10,
+    }
+}

models/data_utils_torch.py

@@ -0,0 +1,1547 @@
+"""Mesh data utilities."""
+from re import I
+# import matplotlib.pyplot as plt
+# from mpl_toolkits import mplot3d  # pylint: disable=unused-import
+# from mpl_toolkits.mplot3d.art3d import Poly3DCollection
+import networkx as nx
+import numpy as np
+import six
+import os
+import math
+import torch
+import models.fields as fields
+# from options.options import opt
+# from polygen_torch.
+# from utils.constants import MASK_GRID_VALIE
+
+try:
+    from torch_cluster import fps
+except:
+    pass
+
+MAX_RANGE = 0.1
+MIN_RANGE = -0.1
+
+# import open3d as o3d
+
+def calculate_correspondences(last_mesh, bending_network, tot_timesteps, delta_bending):
+  # iterate all the timesteps and get the bended 
+  timestep_to_pts = {
+    tot_timesteps - 1: last_mesh.detach().cpu().numpy()
+  }
+  if delta_bending:
+    
+    for i_ts in range(tot_timesteps - 1, -1, -1):
+      if isinstance(bending_network, list):
+        tot_offsets = []
+        for i_bending, cur_bending_net in enumerate(bending_network):
+          cur_def_pts = cur_bending_net(last_mesh, i_ts)
+          tot_offsets.append(cur_def_pts - last_mesh)
+        tot_offsets = torch.stack(tot_offsets, dim=0)
+        tot_offsets = torch.sum(tot_offsets, dim=0)
+        last_mesh = last_mesh + tot_offsets
+      else:
+        last_mesh = bending_network(last_mesh, i_ts)
+      timestep_to_pts[i_ts - 1] = last_mesh.detach().cpu().numpy()
+  elif isinstance(bending_network, fields.BendingNetworkRigidTrans):
+    for i_ts in range(tot_timesteps - 1, -1, -1):
+      last_mesh = bending_network.forward_delta(last_mesh, i_ts)
+      timestep_to_pts[i_ts - 1] = last_mesh.detach().cpu().numpy()
+  else:
+    raise ValueError(f"the function is designed for delta bending")
+  return timestep_to_pts
+  # pass
+
+def joint_infos_to_numpy(joint_infos):
+  joint_infos_np = []
+  for part_joint_info in joint_infos:
+    for zz in ["dir", "center"]:
+      # if isinstance(part_joint_info["axis"][zz], np.array):
+      part_joint_info["axis"][zz] = part_joint_info["axis"][zz].detach().numpy()
+    joint_infos_np.append(part_joint_info)
+  return joint_infos_np
+
+
+def normalie_pc_bbox_batched(pc, rt_stats=False):
+  pc_min = torch.min(pc, dim=1, keepdim=True)[0]
+  pc_max = torch.max(pc, dim=1, keepdim=True)[0]
+  pc_center = 0.5 * (pc_min + pc_max)
+  
+  pc = pc - pc_center
+  extents = pc_max - pc_min
+  scale = torch.sqrt(torch.sum(extents ** 2, dim=-1, keepdim=True))
+  
+  pc = pc / torch.clamp(scale, min=1e-6)
+  if rt_stats:
+    return pc, pc_center, scale
+  else:
+    return pc
+
+
+
+def scale_vertices_to_target_scale(vertices, target_scale):
+  # vertices: bsz x N x 3;
+  # target_scale: bsz x 1
+#   normalized_vertices = normalize_vertices_scale_torch(vertices)
+  normalized_vertices = normalie_pc_bbox_batched(vertices)
+  normalized_vertices = normalized_vertices * target_scale.unsqueeze(1)
+  return normalized_vertices
+
+def compute_normals_o3d(verts, faces): #### assume no batching... ####
+  mesh = o3d.geometry.TriangleMesh()
+  # o3d_mesh_b.vertices = verts_b
+        # o3d_mesh_b.triangles = np.array(faces_b, dtype=np.long)
+  mesh.vertices = verts.detach().cpu().numpy()
+  mesh.triangles = faces.detach().cpu().numpy().astype(np.long)
+  verts_normals = mesh.compute_vertex_normals(normalized=True)
+  verts_normals = torch.from_numpy(verts_normals, dtype=torch.float32).cuda()
+  return verts_normals
+  
+  
+def get_vals_via_nearest_neighbours(pts_src, pts_tar, val_tar):
+  ### n_src x 3 ---> n_src x n_tar
+  dist_src_tar = torch.sum((pts_src.unsqueeze(-2) - pts_tar.unsqueeze(0)) ** 2, dim=-1)
+  minn_dists_src_tar, minn_dists_tar_idxes = torch.min(dist_src_tar, dim=-1) ### n_src
+  selected_src_val = batched_index_select(values=val_tar, indices=minn_dists_tar_idxes, dim=0) ### n_src x dim
+  return selected_src_val
+  
+  
+
+## sample conected componetn start from selected_verts
+def sample_bfs_component(selected_vert, faces, max_num_grids):
+  vert_idx_to_adj_verts = {}
+  for i_f, cur_f in enumerate(faces):
+    # for i0, v0 in enumerate(cur_f):
+    for i0 in range(len(cur_f)):
+      v0 = cur_f[i0] - 1
+      i1 = (i0 + 1) % len(cur_f)
+      v1 = cur_f[i1] - 1
+      if v0 not in vert_idx_to_adj_verts:
+        vert_idx_to_adj_verts[v0] = {v1: 1}
+      else:
+        vert_idx_to_adj_verts[v0][v1] = 1
+      if v1 not in vert_idx_to_adj_verts:
+        vert_idx_to_adj_verts[v1] = {v0: 1}
+      else:
+        vert_idx_to_adj_verts[v1][v0] = 1
+  vert_idx_to_visited = {} # whether visisted here # 
+  vis_que = [selected_vert]
+  vert_idx_to_visited[selected_vert] = 1
+  visited = [selected_vert]
+  while len(vis_que) > 0 and len(visited) < max_num_grids:
+    cur_frnt_vert = vis_que[0]
+    vis_que.pop(0)
+    if cur_frnt_vert in vert_idx_to_adj_verts:
+      cur_frnt_vert_adjs = vert_idx_to_adj_verts[cur_frnt_vert]
+      for adj_vert in cur_frnt_vert_adjs:
+        if adj_vert not in vert_idx_to_visited:
+          vert_idx_to_visited[adj_vert] = 1
+          vis_que.append(adj_vert)
+          visited.append(adj_vert)
+  if len(visited) >= max_num_grids:
+    visited = visited[: max_num_grids - 1]
+  return visited
+
+def select_faces_via_verts(selected_verts, faces):
+  if not isinstance(selected_verts, list):
+    selected_verts = selected_verts.tolist()
+  # selected_verts_dict = {ii: 1 for ii in selected_verts}
+  old_idx_to_new_idx = {v + 1: ii + 1 for ii, v in enumerate(selected_verts)} ####### v + 1: ii + 1 --> for the selected_verts
+  new_faces = []
+  for i_f, cur_f in enumerate(faces):
+    cur_new_f = []
+    valid = True
+    for cur_v in cur_f:
+      if cur_v in old_idx_to_new_idx:
+        cur_new_f.append(old_idx_to_new_idx[cur_v])
+      else:
+        valid  = False
+        break
+    if valid:
+      new_faces.append(cur_new_f)
+  return new_faces 
+      
+
+def convert_grid_content_to_grid_pts(content_value, grid_size):
+  flat_grid = torch.zeros([grid_size ** 3], dtype=torch.long)
+  cur_idx = flat_grid.size(0) - 1
+  while content_value > 0:
+    flat_grid[cur_idx] = content_value % grid_size
+    content_value = content_value // grid_size
+    cur_idx -= 1
+  grid_pts = flat_grid.contiguous().view(grid_size, grid_size, grid_size).contiguous()
+  return grid_pts
+
+# 0.2
+def warp_coord(sampled_gradients, val, reflect=False): # val from [0.0, 1.0] # from the 0.0 
+  # assume single value as inputs
+  grad_values = sampled_gradients.tolist()
+  # mid_val
+  mid_val = grad_values[0] * 0.2 + grad_values[1] * 0.2 + grad_values[2] * 0.1
+  if reflect:
+    grad_values[3] = grad_values[1]
+    grad_values[4] = grad_values[0]
+
+  # if not reflect:
+  accum_val = 0.0
+  for i_val in range(len(grad_values)):
+    if val > 0.2 * (i_val + 1) and i_val < 4: # if i_val == 4, directly use the reamining length * corresponding gradient value
+      accum_val += grad_values[i_val] * 0.2
+    else:
+      accum_val += grad_values[i_val] * (val - 0.2 * i_val)
+      break
+  return accum_val # modified value
+
+def random_shift(vertices, shift_factor=0.25):
+  """Apply random shift to vertices."""
+  # max_shift_pos = tf.cast(255 - tf.reduce_max(vertices, axis=0), tf.float32)
+  
+  # max_shift_pos = tf.maximum(max_shift_pos, 1e-9)
+
+  # max_shift_neg = tf.cast(tf.reduce_min(vertices, axis=0), tf.float32)
+  # max_shift_neg = tf.maximum(max_shift_neg, 1e-9)
+
+  # shift = tfd.TruncatedNormal(
+  #     tf.zeros([1, 3]), shift_factor*255, -max_shift_neg,
+  #     max_shift_pos).sample()
+  # shift = tf.cast(shift, tf.int32)
+  # vertices += shift
+
+  minn_tensor = torch.tensor([1e-9], dtype=torch.float32)
+  
+  max_shift_pos = (255 - torch.max(vertices, dim=0)[0]).float()
+  max_shift_pos = torch.maximum(max_shift_pos, minn_tensor)
+  max_shift_neg = (torch.min(vertices, dim=0)[0]).float()
+  max_shift_neg = torch.maximum(max_shift_neg, minn_tensor)
+  
+  shift = torch.zeros((1, 3), dtype=torch.float32)
+  # torch.nn.init.trunc_normal_(shift, 0., shift_factor * 255., -max_shift_neg, max_shift_pos)
+  for i_s in range(shift.size(-1)):
+    cur_axis_max_shift_neg = max_shift_neg[i_s].item()
+    cur_axis_max_shift_pos = max_shift_pos[i_s].item()
+    cur_axis_shift = torch.zeros((1,), dtype=torch.float32)
+
+    torch.nn.init.trunc_normal_(cur_axis_shift, 0., shift_factor * 255., -cur_axis_max_shift_neg, cur_axis_max_shift_pos)
+    shift[:, i_s] = cur_axis_shift.item()
+    
+  shift = shift.long()
+  vertices += shift
+
+  return vertices
+
+def safe_transpose(tsr, dima, dimb):
+  tsr = tsr.contiguous().transpose(dima, dimb).contiguous()
+  return tsr
+
+def batched_index_select(values, indices, dim = 1):
+  value_dims = values.shape[(dim + 1):]
+  values_shape, indices_shape = map(lambda t: list(t.shape), (values, indices))
+  indices = indices[(..., *((None,) * len(value_dims)))]
+  indices = indices.expand(*((-1,) * len(indices_shape)), *value_dims)
+  value_expand_len = len(indices_shape) - (dim + 1)
+  values = values[(*((slice(None),) * dim), *((None,) * value_expand_len), ...)]
+
+  value_expand_shape = [-1] * len(values.shape)
+  expand_slice = slice(dim, (dim + value_expand_len))
+  value_expand_shape[expand_slice] = indices.shape[expand_slice]
+  values = values.expand(*value_expand_shape)
+
+  dim += value_expand_len
+  return values.gather(dim, indices)
+
+def read_obj_file_ours(obj_fn, sub_one=False):
+  vertices = []
+  faces = []
+  with open(obj_fn, "r") as rf:
+    for line in rf:
+      items = line.strip().split(" ")
+      if items[0] == 'v':
+        cur_verts = items[1:]
+        cur_verts = [float(vv) for vv in cur_verts]
+        vertices.append(cur_verts)
+      elif items[0] == 'f':
+        cur_faces = items[1:] # faces
+        cur_face_idxes = []
+        for cur_f in cur_faces:
+          try:
+            cur_f_idx = int(cur_f.split("/")[0])
+          except:
+            cur_f_idx = int(cur_f.split("//")[0])
+          cur_face_idxes.append(cur_f_idx if not sub_one else cur_f_idx - 1)
+        faces.append(cur_face_idxes)
+    rf.close()
+  vertices = np.array(vertices, dtype=np.float)
+  return vertices, faces
+
+def read_obj_file(obj_file):
+  """Read vertices and faces from already opened file."""
+  vertex_list = []
+  flat_vertices_list = []
+  flat_vertices_indices = {}
+  flat_triangles = []
+
+  for line in obj_file:
+    tokens = line.split()
+    if not tokens:
+      continue
+    line_type = tokens[0]
+    # We skip lines not starting with v or f.
+    if line_type == 'v': # 
+      vertex_list.append([float(x) for x in tokens[1:]])
+    elif line_type == 'f':
+      triangle = []
+      for i in range(len(tokens) - 1):
+        vertex_name = tokens[i + 1]
+        if vertex_name in flat_vertices_indices: # triangles
+          triangle.append(flat_vertices_indices[vertex_name])
+          continue
+        flat_vertex = []
+        for index in six.ensure_str(vertex_name).split('/'):
+          if not index:
+            continue
+          # obj triangle indices are 1 indexed, so subtract 1 here.
+          flat_vertex += vertex_list[int(index) - 1]
+        flat_vertex_index = len(flat_vertices_list)
+        flat_vertices_list.append(flat_vertex)
+        # flat_vertex_index
+        flat_vertices_indices[vertex_name] = flat_vertex_index
+        triangle.append(flat_vertex_index)
+      flat_triangles.append(triangle)
+
+  return np.array(flat_vertices_list, dtype=np.float32), flat_triangles
+
+
+def batched_index_select(values, indices, dim = 1):
+    value_dims = values.shape[(dim + 1):]
+    values_shape, indices_shape = map(lambda t: list(t.shape), (values, indices))
+    indices = indices[(..., *((None,) * len(value_dims)))]
+    indices = indices.expand(*((-1,) * len(indices_shape)), *value_dims)
+    value_expand_len = len(indices_shape) - (dim + 1)
+    values = values[(*((slice(None),) * dim), *((None,) * value_expand_len), ...)]
+
+    value_expand_shape = [-1] * len(values.shape)
+    expand_slice = slice(dim, (dim + value_expand_len))
+    value_expand_shape[expand_slice] = indices.shape[expand_slice]
+    values = values.expand(*value_expand_shape)
+
+    dim += value_expand_len
+    return values.gather(dim, indices)
+
+
+def safe_transpose(x, dim1, dim2):
+    x = x.contiguous().transpose(dim1, dim2).contiguous()
+    return x
+
+def merge_meshes(vertices_list, faces_list):
+  tot_vertices = []
+  tot_faces = []
+  nn_verts = 0
+  for cur_vertices, cur_faces in zip(vertices_list, faces_list):
+    tot_vertices.append(cur_vertices)
+    new_cur_faces = []
+    for cur_face_idx in cur_faces:
+      new_cur_face_idx = [vert_idx + nn_verts for vert_idx in cur_face_idx]
+      new_cur_faces.append(new_cur_face_idx)
+    nn_verts += cur_vertices.shape[0]
+    tot_faces += new_cur_faces # get total-faces
+  tot_vertices = np.concatenate(tot_vertices, axis=0)
+  return tot_vertices, tot_faces
+
+
+def read_obj(obj_path):
+  """Open .obj file from the path provided and read vertices and faces."""
+
+  with open(obj_path) as obj_file:
+    return read_obj_file_ours(obj_path, sub_one=True)
+    # return read_obj_file(obj_file)
+
+
+
+
+def center_vertices(vertices):
+  """Translate the vertices so that bounding box is centered at zero."""
+  vert_min = vertices.min(axis=0)
+  vert_max = vertices.max(axis=0)
+  vert_center = 0.5 * (vert_min + vert_max)
+  return vertices - vert_center
+
+
+def normalize_vertices_scale(vertices):
+  """Scale the vertices so that the long diagonal of the bounding box is one."""
+  vert_min = vertices.min(axis=0)
+  vert_max = vertices.max(axis=0)
+  extents = vert_max - vert_min
+  scale = np.sqrt(np.sum(extents**2)) # normalize the diagonal line to 1.
+  return vertices / scale
+
+
+def get_vertices_center(vertices):
+  vert_min = vertices.min(axis=0)
+  vert_max = vertices.max(axis=0)
+  vert_center = 0.5 * (vert_min + vert_max)
+  return vert_center
+
+def get_batched_vertices_center(vertices):
+  vert_min = vertices.min(axis=1)
+  vert_max = vertices.max(axis=1)
+  vert_center = 0.5 * (vert_min + vert_max)
+  return vert_center
+
+def get_vertices_scale(vertices):
+  vert_min = vertices.min(axis=0)
+  vert_max = vertices.max(axis=0)
+  extents = vert_max - vert_min
+  scale = np.sqrt(np.sum(extents**2))
+  return scale
+
+def quantize_verts(verts, n_bits=8, min_range=None, max_range=None):
+  """Convert vertices in [-1., 1.] to discrete values in [0, n_bits**2 - 1]."""
+  min_range = -0.5 if min_range is None else min_range
+  max_range = 0.5 if max_range is None else max_range
+  range_quantize = 2**n_bits - 1
+  verts_quantize = (verts - min_range) * range_quantize / (
+      max_range - min_range)
+  return verts_quantize.astype('int32')
+
+def quantize_verts_torch(verts, n_bits=8, min_range=None, max_range=None):
+  min_range = -0.5 if min_range is None else min_range
+  max_range = 0.5 if max_range is None else max_range
+  range_quantize = 2**n_bits - 1
+  verts_quantize = (verts - min_range) * range_quantize / (
+      max_range - min_range)
+  return verts_quantize.long()
+
+def dequantize_verts(verts, n_bits=8, add_noise=False, min_range=None, max_range=None):
+  """Convert quantized vertices to floats."""
+  min_range = -0.5 if min_range is None else min_range
+  max_range = 0.5 if max_range is None else max_range
+  range_quantize = 2**n_bits - 1
+  verts = verts.astype('float32')
+  verts = verts * (max_range - min_range) / range_quantize + min_range
+  if add_noise:
+    verts += np.random.uniform(size=verts.shape) * (1 / range_quantize)
+  return verts
+
+def dequantize_verts_torch(verts, n_bits=8, add_noise=False, min_range=None, max_range=None):
+  min_range = -0.5 if min_range is None else min_range
+  max_range = 0.5 if max_range is None else max_range
+  range_quantize = 2**n_bits - 1
+  verts = verts.float()
+  verts = verts * (max_range - min_range) / range_quantize + min_range
+  # if add_noise:
+  #   verts += np.random.uniform(size=verts.shape) * (1 / range_quantize)
+  return verts
+
+
+### dump vertices and faces to the obj file
+def write_obj(vertices, faces, file_path, transpose=True, scale=1.):
+  """Write vertices and faces to obj."""
+  if transpose:
+    vertices = vertices[:, [1, 2, 0]]
+  vertices *= scale
+  if faces is not None:
+    if min(min(faces)) == 0:
+      f_add = 1
+    else:
+      f_add = 0
+  else:
+    faces = []
+  with open(file_path, 'w') as f:
+    for v in vertices:
+      f.write('v {} {} {}\n'.format(v[0], v[1], v[2]))
+    for face in faces:
+      line = 'f'
+      for i in face:
+        line += ' {}'.format(i + f_add)
+      line += '\n'
+      f.write(line)
+
+
+def face_to_cycles(face):
+  """Find cycles in face."""
+  g = nx.Graph()
+  for v in range(len(face) - 1):
+    g.add_edge(face[v], face[v + 1])
+  g.add_edge(face[-1], face[0])
+  return list(nx.cycle_basis(g))
+
+
+def flatten_faces(faces):
+  """Converts from list of faces to flat face array with stopping indices."""
+  if not faces:
+    return np.array([0])
+  else:
+    l = [f + [-1] for f in faces[:-1]]
+    l += [faces[-1] + [-2]]
+    return np.array([item for sublist in l for item in sublist]) + 2  # pylint: disable=g-complex-comprehension
+
+
+def unflatten_faces(flat_faces):
+  """Converts from flat face sequence to a list of separate faces."""
+  def group(seq):
+    g = []
+    for el in seq:
+      if el == 0 or el == -1:
+        yield g
+        g = []
+      else:
+        g.append(el - 1)
+    yield g
+  outputs = list(group(flat_faces - 1))[:-1]
+  # Remove empty faces
+  return [o for o in outputs if len(o) > 2]
+
+
+
+def quantize_process_mesh(vertices, faces, tris=None, quantization_bits=8, remove_du=True):
+  """Quantize vertices, remove resulting duplicates and reindex faces."""
+  vertices = quantize_verts(vertices, quantization_bits)
+  vertices, inv = np.unique(vertices, axis=0, return_inverse=True) # return inverse and unique the vertices
+
+  # 
+  if opt.dataset.sort_dist:
+    if opt.model.debug:
+      print("sorting via dist...")
+    vertices_max = np.max(vertices, axis=0)
+    vertices_min = np.min(vertices, axis=0)
+    dist_vertices = np.minimum(np.abs(vertices - np.array([[vertices_min[0], vertices_min[1], 0]])), np.abs(vertices - np.array([[vertices_max[0], vertices_max[1], 0]])))
+    dist_vertices = np.concatenate([dist_vertices[:, 0:1] + dist_vertices[:, 1:2], dist_vertices[:, 2:]], axis=-1)
+    sort_inds = np.lexsort(dist_vertices.T)
+  else:
+    # Sort vertices by z then y then x.
+    sort_inds = np.lexsort(vertices.T) # sorted indices...
+  vertices = vertices[sort_inds]
+
+  # Re-index faces and tris to re-ordered vertices.
+  faces = [np.argsort(sort_inds)[inv[f]] for f in faces]
+  if tris is not None:
+    tris = np.array([np.argsort(sort_inds)[inv[t]] for t in tris])
+
+  # Merging duplicate vertices and re-indexing the faces causes some faces to
+  # contain loops (e.g [2, 3, 5, 2, 4]). Split these faces into distinct
+  # sub-faces.
+  sub_faces = []
+  for f in faces:
+    cliques = face_to_cycles(f)
+    for c in cliques:
+      c_length = len(c)
+      # Only append faces with more than two verts.
+      if c_length > 2:
+        d = np.argmin(c)
+        # Cyclically permute faces just that first index is the smallest.
+        sub_faces.append([c[(d + i) % c_length] for i in range(c_length)])
+  faces = sub_faces
+  if tris is not None:
+    tris = np.array([v for v in tris if len(set(v)) == len(v)])
+
+  # Sort faces by lowest vertex indices. If two faces have the same lowest
+  # index then sort by next lowest and so on.
+  faces.sort(key=lambda f: tuple(sorted(f)))
+  if tris is not None:
+    tris = tris.tolist()
+    tris.sort(key=lambda f: tuple(sorted(f)))
+    tris = np.array(tris)
+
+  # After removing degenerate faces some vertices are now unreferenced. # Vertices
+  # Remove these. # Vertices
+  num_verts = vertices.shape[0]
+  # print(f"remove_du: {remove_du}")
+  if remove_du: ##### num_verts
+    print("Removing du..")
+    try:
+      vert_connected = np.equal(
+          np.arange(num_verts)[:, None], np.hstack(faces)[None]).any(axis=-1)
+      vertices = vertices[vert_connected]
+    
+
+      # Re-index faces and tris to re-ordered vertices.
+      vert_indices = (
+          np.arange(num_verts) - np.cumsum(1 - vert_connected.astype('int')))
+      faces = [vert_indices[f].tolist() for f in faces]
+    except:
+      pass
+  if tris is not None:
+    tris = np.array([vert_indices[t].tolist() for t in tris])
+
+  return vertices, faces, tris
+
+
+def process_mesh(vertices, faces, quantization_bits=8, recenter_mesh=True, remove_du=True):
+  """Process mesh vertices and faces."""
+
+  
+
+  # Transpose so that z-axis is vertical.
+  vertices = vertices[:, [2, 0, 1]]
+
+  if recenter_mesh:
+    # Translate the vertices so that bounding box is centered at zero.
+    vertices = center_vertices(vertices)
+
+    # Scale the vertices so that the long diagonal of the bounding box is equal
+    # to one.
+    vertices = normalize_vertices_scale(vertices)
+
+  # Quantize and sort vertices, remove resulting duplicates, sort and reindex
+  # faces.
+  vertices, faces, _ = quantize_process_mesh(
+      vertices, faces, quantization_bits=quantization_bits, remove_du=remove_du) ##### quantize_process_mesh
+  
+  # unflatten_faces = np.array(faces, dtype=np.long) ### start from zero
+
+  # Flatten faces and add 'new face' = 1 and 'stop' = 0 tokens.
+  faces = flatten_faces(faces)
+
+  # Discard degenerate meshes without faces.
+  return {
+      'vertices': vertices,
+      'faces': faces,
+  }
+
+
+def process_mesh_ours(vertices, faces, quantization_bits=8, recenter_mesh=True, remove_du=True):
+  """Process mesh vertices and faces."""
+  # Transpose so that z-axis is vertical.
+  vertices = vertices[:, [2, 0, 1]]
+
+  if recenter_mesh:
+    # Translate the vertices so that bounding box is centered at zero.
+    vertices = center_vertices(vertices)
+
+    # Scale the vertices so that the long diagonal of the bounding box is equal
+    # to one.
+    vertices = normalize_vertices_scale(vertices)
+
+  # Quantize and sort vertices, remove resulting duplicates, sort and reindex
+  # faces.
+  quant_vertices, faces, _ = quantize_process_mesh(
+      vertices, faces, quantization_bits=quantization_bits, remove_du=remove_du) ##### quantize_process_mesh
+  vertices = dequantize_verts(quant_vertices) #### dequantize vertices ####
+  ### vertices: nn_verts x 3
+  # try:
+  #   # print("faces", faces)
+  #   unflatten_faces = np.array(faces, dtype=np.long)
+  # except:
+  #   print("faces", faces)
+  #   raise ValueError("Something bad happened when processing meshes...")
+  
+  # Flatten faces and add 'new face' = 1 and 'stop' = 0 tokens.
+
+  faces = flatten_faces(faces)
+
+  # Discard degenerate meshes without faces.
+  return {
+      'vertices': quant_vertices,
+      'faces': faces,
+      # 'unflatten_faces': unflatten_faces,
+      'dequant_vertices': vertices,
+      'class_label': 0
+  }
+
+def read_mesh_from_obj_file(fn, quantization_bits=8, recenter_mesh=True, remove_du=True):
+  vertices, faces = read_obj(fn)
+  # print(vertices.shape)
+  mesh_dict = process_mesh_ours(vertices, faces, quantization_bits=quantization_bits, recenter_mesh=recenter_mesh, remove_du=remove_du)
+  return mesh_dict
+
+def process_mesh_list(vertices, faces, quantization_bits=8, recenter_mesh=True):
+  """Process mesh vertices and faces."""
+
+  vertices = [cur_vert[:, [2, 0, 1]] for cur_vert in vertices]
+
+  tot_vertices = np.concatenate(vertices, axis=0) # center and scale of those vertices
+  vert_center = get_vertices_center(tot_vertices)
+  vert_scale = get_vertices_scale(tot_vertices)
+
+  processed_vertices, processed_faces = [], []
+
+  for cur_verts, cur_faces in zip(vertices, faces):
+    # print(f"current vertices: {cur_verts.shape}, faces: {len(cur_faces)}")
+    normalized_verts = (cur_verts - vert_center) / vert_scale
+    cur_processed_verts, cur_processed_faces, _ = quantize_process_mesh(
+      normalized_verts, cur_faces, quantization_bits=quantization_bits
+    )
+    processed_vertices.append(cur_processed_verts)
+    processed_faces.append(cur_processed_faces)
+  vertices, faces = merge_meshes(processed_vertices, processed_faces)
+  faces = flatten_faces(faces=faces)
+
+
+  # Discard degenerate meshes without faces.
+  return {
+      'vertices': vertices,
+      
+      'faces': faces,
+      
+  }
+
+
+def plot_sampled_meshes(v_sample, f_sample, sv_mesh_folder, cur_step=0, predict_joint=True,):
+  
+  if not os.path.exists(sv_mesh_folder):
+    os.mkdir(sv_mesh_folder)
+
+  part_vertex_samples = [v_sample['left'], v_sample['rgt']]
+  part_face_samples = [f_sample['left'], f_sample['rgt']]
+
+
+  tot_n_samples = part_vertex_samples[0]['vertices'].shape[0]
+  tot_n_part = 2
+
+  if predict_joint:
+    pred_dir = v_sample['joint_dir']
+    pred_pvp = v_sample['joint_pvp']
+    print("pred_dir", pred_dir.shape, pred_dir)
+    print("pred_pvp", pred_pvp.shape, pred_pvp)
+  else:
+    pred_pvp = np.zeros(shape=[tot_n_samples, 3], dtype=np.float32)
+  
+  
+
+
+  tot_mesh_list = []
+  for i_p, (cur_part_v_samples_np, cur_part_f_samples_np) in enumerate(zip(part_vertex_samples, part_face_samples)):
+      mesh_list = []
+      for i_n in range(tot_n_samples):
+          mesh_list.append(
+              {
+                  'vertices': cur_part_v_samples_np['vertices'][i_n][:cur_part_v_samples_np['num_vertices'][i_n]],
+                  'faces': unflatten_faces(
+                      cur_part_f_samples_np['faces'][i_n][:cur_part_f_samples_np['num_face_indices'][i_n]])
+              }
+          )
+      tot_mesh_list.append(mesh_list)
+      # and write this obj file?
+      # write_obj(vertices, faces, file_path, transpose=True, scale=1.):
+      # write mesh objs
+      for i_n in range(tot_n_samples):
+        cur_mesh = mesh_list[i_n]
+        cur_mesh_vertices, cur_mesh_faces = cur_mesh['vertices'], cur_mesh['faces']
+        cur_mesh_sv_fn = os.path.join("./meshes", f"training_step_{cur_step}_part_{i_p}_ins_{i_n}.obj")
+        if cur_mesh_vertices.shape[0] > 0 and len(cur_mesh_faces) > 0:
+          write_obj(cur_mesh_vertices, cur_mesh_faces, cur_mesh_sv_fn, transpose=False, scale=1.)
+        
+      
+
+  ###### plot mesh (predicted) ######
+  tot_samples_mesh_dict = []
+  for i_s in range(tot_n_samples):
+      cur_s_tot_vertices = []
+      cur_s_tot_faces = []
+      cur_s_n_vertices = 0
+      
+      for i_p in range(tot_n_part):
+          cur_s_cur_part_mesh_dict = tot_mesh_list[i_p][i_s]
+          cur_s_cur_part_vertices, cur_s_cur_part_faces = cur_s_cur_part_mesh_dict['vertices'], \
+                                                          cur_s_cur_part_mesh_dict['faces']
+          cur_s_cur_part_new_faces = []
+          for cur_s_cur_part_cur_face in cur_s_cur_part_faces:
+              cur_s_cur_part_cur_new_face = [fid + cur_s_n_vertices for fid in cur_s_cur_part_cur_face]
+              cur_s_cur_part_new_faces.append(cur_s_cur_part_cur_new_face)
+          cur_s_n_vertices += cur_s_cur_part_vertices.shape[0]
+          cur_s_tot_vertices.append(cur_s_cur_part_vertices)
+          cur_s_tot_faces += cur_s_cur_part_new_faces
+
+      cur_s_tot_vertices = np.concatenate(cur_s_tot_vertices, axis=0)
+      cur_s_mesh_dict = {
+          'vertices': cur_s_tot_vertices, 'faces': cur_s_tot_faces
+      }
+      tot_samples_mesh_dict.append(cur_s_mesh_dict)
+
+  for i_s in range(tot_n_samples):
+    cur_mesh = tot_samples_mesh_dict[i_s]
+    cur_mesh_vertices, cur_mesh_faces = cur_mesh['vertices'], cur_mesh['faces']
+    cur_mesh_sv_fn = os.path.join(sv_mesh_folder, f"training_step_{cur_step}_ins_{i_s}.obj")
+    if cur_mesh_vertices.shape[0] > 0 and len(cur_mesh_faces) > 0:
+      write_obj(cur_mesh_vertices, cur_mesh_faces, cur_mesh_sv_fn, transpose=False, scale=1.)
+  ###### plot mesh (predicted) ######
+
+
+  ###### plot mesh (translated) ######
+  tot_samples_mesh_dict = []
+  for i_s in range(tot_n_samples):
+      cur_s_tot_vertices = []
+      cur_s_tot_faces = []
+      cur_s_n_vertices = 0
+      cur_s_pred_pvp = pred_pvp[i_s]
+      
+      for i_p in range(tot_n_part):
+          cur_s_cur_part_mesh_dict = tot_mesh_list[i_p][i_s]
+          cur_s_cur_part_vertices, cur_s_cur_part_faces = cur_s_cur_part_mesh_dict['vertices'], \
+                                                          cur_s_cur_part_mesh_dict['faces']
+          cur_s_cur_part_new_faces = []
+          for cur_s_cur_part_cur_face in cur_s_cur_part_faces:
+              cur_s_cur_part_cur_new_face = [fid + cur_s_n_vertices for fid in cur_s_cur_part_cur_face]
+              cur_s_cur_part_new_faces.append(cur_s_cur_part_cur_new_face)
+          cur_s_n_vertices += cur_s_cur_part_vertices.shape[0]
+
+          if i_p == 1:
+            # min_rngs = cur_s_cur_part_vertices.min(1)
+            # max_rngs = cur_s_cur_part_vertices.max(1)
+            min_rngs = cur_s_cur_part_vertices.min(0)
+            max_rngs = cur_s_cur_part_vertices.max(0)
+            # shifted; cur_s_pred_pvp
+            # shifted = np.array([0., cur_s_pred_pvp[1] - max_rngs[1], cur_s_pred_pvp[2] - min_rngs[2]], dtype=np.float)
+            # shifted = np.reshape(shifted, [1, 3]) # 
+            cur_s_pred_pvp = np.array([0., max_rngs[1], min_rngs[2]], dtype=np.float32)
+            pvp_sample_pred_err = np.sum((cur_s_pred_pvp - pred_pvp[i_s]) ** 2)
+            # print prediction err, pred pvp and real pvp
+            # print("cur_s, sample_pred_pvp_err:", pvp_sample_pred_err.item(), ";real val:", cur_s_pred_pvp, "; pred_val:", pred_pvp[i_s])
+            pred_pvp[i_s] = cur_s_pred_pvp
+            shifted = np.zeros((1, 3), dtype=np.float32)
+            cur_s_cur_part_vertices = cur_s_cur_part_vertices + shifted # shift vertices... # min_rngs
+          # shifted
+          cur_s_tot_vertices.append(cur_s_cur_part_vertices)
+          cur_s_tot_faces += cur_s_cur_part_new_faces
+
+      cur_s_tot_vertices = np.concatenate(cur_s_tot_vertices, axis=0)
+      cur_s_mesh_dict = {
+          'vertices': cur_s_tot_vertices, 'faces': cur_s_tot_faces
+      }
+      tot_samples_mesh_dict.append(cur_s_mesh_dict)
+
+  for i_s in range(tot_n_samples):
+    cur_mesh = tot_samples_mesh_dict[i_s]
+    cur_mesh_vertices, cur_mesh_faces = cur_mesh['vertices'], cur_mesh['faces']
+    cur_mesh_sv_fn = os.path.join(sv_mesh_folder, f"training_step_{cur_step}_ins_{i_s}_shifted.obj")
+    if cur_mesh_vertices.shape[0] > 0 and len(cur_mesh_faces) > 0:
+      write_obj(cur_mesh_vertices, cur_mesh_faces, cur_mesh_sv_fn, transpose=False, scale=1.)
+  ###### plot mesh (translated) ######
+
+
+
+  ###### plot mesh (rotated) ######
+  if predict_joint:
+    from revolute_transform import revoluteTransform
+    tot_samples_mesh_dict = []
+    for i_s in range(tot_n_samples):
+        cur_s_tot_vertices = []
+        cur_s_tot_faces = []
+        cur_s_n_vertices = 0
+      
+        # cur_s_pred_dir = pred_dir[i_s]
+        cur_s_pred_pvp = pred_pvp[i_s]
+        print("current pred dir:", cur_s_pred_dir, "; current pred pvp:", cur_s_pred_pvp)
+        cur_s_pred_dir = np.array([1.0, 0.0, 0.0], dtype=np.float)
+        # cur_s_pred_pvp = cur_s_pred_pvp[[1, 2, 0]]
+
+        for i_p in range(tot_n_part):
+            cur_s_cur_part_mesh_dict = tot_mesh_list[i_p][i_s]
+            cur_s_cur_part_vertices, cur_s_cur_part_faces = cur_s_cur_part_mesh_dict['vertices'], \
+                                                            cur_s_cur_part_mesh_dict['faces']
+                      
+            if i_p == 1:
+              cur_s_cur_part_vertices, _ = revoluteTransform(cur_s_cur_part_vertices, cur_s_pred_pvp, cur_s_pred_dir, 0.5 * np.pi) # reverse revolute vertices of the upper piece
+              cur_s_cur_part_vertices = cur_s_cur_part_vertices[:, :3] # 
+            cur_s_cur_part_new_faces = []
+            for cur_s_cur_part_cur_face in cur_s_cur_part_faces:
+                cur_s_cur_part_cur_new_face = [fid + cur_s_n_vertices for fid in cur_s_cur_part_cur_face]
+                cur_s_cur_part_new_faces.append(cur_s_cur_part_cur_new_face)
+            cur_s_n_vertices += cur_s_cur_part_vertices.shape[0]
+            cur_s_tot_vertices.append(cur_s_cur_part_vertices)
+            # print(f"i_s: {i_s}, i_p: {i_p}, n_vertices: {cur_s_cur_part_vertices.shape[0]}")
+            cur_s_tot_faces += cur_s_cur_part_new_faces
+
+        cur_s_tot_vertices = np.concatenate(cur_s_tot_vertices, axis=0)
+        # print(f"i_s: {i_s}, n_cur_s_tot_vertices: {cur_s_tot_vertices.shape[0]}")
+        cur_s_mesh_dict = {
+            'vertices': cur_s_tot_vertices, 'faces': cur_s_tot_faces
+        }
+        tot_samples_mesh_dict.append(cur_s_mesh_dict)
+    # plot_meshes(tot_samples_mesh_dict, ax_lims=0.5, mesh_sv_fn=f"./figs/training_step_{n}_part_{tot_n_part}_rot.png")  # plot the mesh;
+    for i_s in range(tot_n_samples):
+      cur_mesh = tot_samples_mesh_dict[i_s]
+      cur_mesh_vertices, cur_mesh_faces = cur_mesh['vertices'], cur_mesh['faces']
+      # rotated mesh fn
+      cur_mesh_sv_fn = os.path.join(sv_mesh_folder, f"training_step_{cur_step}_ins_{i_s}_rot.obj")
+      # write object in the file...
+      if cur_mesh_vertices.shape[0] > 0 and len(cur_mesh_faces) > 0:
+        write_obj(cur_mesh_vertices, cur_mesh_faces, cur_mesh_sv_fn, transpose=False, scale=1.)
+
+
+
+
+def sample_pts_from_mesh(vertices, faces, npoints=512, minus_one=True):
+    return vertices
+    sampled_pcts = []
+    pts_to_seg_idx = []
+    # triangles and pts
+    minus_val = 0 if not minus_one else 1
+    for i in range(len(faces)): # 
+        cur_face = faces[i]
+        n_tris = len(cur_face) - 2
+        v_as, v_bs, v_cs = [cur_face[0] for _ in range(n_tris)], cur_face[1: len(cur_face) - 1], cur_face[2: len(cur_face)]
+        for v_a, v_b, v_c in zip(v_as, v_bs, v_cs):
+          
+            v_a, v_b, v_c = vertices[v_a - minus_val], vertices[v_b - minus_val], vertices[v_c - minus_val]
+            ab, ac = v_b - v_a, v_c - v_a
+            cos_ab_ac = (np.sum(ab * ac) / np.clip(np.sqrt(np.sum(ab ** 2)) * np.sqrt(np.sum(ac ** 2)), a_min=1e-9,
+                                                  a_max=9999999.)).item()
+            sin_ab_ac = math.sqrt(min(max(0., 1. - cos_ab_ac ** 2), 1.))
+            cur_area = 0.5 * sin_ab_ac * np.sqrt(np.sum(ab ** 2)).item() * np.sqrt(np.sum(ac ** 2)).item()
+
+            cur_sampled_pts = int(cur_area * npoints)
+            cur_sampled_pts = 1 if cur_sampled_pts == 0 else cur_sampled_pts
+            # if cur_sampled_pts == 0:
+
+            tmp_x, tmp_y = np.random.uniform(0, 1., (cur_sampled_pts,)).tolist(), np.random.uniform(0., 1., (
+            cur_sampled_pts,)).tolist()
+
+            for xx, yy in zip(tmp_x, tmp_y):
+                sqrt_xx, sqrt_yy = math.sqrt(xx), math.sqrt(yy)
+                aa = 1. - sqrt_xx
+                bb = sqrt_xx * (1. - yy)
+                cc = yy * sqrt_xx
+                cur_pos = v_a * aa + v_b * bb + v_c * cc
+                sampled_pcts.append(cur_pos)
+                # pts_to_seg_idx.append(cur_tri_seg)
+
+    # seg_idx_to_sampled_pts = {}
+    sampled_pcts = np.array(sampled_pcts, dtype=np.float)
+
+    return sampled_pcts
+
+
+def fps_fr_numpy(np_pts, n_sampling=4096):
+  pts = torch.from_numpy(np_pts).float().cuda()
+  pts_fps_idx = farthest_point_sampling(pts.unsqueeze(0), n_sampling=n_sampling) # farthes points sampling ##
+  pts = pts[pts_fps_idx].cpu()
+  return pts
+
+
+def farthest_point_sampling(pos: torch.FloatTensor, n_sampling: int):
+    bz, N = pos.size(0), pos.size(1)
+    feat_dim = pos.size(-1)
+    device = pos.device
+    sampling_ratio = float(n_sampling / N)
+    pos_float = pos.float()
+
+    batch = torch.arange(bz, dtype=torch.long).view(bz, 1).to(device)
+    mult_one = torch.ones((N,), dtype=torch.long).view(1, N).to(device)
+
+    batch = batch * mult_one
+    batch = batch.view(-1)
+    pos_float = pos_float.contiguous().view(-1, feat_dim).contiguous() # (bz x N, 3)
+    # sampling_ratio = torch.tensor([sampling_ratio for _ in range(bz)], dtype=torch.float).to(device)
+    # batch = torch.zeros((N, ), dtype=torch.long, device=device)
+    sampled_idx = fps(pos_float, batch, ratio=sampling_ratio, random_start=False)
+    # shape of sampled_idx?
+    return sampled_idx
+
+
+def plot_sampled_meshes_single_part(v_sample, f_sample, sv_mesh_folder, cur_step=0, predict_joint=True,):
+  
+  if not os.path.exists(sv_mesh_folder):
+    os.mkdir(sv_mesh_folder)
+
+  part_vertex_samples = [v_sample]
+  part_face_samples = [f_sample]
+
+
+  tot_n_samples = part_vertex_samples[0]['vertices'].shape[0]
+  tot_n_part = 2
+
+  # not predict joints here
+  # if predict_joint:
+  #   pred_dir = v_sample['joint_dir']
+  #   pred_pvp = v_sample['joint_pvp']
+  #   print("pred_dir", pred_dir.shape, pred_dir)
+  #   print("pred_pvp", pred_pvp.shape, pred_pvp)
+  # else:
+  #   pred_pvp = np.zeros(shape=[tot_n_samples, 3], dtype=np.float32)
+  
+
+  tot_mesh_list = []
+  for i_p, (cur_part_v_samples_np, cur_part_f_samples_np) in enumerate(zip(part_vertex_samples, part_face_samples)):
+      mesh_list = []
+      for i_n in range(tot_n_samples):
+          mesh_list.append(
+              {
+                  'vertices': cur_part_v_samples_np['vertices'][i_n][:cur_part_v_samples_np['num_vertices'][i_n]],
+                  'faces': unflatten_faces(
+                      cur_part_f_samples_np['faces'][i_n][:cur_part_f_samples_np['num_face_indices'][i_n]])
+              }
+          )
+      tot_mesh_list.append(mesh_list)
+
+      for i_n in range(tot_n_samples):
+        cur_mesh = mesh_list[i_n]
+        cur_mesh_vertices, cur_mesh_faces = cur_mesh['vertices'], cur_mesh['faces']
+        # cur_mesh_sv_fn = os.path.join("./meshes", f"training_step_{cur_step}_part_{i_p}_ins_{i_n}.obj")
+        cur_mesh_sv_fn = os.path.join(sv_mesh_folder, f"training_step_{cur_step}_part_{i_p}_ins_{i_n}.obj")
+        print(f"saving to {cur_mesh_sv_fn}, nn_verts: {cur_mesh_vertices.shape[0]}, nn_faces: {len(cur_mesh_faces)}")
+        if cur_mesh_vertices.shape[0] > 0 and len(cur_mesh_faces) > 0:
+          write_obj(cur_mesh_vertices, cur_mesh_faces, cur_mesh_sv_fn, transpose=True, scale=1.)
+        
+
+def plot_sampled_meshes(v_sample, f_sample, sv_mesh_folder, cur_step=0, predict_joint=True,):
+  
+  if not os.path.exists(sv_mesh_folder):
+    os.mkdir(sv_mesh_folder)
+
+  part_vertex_samples = [v_sample]
+  part_face_samples = [f_sample]
+
+
+  tot_n_samples = part_vertex_samples[0]['vertices'].shape[0]
+  # tot_n_part = 2
+
+  # not predict joints here
+  # if predict_joint:
+  #   pred_dir = v_sample['joint_dir']
+  #   pred_pvp = v_sample['joint_pvp']
+  #   print("pred_dir", pred_dir.shape, pred_dir)
+  #   print("pred_pvp", pred_pvp.shape, pred_pvp)
+  # else:
+  #   pred_pvp = np.zeros(shape=[tot_n_samples, 3], dtype=np.float32)
+  
+
+  tot_mesh_list = []
+  for i_p, (cur_part_v_samples_np, cur_part_f_samples_np) in enumerate(zip(part_vertex_samples, part_face_samples)):
+      mesh_list = []
+      for i_n in range(tot_n_samples):
+          mesh_list.append(
+              {
+                  'vertices': cur_part_v_samples_np['vertices'][i_n][:cur_part_v_samples_np['num_vertices'][i_n]],
+                  'faces': unflatten_faces(
+                      cur_part_f_samples_np['faces'][i_n][:cur_part_f_samples_np['num_face_indices'][i_n]])
+              }
+          )
+      tot_mesh_list.append(mesh_list)
+
+      for i_n in range(tot_n_samples):
+        cur_mesh = mesh_list[i_n]
+        cur_mesh_vertices, cur_mesh_faces = cur_mesh['vertices'], cur_mesh['faces']
+        # cur_mesh_sv_fn = os.path.join("./meshes", f"training_step_{cur_step}_part_{i_p}_ins_{i_n}.obj")
+        cur_mesh_sv_fn = os.path.join(sv_mesh_folder, f"training_step_{cur_step}_part_{i_p}_ins_{i_n}.obj")
+        print(f"saving to {cur_mesh_sv_fn}, nn_verts: {cur_mesh_vertices.shape[0]}, nn_faces: {len(cur_mesh_faces)}")
+        if cur_mesh_vertices.shape[0] > 0 and len(cur_mesh_faces) > 0:
+          write_obj(cur_mesh_vertices, cur_mesh_faces, cur_mesh_sv_fn, transpose=True, scale=1.)
+         
+
+def filter_masked_vertices(vertices, mask_indicator):
+  # vertices: n_verts x 3 
+  mask_indicator = np.reshape(mask_indicator, (vertices.shape[0], 3))
+  tot_vertices = []
+  for i_v in range(vertices.shape[0]):
+    cur_vert = vertices[i_v]
+    cur_vert_indicator = mask_indicator[i_v][0].item()
+    if cur_vert_indicator < 0.5:
+      tot_vertices.append(cur_vert)
+  tot_vertices = np.array(tot_vertices)
+  return tot_vertices
+
+
+def plot_sampled_ar_subd_meshes(v_sample, f_sample, sv_mesh_folder, cur_step=0, ):
+  if not os.path.exists(sv_mesh_folder): ### vertices_mask
+    os.mkdir(sv_mesh_folder)
+  ### v_sample: bsz x nn_verts x 3
+  vertices_mask = v_sample['vertices_mask']
+  vertices = v_sample['vertices']
+  faces = f_sample['faces']
+  num_face_indices = f_sample['num_face_indices'] #### num_faces_indices
+  bsz = vertices.shape[0]
+  
+  for i_bsz in range(bsz):
+    cur_vertices = vertices[i_bsz]
+    cur_vertices_mask = vertices_mask[i_bsz]
+    cur_faces = faces[i_bsz]
+    cur_num_face_indices = num_face_indices[i_bsz]
+    cur_nn_verts = cur_vertices_mask.sum(-1).item()
+    cur_nn_verts = int(cur_nn_verts)
+    cur_vertices = cur_vertices[:cur_nn_verts]
+    cur_faces = unflatten_faces(
+                      cur_faces[:int(cur_num_face_indices)])
+
+    cur_num_faces = len(cur_faces)
+    cur_mesh_sv_fn = os.path.join(sv_mesh_folder, f"training_step_{cur_step}_inst_{i_bsz}.obj")
+    # cur_context_mesh_sv_fn = os.path.join(sv_mesh_folder, f"training_step_{cur_step}_part_{i_p}_ins_{i_n}_context.obj")
+    print(f"saving to {cur_mesh_sv_fn}, nn_verts: {cur_nn_verts}, nn_faces: {cur_num_faces}")
+    if cur_nn_verts > 0 and cur_num_faces > 0:
+      write_obj(cur_vertices, cur_faces, cur_mesh_sv_fn, transpose=True, scale=1.)
+    
+  
+
+def plot_sampled_meshes_single_part_for_pretraining(v_sample, f_sample, context, sv_mesh_folder, cur_step=0, predict_joint=True, with_context=True):
+  
+  if not os.path.exists(sv_mesh_folder):
+    os.mkdir(sv_mesh_folder)
+
+  part_vertex_samples = [v_sample]
+  part_face_samples = [f_sample]
+
+  context_vertices = [context['vertices']]
+  context_faces = [context['faces']]
+  context_vertices_mask = [context['vertices_mask']]
+  context_faces_mask = [context['faces_mask']]
+
+
+  tot_n_samples = part_vertex_samples[0]['vertices'].shape[0]
+  tot_n_part = 2
+
+  # not predict joints here
+  # if predict_joint:
+  #   pred_dir = v_sample['joint_dir']
+  #   pred_pvp = v_sample['joint_pvp']
+  #   print("pred_dir", pred_dir.shape, pred_dir)
+  #   print("pred_pvp", pred_pvp.shape, pred_pvp)
+  # else:
+  #   pred_pvp = np.zeros(shape=[tot_n_samples, 3], dtype=np.float32)
+
+  # 
+  
+
+  tot_mesh_list = []
+  for i_p, (cur_part_v_samples_np, cur_part_f_samples_np) in enumerate(zip(part_vertex_samples, part_face_samples)):
+      mesh_list = []
+      context_mesh_list = []
+      for i_n in range(tot_n_samples):
+          mesh_list.append(
+              {
+                  'vertices': cur_part_v_samples_np['vertices'][i_n][:cur_part_v_samples_np['num_vertices'][i_n]],
+                  'faces': unflatten_faces(
+                      cur_part_f_samples_np['faces'][i_n][:cur_part_f_samples_np['num_face_indices'][i_n]])
+              }
+          )
+
+          cur_context_vertices = context_vertices[i_p][i_n]
+          cur_context_faces = context_faces[i_p][i_n]
+          cur_context_vertices_mask = context_vertices_mask[i_p][i_n]
+          cur_context_faces_mask = context_faces_mask[i_p][i_n]
+          cur_nn_vertices = np.sum(cur_context_vertices_mask).item()
+          cur_nn_faces = np.sum(cur_context_faces_mask).item()
+          cur_nn_vertices, cur_nn_faces = int(cur_nn_vertices), int(cur_nn_faces)
+          cur_context_vertices  = cur_context_vertices[:cur_nn_vertices]
+          if 'mask_vertices_flat_indicator' in context:
+            cur_context_vertices_mask_indicator = context['mask_vertices_flat_indicator'][i_n]
+            cur_context_vertices_mask_indicator = cur_context_vertices_mask_indicator[:cur_nn_vertices * 3]
+            cur_context_vertices = filter_masked_vertices(cur_context_vertices, cur_context_vertices_mask_indicator)
+          cur_context_faces = cur_context_faces[:cur_nn_faces] # context faces...
+          context_mesh_dict = {
+            'vertices': dequantize_verts(cur_context_vertices, n_bits=8), 'faces': unflatten_faces(cur_context_faces)
+          }
+          context_mesh_list.append(context_mesh_dict)
+
+      tot_mesh_list.append(mesh_list)
+
+      # if with_context:
+      for i_n in range(tot_n_samples):
+        cur_mesh = mesh_list[i_n]
+        cur_context_mesh = context_mesh_list[i_n]
+        cur_mesh_vertices, cur_mesh_faces = cur_mesh['vertices'], cur_mesh['faces']
+        cur_context_vertices, cur_context_faces = cur_context_mesh['vertices'], cur_context_mesh['faces']
+        # cur_mesh_sv_fn = os.path.join("./meshes", f"training_step_{cur_step}_part_{i_p}_ins_{i_n}.obj")
+        cur_mesh_sv_fn = os.path.join(sv_mesh_folder, f"training_step_{cur_step}_part_{i_p}_ins_{i_n}.obj")
+        cur_context_mesh_sv_fn = os.path.join(sv_mesh_folder, f"training_step_{cur_step}_part_{i_p}_ins_{i_n}_context.obj")
+        print(f"saving to {cur_mesh_sv_fn}, nn_verts: {cur_mesh_vertices.shape[0]}, nn_faces: {len(cur_mesh_faces)}")
+        if cur_mesh_vertices.shape[0] > 0 and len(cur_mesh_faces) > 0:
+          write_obj(cur_mesh_vertices, cur_mesh_faces, cur_mesh_sv_fn, transpose=True, scale=1.)
+        if cur_context_vertices.shape[0] > 0 and len(cur_context_faces) > 0:
+          write_obj(cur_context_vertices, cur_context_faces, cur_context_mesh_sv_fn, transpose=True, scale=1.)
+
+
+def plot_grids_for_pretraining_ml(v_sample, sv_mesh_folder="", cur_step=0, context=None):
+  
+  if not os.path.exists(sv_mesh_folder):
+    os.mkdir(sv_mesh_folder)
+
+  mesh_list = []
+  context_mesh_list = []
+  tot_n_samples = v_sample['vertices'].shape[0]
+
+  for i_n in range(tot_n_samples):
+      mesh_list.append(
+          {
+              'vertices': v_sample['vertices'][i_n][:v_sample['num_vertices'][i_n]],
+              'faces': []
+          }
+      )
+
+      cur_context_vertices = context['vertices'][i_n]
+      cur_context_vertices_mask = context['vertices_mask'][i_n]
+      cur_nn_vertices = np.sum(cur_context_vertices_mask).item()
+      cur_nn_vertices = int(cur_nn_vertices)
+      cur_context_vertices  = cur_context_vertices[:cur_nn_vertices]
+      if 'mask_vertices_flat_indicator' in context:
+        cur_context_vertices_mask_indicator = context['mask_vertices_flat_indicator'][i_n]
+        cur_context_vertices_mask_indicator = cur_context_vertices_mask_indicator[:cur_nn_vertices * 3]
+        cur_context_vertices = filter_masked_vertices(cur_context_vertices, cur_context_vertices_mask_indicator)
+      context_mesh_dict = {
+        'vertices': dequantize_verts(cur_context_vertices, n_bits=8), 'faces': []
+      }
+      context_mesh_list.append(context_mesh_dict)
+
+  # tot_mesh_list.append(mesh_list)
+
+  # if with_context:
+  for i_n in range(tot_n_samples):
+    cur_mesh = mesh_list[i_n]
+    cur_context_mesh = context_mesh_list[i_n]
+    cur_mesh_vertices = cur_mesh['vertices']
+    cur_context_vertices = cur_context_mesh['vertices']
+    # cur_mesh_sv_fn = os.path.join("./meshes", f"training_step_{cur_step}_part_{i_p}_ins_{i_n}.obj")
+    cur_mesh_sv_fn = os.path.join(sv_mesh_folder, f"training_step_{cur_step}_ins_{i_n}.obj")
+    cur_context_mesh_sv_fn = os.path.join(sv_mesh_folder, f"training_step_{cur_step}_ins_{i_n}_context.obj")
+    # print(f"saving to {cur_mesh_sv_fn}, nn_verts: {cur_mesh_vertices.shape[0]}, nn_faces: {len(cur_mesh_faces)}")
+    print(f"saving the sample to {cur_mesh_sv_fn}, context sample to {cur_context_mesh_sv_fn}")
+    if cur_mesh_vertices.shape[0] > 0:
+      write_obj(cur_mesh_vertices, None, cur_mesh_sv_fn, transpose=True, scale=1.)
+    if cur_context_vertices.shape[0] > 0:
+      write_obj(cur_context_vertices, None, cur_context_mesh_sv_fn, transpose=True, scale=1.)
+  
+
+def get_grid_content_from_grids(grid_xyzs, grid_values, grid_size=2):
+  cur_bsz_grid_xyzs = grid_xyzs # grid_length x 3 # grids pts for a sinlge batch
+  cur_bsz_grid_values = grid_values  # grid_length x gs x gs x gs
+  pts = []
+  for i_grid in range(cur_bsz_grid_xyzs.shape[0]): # cur_bsz_grid_xyzs
+    cur_grid_xyz = cur_bsz_grid_xyzs[i_grid].tolist()
+    if cur_grid_xyz[0] == -1 or cur_grid_xyz[1] == -1 or cur_grid_xyz[2] == -1:
+      break
+    if len(cur_bsz_grid_values.shape) > 1: 
+      cur_grid_values = cur_bsz_grid_values[i_grid]
+    else:
+      cur_grid_content = cur_bsz_grid_values[i_grid].item()
+      if cur_grid_content >= MASK_GRID_VALIE:
+        continue
+      inde = 2
+      cur_grid_values = []
+      for i_s in range(grid_size ** 3):
+        cur_mod_value = cur_grid_content % inde
+        cur_grid_content = cur_grid_content // inde
+        cur_grid_values = [cur_mod_value] + cur_grid_values # higher values should be put to the front of the list
+      cur_grid_values = np.array(cur_grid_values, dtype=np.long)
+      cur_grid_values = np.reshape(cur_grid_values, (grid_size, grid_size, grid_size))
+  # if words 
+  # flip words 
+    for i_x in range(cur_grid_values.shape[0]):
+      for i_y in range(cur_grid_values.shape[1]):
+        for i_z in range(cur_grid_values.shape[2]):
+          cur_grid_xyz_value = int(cur_grid_values[i_x, i_y, i_z].item())
+          if cur_grid_xyz_value > 0.5:
+            cur_x, cur_y, cur_z = cur_grid_xyz[0] * grid_size + i_x, cur_grid_xyz[1] * grid_size + i_y, cur_grid_xyz[2] * grid_size + i_z
+            pts.append([cur_x, cur_y, cur_z])
+  return pts
+
+def plot_grids_for_pretraining(v_sample, sv_mesh_folder="", cur_step=0, context=None, sv_mesh_fn=None):
+  
+  ##### plot grids
+  if not os.path.exists(sv_mesh_folder):
+    os.mkdir(sv_mesh_folder)
+
+  # part_vertex_samples = [v_sample] # vertex samples
+  # part_face_samples = [f_sample] # face samples
+
+  grid_xyzs = v_sample['grid_xyzs']
+  grid_values = v_sample['grid_values']
+
+  bsz = grid_xyzs.shape[0]
+  grid_size = opt.vertex_model.grid_size
+  
+
+  for i_bsz in range(bsz):
+    cur_bsz_grid_xyzs = grid_xyzs[i_bsz] # grid_length x 3
+    cur_bsz_grid_values = grid_values[i_bsz] # grid_length x gs x gs x gs
+    pts = []
+    for i_grid in range(cur_bsz_grid_xyzs.shape[0]): # cur_bsz_grid_xyzs
+      cur_grid_xyz = cur_bsz_grid_xyzs[i_grid].tolist()
+      if cur_grid_xyz[0] == -1 or cur_grid_xyz[1] == -1 or cur_grid_xyz[2] == -1:
+        break
+      if len(cur_bsz_grid_values.shape) > 1: 
+        cur_grid_values = cur_bsz_grid_values[i_grid]
+      else:
+        cur_grid_content = cur_bsz_grid_values[i_grid].item()
+        if cur_grid_content >= MASK_GRID_VALIE:
+          continue
+        inde = 2
+        cur_grid_values = []
+        for i_s in range(grid_size ** 3):
+          cur_mod_value = cur_grid_content % inde
+          cur_grid_content = cur_grid_content // inde
+          cur_grid_values = [cur_mod_value] + cur_grid_values # higher values should be put to the front of the list
+        cur_grid_values = np.array(cur_grid_values, dtype=np.long)
+        cur_grid_values = np.reshape(cur_grid_values, (grid_size, grid_size, grid_size))
+    # if 
+      for i_x in range(cur_grid_values.shape[0]):
+        for i_y in range(cur_grid_values.shape[1]):
+          for i_z in range(cur_grid_values.shape[2]):
+            cur_grid_xyz_value = int(cur_grid_values[i_x, i_y, i_z].item())
+            if cur_grid_xyz_value > 0.5:
+              cur_x, cur_y, cur_z = cur_grid_xyz[0] * grid_size + i_x, cur_grid_xyz[1] * grid_size + i_y, cur_grid_xyz[2] * grid_size + i_z
+              pts.append([cur_x, cur_y, cur_z])
+        
+    
+    if len(pts) == 0:
+      print("zzz, len(pts) == 0")
+      continue
+    pts = np.array(pts, dtype=np.float32)
+    # pts = center_vertices(pts)
+    # pts = normalize_vertices_scale(pts)
+    pts = pts[:, [2, 1, 0]]
+    cur_mesh_sv_fn = os.path.join(sv_mesh_folder, f"training_step_{cur_step}_ins_{i_bsz}.obj" if sv_mesh_fn is None else sv_mesh_fn)
+    
+    print(f"Saving obj to {cur_mesh_sv_fn}")
+    write_obj(pts, None, cur_mesh_sv_fn, transpose=True, scale=1.)
+
+
+def plot_grids_for_pretraining_obj_corpus(v_sample, sv_mesh_folder="", cur_step=0, context=None, sv_mesh_fn=None):
+  
+  ##### plot grids
+  if not os.path.exists(sv_mesh_folder):
+    os.mkdir(sv_mesh_folder)
+
+  # part_vertex_samples = [v_sample] # vertex samples
+  # part_face_samples = [f_sample] # face samples
+
+  grid_xyzs = v_sample['grid_xyzs']
+  grid_values = v_sample['grid_values']
+
+  bsz = grid_xyzs.shape[0]
+  grid_size = opt.vertex_model.grid_size
+  
+
+  for i_bsz in range(bsz):
+    cur_bsz_grid_xyzs = grid_xyzs[i_bsz] # grid_length x 3
+    cur_bsz_grid_values = grid_values[i_bsz] # grid_length x gs x gs x gs
+    part_pts = []
+    pts = []
+    for i_grid in range(cur_bsz_grid_xyzs.shape[0]): # cur_bsz_grid_xyzs
+      cur_grid_xyz = cur_bsz_grid_xyzs[i_grid].tolist()
+      ##### grid_xyz; grid_
+      if cur_grid_xyz[0] == -1 and cur_grid_xyz[1] == -1 and cur_grid_xyz[2] == -1:
+        part_pts.append(pts)
+        pts = []
+        continue
+      ##### cur_grid_xyz... #####
+      elif not (cur_grid_xyz[0] >= 0 and cur_grid_xyz[1] >= 0 and cur_grid_xyz[2] >= 0):
+        continue
+      if len(cur_bsz_grid_values.shape) > 1: 
+        cur_grid_values = cur_bsz_grid_values[i_grid]
+      else:
+        cur_grid_content = cur_bsz_grid_values[i_grid].item()
+        if cur_grid_content >= MASK_GRID_VALIE: # mask grid value
+          continue
+        inde = 2
+        cur_grid_values = []
+        for i_s in range(grid_size ** 3):
+          cur_mod_value = cur_grid_content % inde
+          cur_grid_content = cur_grid_content // inde
+          cur_grid_values = [cur_mod_value] + cur_grid_values # higher values should be put to the front of the list
+        cur_grid_values = np.array(cur_grid_values, dtype=np.long)
+        cur_grid_values = np.reshape(cur_grid_values, (grid_size, grid_size, grid_size))
+    # if 
+      for i_x in range(cur_grid_values.shape[0]):
+        for i_y in range(cur_grid_values.shape[1]):
+          for i_z in range(cur_grid_values.shape[2]):
+            cur_grid_xyz_value = int(cur_grid_values[i_x, i_y, i_z].item())
+            ##### gird-xyz-values #####
+            if cur_grid_xyz_value > 0.5: # cur_grid_xyz_value
+              cur_x, cur_y, cur_z = cur_grid_xyz[0] * grid_size + i_x, cur_grid_xyz[1] * grid_size + i_y, cur_grid_xyz[2] * grid_size + i_z
+              pts.append([cur_x, cur_y, cur_z])
+      
+    if len(pts) > 0:
+      part_pts.append(pts)
+      pts = []
+    tot_nn_pts = sum([len(aa) for aa in part_pts])
+    if tot_nn_pts == 0:
+      print("zzz, tot_nn_pts == 0")
+      continue
+
+    for i_p, pts in enumerate(part_pts):
+      if len(pts) == 0:
+        continue
+      pts = np.array(pts, dtype=np.float32)
+      pts = center_vertices(pts)
+      # pts = normalize_vertices_scale(pts)
+      pts = pts[:, [2, 1, 0]]
+      cur_mesh_sv_fn = os.path.join(sv_mesh_folder, f"training_step_{cur_step}_ins_{i_bsz}_ip_{i_p}.obj" if sv_mesh_fn is None else sv_mesh_fn)
+      
+      print(f"Saving {i_p}-th part obj to {cur_mesh_sv_fn}")
+      write_obj(pts, None, cur_mesh_sv_fn, transpose=True, scale=1.)
+
+
+
+def plot_grids_for_pretraining_obj_part(v_sample, sv_mesh_folder="", cur_step=0, context=None, sv_mesh_fn=None):
+  
+  ##### plot grids
+  if not os.path.exists(sv_mesh_folder):
+    os.mkdir(sv_mesh_folder)
+
+  # part_vertex_samples = [v_sample] # vertex samples
+  # part_face_samples = [f_sample] # face samples
+
+  grid_xyzs = v_sample['grid_xyzs']
+  grid_values = v_sample['grid_values']
+
+  bsz = grid_xyzs.shape[0]
+  grid_size = opt.vertex_model.grid_size
+  
+
+  for i_bsz in range(bsz):
+    cur_bsz_grid_xyzs = grid_xyzs[i_bsz] # grid_length x 3
+    cur_bsz_grid_values = grid_values[i_bsz] # grid_length x gs x gs x gs
+    part_pts = []
+    pts = []
+    for i_grid in range(cur_bsz_grid_xyzs.shape[0]): # cur_bsz_grid_xyzs
+      cur_grid_xyz = cur_bsz_grid_xyzs[i_grid].tolist()
+      ##### grid_xyz; grid_
+      if cur_grid_xyz[0] == -1 and cur_grid_xyz[1] == -1 and cur_grid_xyz[2] == -1:
+        part_pts.append(pts)
+        pts = []
+        break
+      elif cur_grid_xyz[0] == -1 and cur_grid_xyz[1] == -1 and cur_grid_xyz[2] == 0:
+        part_pts.append(pts)
+        pts = []
+        continue
+      ##### cur_grid_xyz... #####
+      elif not (cur_grid_xyz[0] >= 0 and cur_grid_xyz[1] >= 0 and cur_grid_xyz[2] >= 0):
+        continue
+      if len(cur_bsz_grid_values.shape) > 1: 
+        cur_grid_values = cur_bsz_grid_values[i_grid]
+      else:
+        cur_grid_content = cur_bsz_grid_values[i_grid].item()
+        if cur_grid_content >= MASK_GRID_VALIE: # invalid jor dummy content value s
+          continue
+        inde = 2
+        cur_grid_values = []
+        for i_s in range(grid_size ** 3):
+          cur_mod_value = cur_grid_content % inde
+          cur_grid_content = cur_grid_content // inde
+          cur_grid_values = [cur_mod_value] + cur_grid_values # higher values should be put to the front of the list
+        cur_grid_values = np.array(cur_grid_values, dtype=np.long)
+        cur_grid_values = np.reshape(cur_grid_values, (grid_size, grid_size, grid_size))
+    # if 
+      for i_x in range(cur_grid_values.shape[0]):
+        for i_y in range(cur_grid_values.shape[1]):
+          for i_z in range(cur_grid_values.shape[2]):
+            cur_grid_xyz_value = int(cur_grid_values[i_x, i_y, i_z].item())
+            ##### gird-xyz-values #####
+            if cur_grid_xyz_value > 0.5: # cur_grid_xyz_value
+              cur_x, cur_y, cur_z = cur_grid_xyz[0] * grid_size + i_x, cur_grid_xyz[1] * grid_size + i_y, cur_grid_xyz[2] * grid_size + i_z
+              pts.append([cur_x, cur_y, cur_z])
+      
+    if len(pts) > 0:
+      part_pts.append(pts)
+      pts = []
+    tot_nn_pts = sum([len(aa) for aa in part_pts])
+    if tot_nn_pts == 0:
+      print("zzz, tot_nn_pts == 0")
+      continue
+
+    for i_p, pts in enumerate(part_pts):
+      if len(pts) == 0:
+        continue
+      pts = np.array(pts, dtype=np.float32)
+      pts = center_vertices(pts)
+      # pts = normalize_vertices_scale(pts)
+      pts = pts[:, [2, 1, 0]]
+      cur_mesh_sv_fn = os.path.join(sv_mesh_folder, f"training_step_{cur_step}_ins_{i_bsz}_ip_{i_p}.obj" if sv_mesh_fn is None else sv_mesh_fn)
+      
+      print(f"Saving {i_p}-th part obj to {cur_mesh_sv_fn}")
+      write_obj(pts, None, cur_mesh_sv_fn, transpose=True, scale=1.)
+
+
+def plot_grids_for_pretraining_ml(v_sample, sv_mesh_folder="", cur_step=0, context=None):
+  
+  if not os.path.exists(sv_mesh_folder):
+    os.mkdir(sv_mesh_folder)
+
+  # part_vertex_samples = [v_sample] # vertex samples
+  # part_face_samples = [f_sample] # face samples
+
+  grid_xyzs = v_sample['grid_xyzs']
+  grid_values = v_sample['grid_values']
+
+  context_grid_xyzs = context['grid_xyzs'] - 1
+  # context_grid_values = context['grid_content']
+  context_grid_values = context['mask_grid_content']
+
+  bsz = grid_xyzs.shape[0]
+  grid_size = opt.vertex_model.grid_size
+  
+
+  for i_bsz in range(bsz):
+    cur_bsz_grid_pts = get_grid_content_from_grids(grid_xyzs[i_bsz], grid_values[i_bsz], grid_size=grid_size)
+    cur_context_grid_pts = get_grid_content_from_grids(context_grid_xyzs[i_bsz], context_grid_values[i_bsz], grid_size=grid_size)
+
+    if len(cur_bsz_grid_pts) > 0 and len(cur_context_grid_pts) > 0:
+      cur_bsz_grid_pts = np.array(cur_bsz_grid_pts, dtype=np.float32)
+      cur_bsz_grid_pts = center_vertices(cur_bsz_grid_pts)
+      cur_bsz_grid_pts = normalize_vertices_scale(cur_bsz_grid_pts)
+      cur_bsz_grid_pts = cur_bsz_grid_pts[:, [2, 1, 0]]
+      #### plot current mesh / sampled points ####
+      cur_mesh_sv_fn = os.path.join(sv_mesh_folder, f"training_step_{cur_step}_ins_{i_bsz}.obj")
+      print(f"Saving predicted grid points to {cur_mesh_sv_fn}")
+      write_obj(cur_bsz_grid_pts, None, cur_mesh_sv_fn, transpose=True, scale=1.)
+
+      cur_context_grid_pts = np.array(cur_context_grid_pts, dtype=np.float32)
+      cur_context_grid_pts = center_vertices(cur_context_grid_pts)
+      cur_context_grid_pts = normalize_vertices_scale(cur_context_grid_pts)
+      cur_context_grid_pts = cur_context_grid_pts[:, [2, 1, 0]]
+      #### plot current mesh / sampled points ####
+      cur_context_mesh_sv_fn = os.path.join(sv_mesh_folder, f"training_step_{cur_step}_ins_{i_bsz}_context.obj")
+      print(f"Saving context grid points to {cur_context_mesh_sv_fn}")
+      write_obj(cur_context_grid_pts, None, cur_context_mesh_sv_fn, transpose=True, scale=1.)
+
+    # print(f"Saving obj to {cur_mesh_sv_fn}")
+    # write_obj(pts, None, cur_mesh_sv_fn, transpose=True, scale=1.)
+
+    # if len(cur_bsz_grid_pts) == 0:
+    #   print("zzz, len(pts) == 0")
+    #   continue
+    # pts = np.array(pts, dtype=np.float32)
+    # pts = center_vertices(pts)
+    # pts = normalize_vertices_scale(pts)
+    # pts = pts[:, [2, 1, 0]]
+    # cur_mesh_sv_fn = os.path.join(sv_mesh_folder, f"training_step_{cur_step}_ins_{i_bsz}.obj")
+    
+    # print(f"Saving obj to {cur_mesh_sv_fn}")
+    # write_obj(pts, None, cur_mesh_sv_fn, transpose=True, scale=1.)
+
+
+
+def plot_sampled_meshes_single_part_for_sampling(v_sample, f_sample, sv_mesh_folder, cur_step=0, predict_joint=True,):
+  
+  if not os.path.exists(sv_mesh_folder):
+    os.mkdir(sv_mesh_folder)
+
+  part_vertex_samples = [v_sample]
+  part_face_samples = [f_sample]
+
+
+  tot_n_samples = part_vertex_samples[0]['vertices'].shape[0]
+  tot_n_part = 2
+
+  # not predict joints here
+  # if predict_joint:
+  #   pred_dir = v_sample['joint_dir']
+  #   pred_pvp = v_sample['joint_pvp']
+  #   print("pred_dir", pred_dir.shape, pred_dir)
+  #   print("pred_pvp", pred_pvp.shape, pred_pvp)
+  # else:
+  #   pred_pvp = np.zeros(shape=[tot_n_samples, 3], dtype=np.float32)
+  
+
+  tot_mesh_list = []
+  for i_p, (cur_part_v_samples_np, cur_part_f_samples_np) in enumerate(zip(part_vertex_samples, part_face_samples)):
+      mesh_list = []
+      for i_n in range(tot_n_samples):
+          mesh_list.append(
+              {
+                  'vertices': cur_part_v_samples_np['vertices'][i_n][:cur_part_v_samples_np['num_vertices'][i_n]],
+                  'faces': unflatten_faces(
+                      cur_part_f_samples_np['faces'][i_n][:cur_part_f_samples_np['num_face_indices'][i_n]])
+              }
+          )
+      tot_mesh_list.append(mesh_list)
+
+      for i_n in range(tot_n_samples):
+        cur_mesh = mesh_list[i_n]
+        cur_mesh_vertices, cur_mesh_faces = cur_mesh['vertices'], cur_mesh['faces']
+        # cur_mesh_sv_fn = os.path.join("./meshes", f"training_step_{cur_step}_part_{i_p}_ins_{i_n}.obj")
+        cur_mesh_sv_fn = os.path.join(sv_mesh_folder, f"step_{cur_step}_part_{i_p}_ins_{i_n}.obj")
+        print(f"saving to {cur_mesh_sv_fn}, nn_verts: {cur_mesh_vertices.shape[0]}, nn_faces: {len(cur_mesh_faces)}")
+        if cur_mesh_vertices.shape[0] > 0 and len(cur_mesh_faces) > 0:
+          write_obj(cur_mesh_vertices, cur_mesh_faces, cur_mesh_sv_fn, transpose=True, scale=1.)
+        
+     

models/dataset.py

@@ -0,0 +1,359 @@
+import torch
+import torch.nn.functional as F
+import cv2 as cv
+import numpy as np
+import os
+from glob import glob
+from icecream import ic
+from scipy.spatial.transform import Rotation as Rot
+from scipy.spatial.transform import Slerp
+
+
+# This function is borrowed from IDR: https://github.com/lioryariv/idr
+def load_K_Rt_from_P(filename, P=None):
+    if P is None:
+        lines = open(filename).read().splitlines()
+        if len(lines) == 4:
+            lines = lines[1:]
+        lines = [[x[0], x[1], x[2], x[3]] for x in (x.split(" ") for x in lines)]
+        P = np.asarray(lines).astype(np.float32).squeeze()
+
+    out = cv.decomposeProjectionMatrix(P)
+    K = out[0]
+    R = out[1]
+    t = out[2]
+
+    K = K / K[2, 2]
+    intrinsics = np.eye(4)
+    intrinsics[:3, :3] = K
+
+    pose = np.eye(4, dtype=np.float32)
+    pose[:3, :3] = R.transpose()
+    pose[:3, 3] = (t[:3] / t[3])[:, 0]
+
+    return intrinsics, pose
+
+def filter_iamges_via_pixel_values(data_dir):
+    images_lis = sorted(glob(os.path.join(data_dir, 'image/*.png'))) ## images lis ##
+    n_images = len(images_lis)
+    images_np = np.stack([cv.imread(im_name) for im_name in images_lis]) / 255.0
+    print(f"images_np: {images_np.shape}")
+    # nn_frames x res x res x 3 #
+    images_np = 1. - images_np
+    has_density_values = (np.sum(images_np, axis=-1) > 0.7).astype(np.float32)
+    has_density_values = np.sum(np.sum(has_density_values, axis=-1), axis=-1)
+    tot_res_nns = float(images_np.shape[1] * images_np.shape[2])
+    has_density_ratio = has_density_values / tot_res_nns ### has density ratio and ratio # 
+    print(f"has_density_values: {has_density_values.shape}")
+    paried_has_density_ratio_list = [(i_fr, has_density_ratio[i_fr].item()) for i_fr in range(has_density_ratio.shape[0])]
+    paried_has_density_ratio_list = sorted(paried_has_density_ratio_list, key=lambda ii: ii[1], reverse=True)
+    mid_rnk_value = len(paried_has_density_ratio_list) // 4
+    print(f"mid value of the density ratio")
+    print(paried_has_density_ratio_list[mid_rnk_value])
+    iamge_idx = paried_has_density_ratio_list[mid_rnk_value][0]
+    print(f"iamge idx: {images_lis[iamge_idx]}")
+    print(paried_has_density_ratio_list[:mid_rnk_value])
+    tot_selected_img_idx_list = [ii[0] for ii in paried_has_density_ratio_list[:mid_rnk_value]]
+    tot_selected_img_idx_list =sorted(tot_selected_img_idx_list)
+    print(len(tot_selected_img_idx_list))
+    # print(tot_selected_img_idx_list[54])
+    print(tot_selected_img_idx_list)
+    
+    
+
+class Dataset:
+    def __init__(self, conf):
+        super(Dataset, self).__init__()
+        print('Load data: Begin')
+        self.device = torch.device('cuda')
+        self.conf = conf
+        
+        self.selected_img_idxes_list = [0, 1, 5, 6, 7, 8, 9, 13, 14, 15, 35, 36, 42, 43, 44, 48, 49, 50, 51, 55, 56, 57, 61, 62, 63, 69, 84, 90, 91, 92, 96, 97]
+        # self.selected_img_idxes_list = [0, 1, 5, 6, 7, 8, 9, 12, 13, 14, 15, 20, 21, 22, 23, 26, 27, 28, 29, 35, 36, 37, 40, 41, 70, 71, 79, 82, 83, 84, 85, 92, 93, 96, 97, 98, 99, 105, 106, 107, 110, 111, 112, 113, 118, 119, 120, 121, 124, 125, 133, 134, 135, 139, 174, 175, 176, 177, 180, 188, 189, 190, 191, 194, 195]
+        
+        self.selected_img_idxes_list = [0, 1, 6, 7, 8, 9, 12, 13, 14, 15, 20, 21, 22, 23, 26, 27, 36, 40, 41, 70, 71, 78, 82, 83, 84, 85, 90, 91, 92, 93, 96, 97]
+        
+        self.selected_img_idxes_list = [0, 1, 6, 7, 8, 9, 12, 13, 14, 15, 20, 21, 22, 23, 26, 27, 36, 40, 41, 70, 71, 78, 82, 83, 84, 85, 90, 91, 92, 93, 96, 97, 98, 99, 104, 105, 106, 107, 110, 111, 112, 113, 118, 119, 120, 121, 124, 125, 134, 135, 139, 174, 175, 176, 177, 180, 181, 182, 183, 188, 189, 190, 191, 194, 195]
+        
+        self.selected_img_idxes_list = [0, 1, 6, 7, 8, 9, 12, 13, 14, 20, 21, 22, 23, 26, 27, 70, 78, 83, 84, 85, 91, 92, 93, 96, 97, 98, 99, 105, 106, 107, 110, 111, 112, 113, 119, 120, 121, 124, 125, 175, 176, 181, 182, 188, 189, 190, 191, 194, 195]
+        # or the timestep to the dataset instance ## # selected img idxes list #
+        self.selected_img_idxes = np.array(self.selected_img_idxes_list).astype(np.int32)
+        
+        
+        
+       
+
+        self.data_dir = conf.get_string('data_dir')
+        self.render_cameras_name = conf.get_string('render_cameras_name')
+        self.object_cameras_name = conf.get_string('object_cameras_name')
+
+        ## camera outside sphere ## 
+        self.camera_outside_sphere = conf.get_bool('camera_outside_sphere', default=True)
+        self.scale_mat_scale = conf.get_float('scale_mat_scale', default=1.1)
+
+        camera_dict = np.load(os.path.join(self.data_dir, self.render_cameras_name))
+        # camera_dict = np.load("/home/xueyi/diffsim/NeuS/public_data/dtu_scan24/cameras_sphere.npz")
+        self.camera_dict = camera_dict # rendr camera dict #
+        # render camera dict # # number of pixels in the views -> very thin geometry is not useful 
+        self.images_lis = sorted(glob(os.path.join(self.data_dir, 'image/*.png')))
+        
+        # iamges_lis # and the images_lis and the images_lis #
+        # self.images_lis = self.images_lis[:1] # totoal views and poses of the camera; # and select cameras for rendering #
+        
+        self.n_images = len(self.images_lis)
+        self.images_np = np.stack([cv.imread(im_name) for im_name in self.images_lis]) / 256.0
+        
+        
+        self.selected_img_idxes_list = list(range(self.images_np.shape[0]))
+        self.selected_img_idxes = np.array(self.selected_img_idxes_list).astype(np.int32)
+        
+        self.images_np = self.images_np[self.selected_img_idxes] ## get selected iamges_np #
+        
+        ### if we deal with the backgound carefully ### ### get 
+        self.images_np = np.stack([cv.imread(im_name) for im_name in self.images_lis]) / 255.0
+        self.images_np = self.images_np[self.selected_img_idxes]
+        self.images_np = 1. - self.images_np ### 
+        
+        
+        self.masks_lis = sorted(glob(os.path.join(self.data_dir, 'mask/*.png')))
+        
+        # self.masks_lis = self.masks_lis[:1]
+        
+        try:
+            self.masks_np = np.stack([cv.imread(im_name) for im_name in self.masks_lis]) / 256.0
+            self.masks_np = self.masks_np[self.selected_img_idxes]
+        except:
+            self.masks_np = self.images_np.copy()
+
+
+        
+        
+
+
+        # world_mat is a projection matrix from world to image
+        self.world_mats_np = [camera_dict['world_mat_%d' % idx].astype(np.float32) for idx in range(self.n_images)]
+
+        self.scale_mats_np = []
+
+        # scale_mat: used for coordinate normalization, we assume the scene to render is inside a unit sphere at origin.
+        self.scale_mats_np = [camera_dict['scale_mat_%d' % idx].astype(np.float32) for idx in range(self.n_images)]
+
+        self.intrinsics_all = []
+        self.pose_all = []
+
+        # for idx, (scale_mat, world_mat) in enumerate(zip(self.scale_mats_np, self.world_mats_np)):
+        for idx  in self.selected_img_idxes_list:
+            scale_mat = self.scale_mats_np[idx]
+            world_mat = self.world_mats_np[idx]
+            
+            if "hand" in self.data_dir:
+                intrinsics = np.eye(4)
+                fov = 512. / 2. # * 2
+                res = 512.
+                intrinsics[:3, :3] = np.array([
+                    [fov, 0, 0.5* res], # res #
+                    [0, fov, 0.5* res], # res #
+                    [0, 0, 1]
+                ], dtype=np.float32)
+                pose = camera_dict['camera_mat_%d' % idx].astype(np.float32)
+            else:
+                P = world_mat @ scale_mat
+                P = P[:3, :4]
+                intrinsics, pose = load_K_Rt_from_P(None, P)
+                
+            self.intrinsics_all.append(torch.from_numpy(intrinsics).float())
+            self.pose_all.append(torch.from_numpy(pose).float())
+
+        ### images, masks, 
+        self.images = torch.from_numpy(self.images_np.astype(np.float32)).cpu()  # [n_images, H, W, 3] #
+        self.masks  = torch.from_numpy(self.masks_np.astype(np.float32)).cpu()   # [n_images, H, W, 3] #
+        self.intrinsics_all = torch.stack(self.intrinsics_all).to(self.device)   # [n_images, 4, 4] # optimize sdf field # rigid model hand
+        self.intrinsics_all_inv = torch.inverse(self.intrinsics_all)  # [n_images, 4, 4]
+        self.focal = self.intrinsics_all[0][0, 0]
+        self.pose_all = torch.stack(self.pose_all).to(self.device)  # [n_images, 4, 4]
+        self.H, self.W = self.images.shape[1], self.images.shape[2]
+        self.image_pixels = self.H * self.W
+
+        object_bbox_min = np.array([-1.01, -1.01, -1.01, 1.0])
+        object_bbox_max = np.array([ 1.01,  1.01,  1.01, 1.0])
+        # Object scale mat: region of interest to **extract mesh**
+        object_scale_mat = np.load(os.path.join(self.data_dir, self.object_cameras_name))['scale_mat_0']
+        object_bbox_min = np.linalg.inv(self.scale_mats_np[0]) @ object_scale_mat @ object_bbox_min[:, None]
+        object_bbox_max = np.linalg.inv(self.scale_mats_np[0]) @ object_scale_mat @ object_bbox_max[:, None]
+        self.object_bbox_min = object_bbox_min[:3, 0]
+        self.object_bbox_max = object_bbox_max[:3, 0]
+        
+        self.n_images = self.images.size(0)
+
+        print('Load data: End')
+        
+    def get_rays(H, W, K, c2w, inverse_y, flip_x, flip_y, mode='center'):
+        i, j = torch.meshgrid( # meshgrid #
+            torch.linspace(0, W-1, W, device=c2w.device),
+            torch.linspace(0, H-1, H, device=c2w.device))
+        i = i.t().float()
+        j = j.t().float()
+        if mode == 'lefttop':
+            pass
+        elif mode == 'center':
+            i, j = i+0.5, j+0.5
+        elif mode == 'random':
+            i = i+torch.rand_like(i)
+            j = j+torch.rand_like(j)
+        else:
+            raise NotImplementedError
+
+        if flip_x:
+            i = i.flip((1,))
+        if flip_y:
+            j = j.flip((0,))
+        if inverse_y:
+            dirs = torch.stack([(i-K[0][2])/K[0][0], (j-K[1][2])/K[1][1], torch.ones_like(i)], -1)
+        else:
+            dirs = torch.stack([(i-K[0][2])/K[0][0], -(j-K[1][2])/K[1][1], -torch.ones_like(i)], -1)
+        # Rotate ray directions from camera frame to the world frame
+        rays_d = torch.sum(dirs[..., np.newaxis, :] * c2w[:3,:3], -1)  # dot product, equals to: [c2w.dot(dir) for dir in dirs]
+        # Translate camera frame's origin to the world frame. It is the origin of all rays.
+        rays_o = c2w[:3,3].expand(rays_d.shape)
+        return rays_o, rays_d
+
+    def gen_rays_at(self, img_idx, resolution_level=1):
+        """
+        Generate rays at world space from one camera.
+        """
+        l = resolution_level
+        tx = torch.linspace(0, self.W - 1, self.W // l)
+        ty = torch.linspace(0, self.H - 1, self.H // l)
+        pixels_x, pixels_y = torch.meshgrid(tx, ty)
+        
+        ##### previous method #####
+        # p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1) # W, H, 3
+        # # p = torch.stack([pixels_x, pixels_y, -1. * torch.ones_like(pixels_y)], dim=-1) # W, H, 3
+        # p = torch.matmul(self.intrinsics_all_inv[img_idx, None, None, :3, :3], p[:, :, :, None]).squeeze()  # W, H, 3
+        # rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)  # W, H, 3
+        # rays_v = torch.matmul(self.pose_all[img_idx, None, None, :3, :3], rays_v[:, :, :, None]).squeeze()  # W, H, 3
+        # rays_o = self.pose_all[img_idx, None, None, :3, 3].expand(rays_v.shape)  # W, H, 3
+        ##### previous method #####
+        
+        fov = 512.; res = 512.
+        K = np.array([
+            [fov, 0, 0.5* res],
+            [0, fov, 0.5* res],
+            [0, 0, 1]
+        ], dtype=np.float32)
+        K = torch.from_numpy(K).float().cuda()
+        
+        
+        # ### `center` mode ### #
+        c2w = self.pose_all[img_idx]
+        pixels_x, pixels_y = pixels_x+0.5, pixels_y+0.5
+        
+        dirs = torch.stack([(pixels_x-K[0][2])/K[0][0], -(pixels_y-K[1][2])/K[1][1], -torch.ones_like(pixels_x)], -1)
+        rays_v = torch.sum(dirs[..., np.newaxis, :] * c2w[:3,:3], -1) 
+        rays_o = c2w[:3,3].expand(rays_v.shape)
+        # dirs = torch.stack([(i-K[0][2])/K[0][0], -(j-K[1][2])/K[1][1], -torch.ones_like(i)], -1)
+        
+        # p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1) # W, H, 3
+        # # p = torch.stack([pixels_x, pixels_y, -1. * torch.ones_like(pixels_y)], dim=-1) # W, H, 3
+        # p = torch.matmul(self.intrinsics_all_inv[img_idx, None, None, :3, :3], p[:, :, :, None]).squeeze()  # W, H, 3
+        # rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)  # W, H, 3
+        # rays_v = torch.matmul(self.pose_all[img_idx, None, None, :3, :3], rays_v[:, :, :, None]).squeeze()  # W, H, 3
+        # rays_o = self.pose_all[img_idx, None, None, :3, 3].expand(rays_v.shape)  # W, H, 3
+        return rays_o.transpose(0, 1), rays_v.transpose(0, 1)
+
+    def gen_random_rays_at(self, img_idx, batch_size):
+        """
+        Generate random rays at world space from one camera.
+        """
+        pixels_x = torch.randint(low=0, high=self.W, size=[batch_size])
+        pixels_y = torch.randint(low=0, high=self.H, size=[batch_size])
+        color = self.images[img_idx][(pixels_y, pixels_x)]    # batch_size, 3
+        
+        mask = self.masks[img_idx][(pixels_y, pixels_x)]      # batch_size, 3
+        
+        
+        ##### previous method #####
+        # p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1).float()  # batch_size, 3
+        # # p = torch.stack([pixels_x, pixels_y, -1. * torch.ones_like(pixels_y)], dim=-1).float()  # batch_size, 3
+        # p = torch.matmul(self.intrinsics_all_inv[img_idx, None, :3, :3], p[:, :, None]).squeeze() # batch_size, 3
+        # rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)    # batch_size, 3
+        # rays_v = torch.matmul(self.pose_all[img_idx, None, :3, :3], rays_v[:, :, None]).squeeze()  # batch_size, 3
+        # rays_o = self.pose_all[img_idx, None, :3, 3].expand(rays_v.shape) # batch_size, 3
+        ##### previous method #####
+        
+        fov = 512.; res = 512.
+        K = np.array([
+            [fov, 0, 0.5* res],
+            [0, fov, 0.5* res],
+            [0, 0, 1]
+        ], dtype=np.float32)
+        K = torch.from_numpy(K).float().cuda()
+        
+        
+        # ### `center` mode ### #
+        c2w = self.pose_all[img_idx]
+        pixels_x, pixels_y = pixels_x+0.5, pixels_y+0.5
+        
+        dirs = torch.stack([(pixels_x-K[0][2])/K[0][0], -(pixels_y-K[1][2])/K[1][1], -torch.ones_like(pixels_x)], -1)
+        rays_v = torch.sum(dirs[..., np.newaxis, :] * c2w[:3,:3], -1) 
+        rays_o = c2w[:3,3].expand(rays_v.shape)
+        
+        
+        return torch.cat([rays_o.cpu(), rays_v.cpu(), color, mask[:, :1]], dim=-1).cuda()    # batch_size, 10
+
+    def gen_rays_between(self, idx_0, idx_1, ratio, resolution_level=1):
+        """
+        Interpolate pose between two cameras.
+        """
+        l = resolution_level
+        tx = torch.linspace(0, self.W - 1, self.W // l)
+        ty = torch.linspace(0, self.H - 1, self.H // l)
+        pixels_x, pixels_y = torch.meshgrid(tx, ty)
+        p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1)  # W, H, 3
+        p = torch.matmul(self.intrinsics_all_inv[0, None, None, :3, :3], p[:, :, :, None]).squeeze()  # W, H, 3
+        rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)  # W, H, 3
+        trans = self.pose_all[idx_0, :3, 3] * (1.0 - ratio) + self.pose_all[idx_1, :3, 3] * ratio
+        pose_0 = self.pose_all[idx_0].detach().cpu().numpy()
+        pose_1 = self.pose_all[idx_1].detach().cpu().numpy()
+        pose_0 = np.linalg.inv(pose_0)
+        pose_1 = np.linalg.inv(pose_1)
+        rot_0 = pose_0[:3, :3]
+        rot_1 = pose_1[:3, :3]
+        rots = Rot.from_matrix(np.stack([rot_0, rot_1]))
+        key_times = [0, 1]
+        slerp = Slerp(key_times, rots)
+        rot = slerp(ratio)
+        pose = np.diag([1.0, 1.0, 1.0, 1.0])
+        pose = pose.astype(np.float32)
+        pose[:3, :3] = rot.as_matrix()
+        pose[:3, 3] = ((1.0 - ratio) * pose_0 + ratio * pose_1)[:3, 3]
+        pose = np.linalg.inv(pose)
+        rot = torch.from_numpy(pose[:3, :3]).cuda()
+        trans = torch.from_numpy(pose[:3, 3]).cuda()
+        rays_v = torch.matmul(rot[None, None, :3, :3], rays_v[:, :, :, None]).squeeze()  # W, H, 3
+        rays_o = trans[None, None, :3].expand(rays_v.shape)  # W, H, 3
+        return rays_o.transpose(0, 1), rays_v.transpose(0, 1)
+
+    def near_far_from_sphere(self, rays_o, rays_d):
+        a = torch.sum(rays_d**2, dim=-1, keepdim=True)
+        b = 2.0 * torch.sum(rays_o * rays_d, dim=-1, keepdim=True)
+        mid = 0.5 * (-b) / a
+        near = mid - 1.0
+        far = mid + 1.0
+        return near, far
+
+    ## iamge_at ##
+    def image_at(self, idx, resolution_level):
+        if self.selected_img_idxes_list is not None:
+            img = cv.imread(self.images_lis[self.selected_img_idxes_list[idx]])
+        else:
+            img = cv.imread(self.images_lis[idx])
+        return (cv.resize(img, (self.W // resolution_level, self.H // resolution_level))).clip(0, 255)
+
+
+if __name__=='__main__':
+    data_dir = "/data/datasets/genn/diffsim/diffredmax/save_res/goal_optimize_model_hand_sphere_test_obj_type_active_nfr_10_view_divide_0.5_n_views_7_three_planes_False_recon_dvgo_new_Nposes_7_routine_2"
+    data_dir = "/data/datasets/genn/diffsim/neus/public_data/hand_test"
+    data_dir = "/data2/datasets/diffsim/neus/public_data/hand_test_routine_2"
+    data_dir = "/data2/datasets/diffsim/neus/public_data/hand_test_routine_2_light_color"
+    filter_iamges_via_pixel_values(data_dir=data_dir)

models/dataset_wtime.py

@@ -0,0 +1,403 @@
+import torch
+import torch.nn.functional as F
+import cv2 as cv
+import numpy as np
+import os
+from glob import glob
+from icecream import ic
+from scipy.spatial.transform import Rotation as Rot
+from scipy.spatial.transform import Slerp
+
+
+# This function is borrowed from IDR: https://github.com/lioryariv/idr
+def load_K_Rt_from_P(filename, P=None):
+    if P is None:
+        lines = open(filename).read().splitlines()
+        if len(lines) == 4:
+            lines = lines[1:]
+        lines = [[x[0], x[1], x[2], x[3]] for x in (x.split(" ") for x in lines)]
+        P = np.asarray(lines).astype(np.float32).squeeze()
+
+    out = cv.decomposeProjectionMatrix(P)
+    K = out[0]
+    R = out[1]
+    t = out[2]
+
+    K = K / K[2, 2]
+    intrinsics = np.eye(4)
+    intrinsics[:3, :3] = K
+
+    pose = np.eye(4, dtype=np.float32)
+    pose[:3, :3] = R.transpose()
+    pose[:3, 3] = (t[:3] / t[3])[:, 0]
+
+    return intrinsics, pose
+
+def filter_iamges_via_pixel_values(data_dir):
+    images_lis = sorted(glob(os.path.join(data_dir, 'image/*.png'))) ## images lis ##
+    n_images = len(images_lis)
+    images_np = np.stack([cv.imread(im_name) for im_name in images_lis]) / 255.0
+    print(f"images_np: {images_np.shape}")
+    # nn_frames x res x res x 3 #
+    images_np = 1. - images_np
+    has_density_values = (np.sum(images_np, axis=-1) > 0.7).astype(np.float32)
+    has_density_values = np.sum(np.sum(has_density_values, axis=-1), axis=-1)
+    tot_res_nns = float(images_np.shape[1] * images_np.shape[2])
+    has_density_ratio = has_density_values / tot_res_nns ### has density ratio and ratio # 
+    print(f"has_density_values: {has_density_values.shape}")
+    paried_has_density_ratio_list = [(i_fr, has_density_ratio[i_fr].item()) for i_fr in range(has_density_ratio.shape[0])]
+    paried_has_density_ratio_list = sorted(paried_has_density_ratio_list, key=lambda ii: ii[1], reverse=True)
+    mid_rnk_value = len(paried_has_density_ratio_list) // 4
+    print(f"mid value of the density ratio")
+    print(paried_has_density_ratio_list[mid_rnk_value])
+    iamge_idx = paried_has_density_ratio_list[mid_rnk_value][0]
+    print(f"iamge idx: {images_lis[iamge_idx]}")
+    print(paried_has_density_ratio_list[:mid_rnk_value])
+    tot_selected_img_idx_list = [ii[0] for ii in paried_has_density_ratio_list[:mid_rnk_value]]
+    tot_selected_img_idx_list =sorted(tot_selected_img_idx_list)
+    print(len(tot_selected_img_idx_list))
+    # print(tot_selected_img_idx_list[54])
+    print(tot_selected_img_idx_list)
+    
+    
+
+class Dataset:
+    def __init__(self, conf, time_idx, mode='train'):
+        super(Dataset, self).__init__()
+        print('Load data: Begin')
+        self.device = torch.device('cuda')
+        self.conf = conf
+        
+        self.selected_img_idxes_list = [0, 1, 5, 6, 7, 8, 9, 13, 14, 15, 35, 36, 42, 43, 44, 48, 49, 50, 51, 55, 56, 57, 61, 62, 63, 69, 84, 90, 91, 92, 96, 97]
+        # self.selected_img_idxes_list = [0, 1, 5, 6, 7, 8, 9, 12, 13, 14, 15, 20, 21, 22, 23, 26, 27, 28, 29, 35, 36, 37, 40, 41, 70, 71, 79, 82, 83, 84, 85, 92, 93, 96, 97, 98, 99, 105, 106, 107, 110, 111, 112, 113, 118, 119, 120, 121, 124, 125, 133, 134, 135, 139, 174, 175, 176, 177, 180, 188, 189, 190, 191, 194, 195]
+        
+        self.selected_img_idxes_list = [0, 1, 6, 7, 8, 9, 12, 13, 14, 15, 20, 21, 22, 23, 26, 27, 36, 40, 41, 70, 71, 78, 82, 83, 84, 85, 90, 91, 92, 93, 96, 97]
+        
+        self.selected_img_idxes_list = [0, 1, 6, 7, 8, 9, 12, 13, 14, 15, 20, 21, 22, 23, 26, 27, 36, 40, 41, 70, 71, 78, 82, 83, 84, 85, 90, 91, 92, 93, 96, 97, 98, 99, 104, 105, 106, 107, 110, 111, 112, 113, 118, 119, 120, 121, 124, 125, 134, 135, 139, 174, 175, 176, 177, 180, 181, 182, 183, 188, 189, 190, 191, 194, 195]
+        # selected img idxes list #
+        self.selected_img_idxes_list = [0, 1, 6, 7, 8, 9, 12, 13, 14, 20, 21, 22, 23, 26, 27, 70, 78, 83, 84, 85, 91, 92, 93, 96, 97, 98, 99, 105, 106, 107, 110, 111, 112, 113, 119, 120, 121, 124, 125, 175, 176, 181, 182, 188, 189, 190, 191, 194, 195]
+        # or the timestep to the dataset instance ## # selected img idxes list #
+        self.selected_img_idxes = np.array(self.selected_img_idxes_list).astype(np.int32)
+        
+        
+        
+       
+
+        self.data_dir = conf.get_string('data_dir')
+        
+        self.data_dir = os.path.join(self.data_dir, f"{time_idx}") # the time_idx #
+        
+        self.render_cameras_name = conf.get_string('render_cameras_name')
+        self.object_cameras_name = conf.get_string('object_cameras_name')
+
+        ## camera outside sphere ## 
+        self.camera_outside_sphere = conf.get_bool('camera_outside_sphere', default=True)
+        self.scale_mat_scale = conf.get_float('scale_mat_scale', default=1.1)
+
+        camera_dict = np.load(os.path.join(self.data_dir, self.render_cameras_name))
+        # camera_dict = np.load("/home/xueyi/diffsim/NeuS/public_data/dtu_scan24/cameras_sphere.npz")
+        self.camera_dict = camera_dict # rendr camera dict #
+        # render camera dict # # number of pixels in the views -> very thin geometry is not useful 
+        self.images_lis = sorted(glob(os.path.join(self.data_dir, 'image/*.png')))
+        
+        # iamges_lis # and the images_lis and the images_lis #
+        # self.images_lis = self.images_lis[:1] # totoal views and poses of the camera; # and select cameras for rendering #
+        
+        self.n_images = len(self.images_lis)
+        
+        if mode == 'train_from_model_rules':
+            self.images_np = cv.imread(self.images_lis[0]) / 256.0
+            print(self.images_np.shape)
+            self.images_np = np.reshape(self.images_np, (1, self.images_np.shape[0], self.images_np.shape[1], self.images_np.shape[2]))
+            self.images_np = [self.images_np for _ in range(len(self.images_lis))]
+            self.images_np = np.concatenate(self.images_np, axis=0)
+        else:
+            presaved_imags_npy_fn = os.path.join(self.data_dir, "processed_images.npy")
+            if not os.path.exists(presaved_imags_npy_fn):
+                self.images_np = []
+                for i_im_idx, im_name in enumerate(self.images_lis):
+                    print(f"loading {i_im_idx} / {len(self.images_lis)}")
+                    cur_im = cv.imread(im_name) # for im_name in self.images_lis
+                    self.images_np.append(cur_im)
+                self.images_np = np.stack(self.images_np) / 256.0
+                np.save(presaved_imags_npy_fn, self.images_np)
+            else:
+                print(f"Loading from {presaved_imags_npy_fn}")
+                self.images_np = np.load(presaved_imags_npy_fn, allow_pickle=True)
+
+        # self.images_np = np.stack([cv.imread(im_name) for im_name in self.images_lis]) / 256.0
+        
+        
+        # self.selected_img_idxes_list = list(range(self.images_np.shape[0]))
+        # self.selected_img_idxes = np.array(self.selected_img_idxes_list).astype(np.int32)
+        
+        # get 
+        self.images_np = self.images_np[self.selected_img_idxes] ## get selected iamges_np #
+        
+        ### if we deal with the backgound carefully ### ### get 
+        # self.images_np = np.stack([cv.imread(im_name) for im_name in self.images_lis]) / 255.0
+        # self.images_np = self.images_np[self.selected_img_idxes]
+        self.images_np = 1. - self.images_np ### 
+        
+        
+        
+        self.masks_lis = sorted(glob(os.path.join(self.data_dir, 'mask/*.png')))
+        
+        if mode == 'train_from_model_rules':
+            self.masks_np = cv.imread(self.masks_lis[0]) / 256.0
+            print("masks shape:", self.masks_np.shape)
+            self.masks_np = np.reshape(self.masks_np, (1, self.masks_np.shape[0], self.masks_np.shape[1], self.masks_np.shape[2])) # .repeat(len(self.masks_lis), 1, 1)
+            self.masks_np = [self.masks_np for _ in range(len(self.masks_lis))]
+            self.masks_np = np.concatenate(self.masks_np, axis=0)
+        else:
+            presaved_masks_npy_fn = os.path.join(self.data_dir, "processed_masks.npy")
+            # self.masks_lis = self.masks_lis[:1]
+            
+            if not os.path.exists(presaved_masks_npy_fn):
+                try:
+                    self.masks_np = np.stack([cv.imread(im_name) for im_name in self.masks_lis]) / 256.0
+                    self.masks_np = self.masks_np[self.selected_img_idxes]
+                except:
+                    self.masks_np = self.images_np.copy()
+                np.save(presaved_masks_npy_fn, self.masks_np)
+            else:
+                print(f"Loading from {presaved_masks_npy_fn}")
+                self.masks_np = np.load(presaved_masks_npy_fn, allow_pickle=True)
+
+
+        
+        
+
+
+        # world_mat is a projection matrix from world to image
+        self.world_mats_np = [camera_dict['world_mat_%d' % idx].astype(np.float32) for idx in range(self.n_images)]
+
+        self.scale_mats_np = []
+
+        # scale_mat: used for coordinate normalization, we assume the scene to render is inside a unit sphere at origin.
+        self.scale_mats_np = [camera_dict['scale_mat_%d' % idx].astype(np.float32) for idx in range(self.n_images)]
+
+        self.intrinsics_all = []
+        self.pose_all = []
+
+        # for idx, (scale_mat, world_mat) in enumerate(zip(self.scale_mats_np, self.world_mats_np)):
+        for idx  in self.selected_img_idxes_list:
+            scale_mat = self.scale_mats_np[idx]
+            world_mat = self.world_mats_np[idx]
+            
+            if "hand" in self.data_dir:
+                intrinsics = np.eye(4)
+                fov = 512. / 2. # * 2
+                res = 512.
+                intrinsics[:3, :3] = np.array([
+                    [fov, 0, 0.5* res], # res #
+                    [0, fov, 0.5* res], # res #
+                    [0, 0, 1]
+                ], dtype=np.float32)
+                pose = camera_dict['camera_mat_%d' % idx].astype(np.float32)
+            else:
+                P = world_mat @ scale_mat
+                P = P[:3, :4]
+                intrinsics, pose = load_K_Rt_from_P(None, P)
+                
+            self.intrinsics_all.append(torch.from_numpy(intrinsics).float())
+            self.pose_all.append(torch.from_numpy(pose).float())
+
+        ### images, masks, 
+        self.images = torch.from_numpy(self.images_np.astype(np.float32)).cpu()  # [n_images, H, W, 3] #
+        self.masks  = torch.from_numpy(self.masks_np.astype(np.float32)).cpu()   # [n_images, H, W, 3] #
+        self.intrinsics_all = torch.stack(self.intrinsics_all).to(self.device)   # [n_images, 4, 4] # optimize sdf field # rigid model hand
+        self.intrinsics_all_inv = torch.inverse(self.intrinsics_all)  # [n_images, 4, 4]
+        self.focal = self.intrinsics_all[0][0, 0]
+        self.pose_all = torch.stack(self.pose_all).to(self.device)  # [n_images, 4, 4]
+        self.H, self.W = self.images.shape[1], self.images.shape[2]
+        self.image_pixels = self.H * self.W
+
+        object_bbox_min = np.array([-1.01, -1.01, -1.01, 1.0])
+        object_bbox_max = np.array([ 1.01,  1.01,  1.01, 1.0])
+        # Object scale mat: region of interest to **extract mesh**
+        object_scale_mat = np.load(os.path.join(self.data_dir, self.object_cameras_name))['scale_mat_0']
+        object_bbox_min = np.linalg.inv(self.scale_mats_np[0]) @ object_scale_mat @ object_bbox_min[:, None]
+        object_bbox_max = np.linalg.inv(self.scale_mats_np[0]) @ object_scale_mat @ object_bbox_max[:, None]
+        self.object_bbox_min = object_bbox_min[:3, 0]
+        self.object_bbox_max = object_bbox_max[:3, 0]
+        
+        self.n_images = self.images.size(0)
+
+        print('Load data: End')
+        
+    def get_rays(H, W, K, c2w, inverse_y, flip_x, flip_y, mode='center'):
+        i, j = torch.meshgrid( # meshgrid #
+            torch.linspace(0, W-1, W, device=c2w.device),
+            torch.linspace(0, H-1, H, device=c2w.device))
+        i = i.t().float()
+        j = j.t().float()
+        if mode == 'lefttop':
+            pass
+        elif mode == 'center':
+            i, j = i+0.5, j+0.5
+        elif mode == 'random':
+            i = i+torch.rand_like(i)
+            j = j+torch.rand_like(j)
+        else:
+            raise NotImplementedError
+
+        if flip_x:
+            i = i.flip((1,))
+        if flip_y:
+            j = j.flip((0,))
+        if inverse_y:
+            dirs = torch.stack([(i-K[0][2])/K[0][0], (j-K[1][2])/K[1][1], torch.ones_like(i)], -1)
+        else:
+            dirs = torch.stack([(i-K[0][2])/K[0][0], -(j-K[1][2])/K[1][1], -torch.ones_like(i)], -1)
+        # Rotate ray directions from camera frame to the world frame
+        rays_d = torch.sum(dirs[..., np.newaxis, :] * c2w[:3,:3], -1)  # dot product, equals to: [c2w.dot(dir) for dir in dirs]
+        # Translate camera frame's origin to the world frame. It is the origin of all rays.
+        rays_o = c2w[:3,3].expand(rays_d.shape)
+        return rays_o, rays_d
+
+    def gen_rays_at(self, img_idx, resolution_level=1):
+        """
+        Generate rays at world space from one camera.
+        """
+        l = resolution_level
+        tx = torch.linspace(0, self.W - 1, self.W // l)
+        ty = torch.linspace(0, self.H - 1, self.H // l)
+        pixels_x, pixels_y = torch.meshgrid(tx, ty)
+        
+        ##### previous method #####
+        # p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1) # W, H, 3
+        # # p = torch.stack([pixels_x, pixels_y, -1. * torch.ones_like(pixels_y)], dim=-1) # W, H, 3
+        # p = torch.matmul(self.intrinsics_all_inv[img_idx, None, None, :3, :3], p[:, :, :, None]).squeeze()  # W, H, 3
+        # rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)  # W, H, 3
+        # rays_v = torch.matmul(self.pose_all[img_idx, None, None, :3, :3], rays_v[:, :, :, None]).squeeze()  # W, H, 3
+        # rays_o = self.pose_all[img_idx, None, None, :3, 3].expand(rays_v.shape)  # W, H, 3
+        ##### previous method #####
+        
+        fov = 512.; res = 512.
+        K = np.array([
+            [fov, 0, 0.5* res],
+            [0, fov, 0.5* res],
+            [0, 0, 1]
+        ], dtype=np.float32)
+        K = torch.from_numpy(K).float().cuda()
+        
+        
+        # ### `center` mode ### #
+        c2w = self.pose_all[img_idx]
+        pixels_x, pixels_y = pixels_x+0.5, pixels_y+0.5
+        
+        dirs = torch.stack([(pixels_x-K[0][2])/K[0][0], -(pixels_y-K[1][2])/K[1][1], -torch.ones_like(pixels_x)], -1)
+        rays_v = torch.sum(dirs[..., np.newaxis, :] * c2w[:3,:3], -1) 
+        rays_o = c2w[:3,3].expand(rays_v.shape)
+        # dirs = torch.stack([(i-K[0][2])/K[0][0], -(j-K[1][2])/K[1][1], -torch.ones_like(i)], -1)
+        
+        # p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1) # W, H, 3
+        # # p = torch.stack([pixels_x, pixels_y, -1. * torch.ones_like(pixels_y)], dim=-1) # W, H, 3
+        # p = torch.matmul(self.intrinsics_all_inv[img_idx, None, None, :3, :3], p[:, :, :, None]).squeeze()  # W, H, 3
+        # rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)  # W, H, 3
+        # rays_v = torch.matmul(self.pose_all[img_idx, None, None, :3, :3], rays_v[:, :, :, None]).squeeze()  # W, H, 3
+        # rays_o = self.pose_all[img_idx, None, None, :3, 3].expand(rays_v.shape)  # W, H, 3
+        return rays_o.transpose(0, 1), rays_v.transpose(0, 1)
+
+    def gen_random_rays_at(self, img_idx, batch_size):
+        """
+        Generate random rays at world space from one camera.
+        """
+        img_idx = img_idx.cpu()
+        pixels_x = torch.randint(low=0, high=self.W, size=[batch_size]).cpu()
+        pixels_y = torch.randint(low=0, high=self.H, size=[batch_size]).cpu()
+        
+        # print(self.images.device, img_idx.device, pixels_y.device)
+        color = self.images[img_idx][(pixels_y, pixels_x)]    # batch_size, 3
+        
+        mask = self.masks[img_idx][(pixels_y, pixels_x)]      # batch_size, 3
+        
+        
+        ##### previous method #####
+        # p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1).float()  # batch_size, 3
+        # # p = torch.stack([pixels_x, pixels_y, -1. * torch.ones_like(pixels_y)], dim=-1).float()  # batch_size, 3
+        # p = torch.matmul(self.intrinsics_all_inv[img_idx, None, :3, :3], p[:, :, None]).squeeze() # batch_size, 3
+        # rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)    # batch_size, 3
+        # rays_v = torch.matmul(self.pose_all[img_idx, None, :3, :3], rays_v[:, :, None]).squeeze()  # batch_size, 3
+        # rays_o = self.pose_all[img_idx, None, :3, 3].expand(rays_v.shape) # batch_size, 3
+        ##### previous method #####
+        
+        fov = 512.; res = 512.
+        K = np.array([
+            [fov, 0, 0.5* res],
+            [0, fov, 0.5* res],
+            [0, 0, 1]
+        ], dtype=np.float32)
+        K = torch.from_numpy(K).float().cuda()
+        
+        
+        # ### `center` mode ### #
+        c2w = self.pose_all[img_idx]
+        
+        pixels_x = pixels_x.cuda()
+        pixels_y = pixels_y.cuda()
+        pixels_x, pixels_y = pixels_x+0.5, pixels_y+0.5
+        
+        dirs = torch.stack([(pixels_x-K[0][2])/K[0][0], -(pixels_y-K[1][2])/K[1][1], -torch.ones_like(pixels_x)], -1)
+        rays_v = torch.sum(dirs[..., np.newaxis, :] * c2w[:3,:3], -1) 
+        rays_o = c2w[:3,3].expand(rays_v.shape)
+        
+        
+        return torch.cat([rays_o.cpu(), rays_v.cpu(), color, mask[:, :1]], dim=-1).cuda()    # batch_size, 10
+
+    def gen_rays_between(self, idx_0, idx_1, ratio, resolution_level=1):
+        """
+        Interpolate pose between two cameras.
+        """
+        l = resolution_level
+        tx = torch.linspace(0, self.W - 1, self.W // l)
+        ty = torch.linspace(0, self.H - 1, self.H // l)
+        pixels_x, pixels_y = torch.meshgrid(tx, ty)
+        p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1)  # W, H, 3
+        p = torch.matmul(self.intrinsics_all_inv[0, None, None, :3, :3], p[:, :, :, None]).squeeze()  # W, H, 3
+        rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)  # W, H, 3
+        trans = self.pose_all[idx_0, :3, 3] * (1.0 - ratio) + self.pose_all[idx_1, :3, 3] * ratio
+        pose_0 = self.pose_all[idx_0].detach().cpu().numpy()
+        pose_1 = self.pose_all[idx_1].detach().cpu().numpy()
+        pose_0 = np.linalg.inv(pose_0)
+        pose_1 = np.linalg.inv(pose_1)
+        rot_0 = pose_0[:3, :3]
+        rot_1 = pose_1[:3, :3]
+        rots = Rot.from_matrix(np.stack([rot_0, rot_1]))
+        key_times = [0, 1]
+        slerp = Slerp(key_times, rots)
+        rot = slerp(ratio)
+        pose = np.diag([1.0, 1.0, 1.0, 1.0])
+        pose = pose.astype(np.float32)
+        pose[:3, :3] = rot.as_matrix()
+        pose[:3, 3] = ((1.0 - ratio) * pose_0 + ratio * pose_1)[:3, 3]
+        pose = np.linalg.inv(pose)
+        rot = torch.from_numpy(pose[:3, :3]).cuda()
+        trans = torch.from_numpy(pose[:3, 3]).cuda()
+        rays_v = torch.matmul(rot[None, None, :3, :3], rays_v[:, :, :, None]).squeeze()  # W, H, 3
+        rays_o = trans[None, None, :3].expand(rays_v.shape)  # W, H, 3
+        return rays_o.transpose(0, 1), rays_v.transpose(0, 1)
+
+    def near_far_from_sphere(self, rays_o, rays_d):
+        a = torch.sum(rays_d**2, dim=-1, keepdim=True)
+        b = 2.0 * torch.sum(rays_o * rays_d, dim=-1, keepdim=True)
+        mid = 0.5 * (-b) / a
+        near = mid - 1.0
+        far = mid + 1.0
+        return near, far
+
+    def image_at(self, idx, resolution_level):
+        if self.selected_img_idxes_list is not None:
+            img = cv.imread(self.images_lis[self.selected_img_idxes_list[idx]])
+        else:
+            img = cv.imread(self.images_lis[idx])
+        return (cv.resize(img, (self.W // resolution_level, self.H // resolution_level))).clip(0, 255)
+
+
+if __name__=='__main__':
+    data_dir = "/data/datasets/genn/diffsim/diffredmax/save_res/goal_optimize_model_hand_sphere_test_obj_type_active_nfr_10_view_divide_0.5_n_views_7_three_planes_False_recon_dvgo_new_Nposes_7_routine_2"
+    data_dir = "/data/datasets/genn/diffsim/neus/public_data/hand_test"
+    data_dir = "/data2/datasets/diffsim/neus/public_data/hand_test_routine_2"
+    data_dir = "/data2/datasets/diffsim/neus/public_data/hand_test_routine_2_light_color"
+    filter_iamges_via_pixel_values(data_dir=data_dir)

models/dyn_model_act_v2_deformable.py

@@ -0,0 +1,1582 @@
+
+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
+
+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).cuda()
+        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 = "/home/xueyi/diffsim/NeuS/rsc/mano" ## mano models of the mesh root ##
+        if not os.path.exists(mesh_root):
+            mesh_root = "/data/xueyi/diffsim/NeuS/rsc/mano"
+        self.mesh_root = mesh_root
+        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 = 
+        mesh = trimesh.load_mesh(self.geometry_mesh_fn)
+        vertices = mesh.vertices
+        faces = mesh.faces
+        
+        vertices = torch.from_numpy(vertices).float().cuda()
+        faces =torch.from_numpy(faces).long().cuda()
+        
+        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):
+        # cur_vertices = torch.matmul(parent_rot, self.vertices.transpose(1, 0)).contiguous().transpose(1, 0).contiguous() + parent_trans.unsqueeze(0)
+        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)
+        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_save_fn = f"{self.mesh_nm}_expanded_pts_factor_{expand_factor}_nnexp_{nn_expand_pts}.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)).cuda()
+                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().cuda()
+        self.cur_expanded_visual_pts = cur_expanded_visual_pts # expanded visual pts #
+        return self.cur_expanded_visual_pts
+        # cur_pts #
+        # use r as the search direction # # expande save fn #
+    def get_transformed_visual_pts(self, visual_pts_list):
+        visual_pts_list.append(self.cur_expanded_visual_pts) # 
+        return visual_pts_list
+            
+        
+
+## 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 #
+## Link_urdf ##
+class Link_urdf: # get_transformed_visual_pts #
+    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
+        
+        ### dyn_model_act ###
+        # parent_rot_mtx, parent_trans_vec #
+        # parent_rot_mtx, parent_trans_vec # # link urdf # 
+        # self.parent_rot_mtx = nn.Parameter(torch.eye(n=3, dtype=torch.float32).cuda(), requires_grad=True)
+        # self.parent_trans_vec = nn.Parameter(torch.zeros((3,), dtype=torch.float32).cuda(), requires_grad=True)
+        # self.curr_rot_mtx = nn.Parameter(torch.eye(n=3, dtype=torch.float32).cuda(), requires_grad=True)
+        # self.curr_trans_vec = nn.Parameter(torch.zeros((3,), dtype=torch.float32).cuda(), requires_grad=True)
+        # # 
+        # self.tot_rot_mtx = nn.Parameter(torch.eye(n=3, dtype=torch.float32).cuda(), requires_grad=True)
+        # self.tot_trans_vec = nn.Parameter(torch.zeros((3,), dtype=torch.float32).cuda(), requires_grad=True)
+
+    # expand visual pts #
+    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 = cur_link_struct.expand_visual_pts(expanded_visual_pts, link_name_to_visited, link_name_to_link_struct)
+        return expanded_visual_pts
+    
+    def get_transformed_visual_pts(self, visual_pts_list, 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_child_name = self.children[cur_joint_name]
+                cur_child = link_name_to_link_struct[cur_child_name] # parent and the child_visual, cur_child.visual #
+                # parent #
+                # 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 ## transformed_visual_pts ####
+                if cur_child_name in link_name_to_visited:
+                    continue
+                # cur_child_visual_pts = {'vertices': [], 'faces': [], 'link_idxes': [], 'transformed_joint_pos': [], 'joint_link_idxes': []}
+                cur_child_visual_pts_list = []
+                cur_child_visual_pts_list = cur_child.get_transformed_visual_pts(cur_child_visual_pts_list, link_name_to_visited, link_name_to_link_struct)
+                
+                if len(cur_child_visual_pts_list) > 0:
+                    cur_child_visual_pts = torch.cat(cur_child_visual_pts_list, dim=0)
+                    # 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_visual_pts = cur_child_visual_pts + cur_joint.origin_xyz.unsqueeze(0)
+                    cur_joint_rot, cur_joint_trans = cur_joint.compute_transformation_from_current_state()
+                    cur_child_visual_pts = torch.matmul(cur_joint_rot, cur_child_visual_pts.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).cuda()], dim=-1
+                    #     )
+                    # else:
+                    #     joint_link_idxes = torch.tensor([cur_child.link_idx], dtype=torch.long).cuda().view(1,)
+                    
+                    
+                    visual_pts_list.append(cur_child_visual_pts)
+                    # 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)
+                    
+        
+        if self.visual is not None:
+            # get_transformed_visual_pts #
+            visual_pts_list = self.visual.get_transformed_visual_pts(visual_pts_list)
+        
+        # for cur_link in self.children:
+        #     cur_link_name = cur_link.name
+        #     if cur_link_name in link_name_to_visited: # link name to visited # 
+        #         continue
+        #     visual_pts_list = cur_link.get_transformed_visual_pts(visual_pts_list, link_name_to_visited, link_name_to_link_struct)
+        return visual_pts_list
+    
+    # use both the articulated motion and the frre form 
+    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 get_init_visual_meshes(self, parent_rot, parent_trans, init_visual_meshes, link_name_to_link_struct, link_name_to_visited):
+        link_name_to_visited[self.name] = 1
+        
+        # 'transformed_joint_pos': [], 'link_idxes': []
+        if self.joint is not None:
+            # 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]
+                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': []}
+                cur_child_visual_pts = 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)
+                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).cuda()], dim=-1
+                        )
+                    else:
+                        joint_link_idxes = torch.tensor([cur_child.link_idx], dtype=torch.long).cuda().view(1,)
+                    
+                    
+                    
+                    # 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
+        else:
+            joint_origin_xyz = torch.tensor([0., 0., 0.], dtype=torch.float32).cuda()
+        # self.parent_rot_mtx = parent_rot
+        # self.parent_trans_vec = parent_trans + joint_origin_xyz
+        
+        
+        if self.visual is not None:
+            init_visual_meshes = self.visual.get_init_visual_meshes(parent_rot, parent_trans, init_visual_meshes)
+            cur_visual_mesh_pts_nn = self.visual.vertices.size(0)
+            cur_link_idxes = torch.zeros((cur_visual_mesh_pts_nn, ), dtype=torch.long).cuda()+ self.link_idx
+            init_visual_meshes['link_idxes'].append(cur_link_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 ## 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).cuda()
+        
+        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)
+
+        
+        
+    # 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
+        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_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).cuda()
+            self.joint.state = self.joint.state + update_quaternion(self.rot_vec, self.joint.state)
+            self.joint.timestep_to_states[0] = self.joint.state.detach()
+            self.joint.timestep_to_vels[0] = torch.zeros((3,), dtype=torch.float32).cuda().detach() ## velocity ##
+        for cur_link in self.children:
+            cur_link.set_init_states_target_value(init_states)
+    
+    # should forward for one single step -> use the action #
+    def set_init_states(self, ):
+        self.joint.state = torch.tensor([1, 0, 0, 0], dtype=torch.float32).cuda()
+        self.joint.timestep_to_states[0] = self.joint.state.detach()
+        self.joint.timestep_to_vels[0] = torch.zeros((3,), dtype=torch.float32).cuda().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) -> 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
+        
+        # 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).cuda(), requires_grad=True
+        )
+        self.action = nn.Parameter(
+            torch.zeros((1,), dtype=torch.float32, requires_grad=True).cuda(), requires_grad=True
+        )
+        # self.rot_mtx = np.eye(3, dtypes=np.float32)
+        # self.trans_vec = np.zeros((3,), dtype=np.float32) ## rot m
+        self.rot_mtx = nn.Parameter(torch.eye(n=3, dtype=torch.float32, requires_grad=True).cuda(), requires_grad=True)
+        self.trans_vec  = nn.Parameter(torch.zeros((3,), dtype=torch.float32,  requires_grad=True).cuda(), requires_grad=True)
+        
+    def set_initial_state(self, state):
+        # joint angle as the state value #
+        self.timestep_to_vels[0] = torch.zeros((3,), dtype=torch.float32).cuda().detach() ## velocity ##
+        delta_rot_vec = self.axis_xyz * state
+        # self.timestep_to_states[0] = state.detach()
+        cur_state = torch.tensor([1., 0., 0., 0.], dtype=torch.float32).cuda()
+        init_state = cur_state + update_quaternion(delta_rot_vec, cur_state)
+        self.timestep_to_states[0] = init_state.detach()
+        self.state = init_state
+        
+    def set_delta_state_and_update(self, state, cur_timestep):
+        self.timestep_to_vels[cur_timestep] = torch.zeros((3,), dtype=torch.float32).cuda().detach() 
+        delta_rot_vec = self.axis_xyz * state
+        if cur_timestep == 0:
+            prev_state = torch.tensor([1., 0., 0., 0.], dtype=torch.float32).cuda()
+        else:
+            prev_state = self.timestep_to_states[cur_timestep - 1].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).cuda()
+            trans_vec = torch.zeros((3,), dtype=torch.float32).cuda()
+            # 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 #
+            torque = self.action * self.axis_xyz
+            
+            # # Compute inertia matrix #
+            # inertial = torch.zeros((3, 3), dtype=torch.float32).cuda()
+            # 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).cuda()
+            #     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 # 
+            
+            # timestep_to_vels, timestep_to_states, state #
+            
+            # TODO: dt should be an optimizable 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
+            else:
+                cur_angular_vel = delta_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)
+            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 #
+
+
+# get_transformed_visual_pts() --- transformed_visual_pts ##
+# use the transformed visual # the articulated motion field # 
+# then we should add the free motion field here # # add the free motion field #  # hwo to use that? #
+# another rules for optimizing articulation motion field # 
+# -> the articulated model predicted transformations #
+# -> the free motion field -> the motion field predicted by the network for each timestep -> an implicit motion field #
+
+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) -> 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
+        
+        # 
+        # particles 
+        # sample particles 
+        # how to sample particles? 
+        # how to expand the particles? # -> you can use weights in the model dict # 
+        # from grids and jample from grids # 
+        # link idx to the 
+        # robot # 
+        # init vertices, init faces # 
+        # expande the aprticles # 
+        # expanede particles # 
+        # use particles to conduct the simulation # 
+        
+        
+        
+        self.init_vertices, self.init_faces = self.get_init_visual_pts()
+        
+        
+        init_visual_pts_sv_fn = "robot_expanded_visual_pts.npy"
+        np.save(init_visual_pts_sv_fn, self.init_vertices.detach().cpu().numpy())
+        
+        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
+    
+    
+    
+    
+    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
+    
+    
+    # get_transformed_visual_pts() # get_transformed_visual_pts of the visual pts ### get_transformed_visual_pts ## get_transformed_visual_pts ### # 
+    def get_transformed_visual_pts(self, ):
+        init_visual_pts = []
+        link_name_to_visited = {}
+        
+        palm_idx = self.link_name_to_link_idxes["palm"]
+        palm_link = self.links[palm_idx]
+        
+        ### init_visual_pts # from the pal mink to get the total transformed visual pts ## 
+        init_visual_pts = palm_link.get_transformed_visual_pts(init_visual_pts, link_name_to_visited, self.link_name_to_link_struct)
+        
+        init_visual_pts = torch.cat(init_visual_pts, dim=0) ## get the inita visual pts from the palm link ### 
+        return init_visual_pts
+    
+    
+    ### samping issue? --- TODO`    `
+    def get_init_visual_pts(self, ):
+        init_visual_meshes = {
+            'vertices': [], 'faces': [], 'link_idxes': [], 'transformed_joint_pos': [], 'link_idxes': [],  'transformed_joint_pos': [], 'joint_link_idxes': []
+        }
+        init_parent_rot = torch.eye(3, dtype=torch.float32).cuda()
+        init_parent_trans = torch.zeros((3,), dtype=torch.float32).cuda()
+        
+        palm_idx = self.link_name_to_link_idxes["palm"]
+        palm_link = self.links[palm_idx]
+        
+        link_name_to_visited = {}
+        
+        init_visual_meshes = 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)
+        
+        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) ### 
+        
+        
+        
+        # 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'])
+        return init_vertices, init_faces
+    
+    
+    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)
+    
+    # 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,):
+        # actions 
+        # self.actions_joint_name_to_joint_idx as the action joint name to joint idx 
+        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_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)
+        
+    ### TODO: add the contact torque when calculating the nextstep states ###
+    ### TODO: not an accurate implementation since differen joints should be jconsidered for one single link ###
+    ### 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.cuda(), 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")
+    inertial_pos = origin.attrib["xyz"]
+    inertial_pos = parse_nparray_from_string(inertial_pos)
+    
+    inertial_rpy = origin.attrib["rpy"]
+    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).cuda()
+    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")
+        mesh_fn = geometry_mesh.attrib["filename"]
+        mesh_scale = geometry_mesh.attrib["scale"]
+        
+        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 = joint_origin.attrib["xyz"]
+        origin_xyz = parse_nparray_from_string(origin_xyz)
+    except:
+        origin_xyz = torch.tensor([0., 0., 0.], dtype=torch.float32).cuda()
+        
+    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).cuda()
+    
+    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)
+        joint_velocity = joint_limit.attrib["velocity"]
+        joint_velocity = float(joint_velocity)
+    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)
+    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")
+    # 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}")
+        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")
+    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
+        
+        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
+        
+        
+        #### 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
+        
+        
+    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_robots.append(cur_robot_obj)
+    
+    # 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
+
+
+
+class RobotAgent: # robot and the robot #
+    def __init__(self, xml_fn) -> None:
+        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
+        ).cuda()
+        torch.nn.init.ones_(self.time_constant.weight) #
+        self.time_constant.weight.data = self.time_constant.weight.data * 0.2 ### time_constant data #
+        
+        self.optimizable_actions = nn.Embedding(
+            num_embeddings=100, embedding_dim=60,
+        ).cuda()
+        torch.nn.init.zeros_(self.optimizable_actions.weight) #
+        
+        self.learning_rate = 5e-4
+    
+        self.active_robot = active_robot
+        
+        # # get init states # #
+        self.set_init_states()
+        init_visual_pts = self.get_init_state_visual_pts()
+        self.init_visual_pts = init_visual_pts
+        
+        
+    def set_init_states_target_value(self, init_states):
+        # glb_rot = torch.eye(n=3, dtype=torch.float32).cuda()
+        # glb_trans = torch.zeros((3,), dtype=torch.float32).cuda() ### glb_trans #### and the rot 3##
+
+        # tot_init_states = {}
+        # tot_init_states['glb_rot'] = glb_rot;
+        # tot_init_states['glb_trans'] = glb_trans; 
+        # tot_init_states['links_init_states'] = init_states
+        # self.active_robot.set_init_states_target_value(tot_init_states)
+        # init_joint_states = torch.zeros((60, ), dtype=torch.float32).cuda()
+        self.active_robot.set_initial_state(init_states)
+    
+    def set_init_states(self):
+        # glb_rot = torch.eye(n=3, dtype=torch.float32).cuda()
+        # glb_trans = torch.zeros((3,), dtype=torch.float32).cuda() ### glb_trans #### and the rot 3##
+        
+        # ### random rotation ###
+        # # glb_rot_np = R.random().as_matrix()
+        # # glb_rot = torch.from_numpy(glb_rot_np).float().cuda()
+        # ### random rotation ###
+        
+        # # glb_rot, glb_trans #
+        # init_states = {} # 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).cuda()
+        self.active_robot.set_initial_state(init_joint_states)
+    
+    def get_init_state_visual_pts(self, ):
+        # 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)
+        cur_verts, cur_faces = self.active_robot.get_init_visual_pts()
+        self.faces = cur_faces
+        # init_visual_pts = visual_pts_list
+        return cur_verts
+    
+    def set_actions_and_update_states(self, actions, cur_timestep):
+        # 
+        time_cons = self.time_constant(torch.zeros((1,), dtype=torch.long).cuda()) ### time constant of the system ##
+        self.active_robot.set_actions_and_update_states(actions, cur_timestep, time_cons) ### 
+        pass
+    
+    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).cuda()
+            actions['delta_glb_trans'] = torch.zeros((3,), dtype=torch.float32).cuda()
+            actions_link_actions = torch.ones((22, ), dtype=torch.float32).cuda()
+            # 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().cuda() 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).cuda()
+                # actions['delta_glb_trans'] = torch.zeros((3,), dtype=torch.float32).cuda()
+                actions_link_actions = self.optimizable_actions(torch.zeros((1,), dtype=torch.long).cuda() + cur_ts).squeeze(0)
+                # actions_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
+
+
+
+#### Big TODO: the external contact forces from the manipulated object to the robot ####
+if __name__=='__main__':
+    
+    urdf_fn = "/home/xueyi/diffsim/NeuS/rsc/mano/mano_mean_nocoll_simplified.urdf"
+    robot_agent = RobotAgent(urdf_fn)
+    
+    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).cuda()
+    cur_robot.set_initial_state(init_joint_states)
+    ##### Test the set initial state function #####
+    
+    
+    
+    
+    cur_zeros_actions = torch.zeros((60, ), dtype=torch.float32).cuda()
+    cur_ones_actions = torch.ones((60, ), dtype=torch.float32).cuda() #  * 100
+    
+    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)
+    

models/dyn_model_utils.py

@@ -0,0 +1,1369 @@
+
+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
+### 
+
+## calculate transformation to the frame ##
+
+## the joint idx ##
+##### th_cuda_idx #####
+
+# name to the main axis? # 
+
+# def get_body_name_to_main_axis()
+# another one is just getting the joint offset positions? 
+# and after the revolute transformation # # all revolute joint points ####
+def get_body_name_to_main_axis():
+    # negative y; positive x #
+    body_name_to_main_axis = {
+        "body2": -2, "body6": 1, "body10": 1, "body14": 1, "body17": 1
+    }
+    return body_name_to_main_axis ## get the body name to main axis ##
+
+## insert one 
+def plane_rotation_matrix_from_angle_xz(angle):
+    ## angle of 
+    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
+    )
+    # 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 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 ## 
+        # rot_mtx = np.stack(
+        #     [np.array([cos_ + u_xx * (1 - cos_), u_xy * (1. - cos_) + u_z * sin_, u_xz * (1. - cos_) - u_y * sin_], dtype=np.float32), 
+        #      np.array([u_xy * (1. - cos_) - u_z * sin_, cos_ + u_yy * (1. - cos_), u_yz * (1. - cos_) + u_x * sin_], dtype=np.float32), 
+        #      np.array([u_xz * (1. - cos_) + u_y * sin_, u_yz * (1. - cos_) - u_x * sin_, cos_ + u_zz * (1. - cos_)], dtype=np.float32)
+        #     ], axis=-1 ### np stack 
+        # ) ## a single
+        
+        # rot_mtx = torch.stack(
+        #     [
+        #         torch.tensor([cos_ + u_xx * (1 - cos_), u_xy * (1. - cos_) + u_z * sin_, u_xz * (1. - cos_) - u_y * sin_], dtype=torch.float32),
+        #         torch.tensor([u_xy * (1. - cos_) - u_z * sin_, cos_ + u_yy * (1. - cos_), u_yz * (1. - cos_) + u_x * sin_], dtype=torch.float32), 
+        #         torch.tensor([u_xz * (1. - cos_) + u_y * sin_, u_yz * (1. - cos_) - u_x * sin_, cos_ + u_zz * (1. - cos_)], dtype=torch.float32)
+        #     ], dim=-1 ## stack those torch tensors ##
+        # )
+        
+        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 ##
+        )
+        
+        # rot_mtx = torch.stack(
+        #     [
+                
+        #         torch.tensor([cos_ + u_xx * (1 - cos_), u_xy * (1. - cos_) + u_z * sin_, u_xz * (1. - cos_) - u_y * sin_], dtype=torch.float32),
+        #         torch.tensor([u_xy * (1. - cos_) - u_z * sin_, cos_ + u_yy * (1. - cos_), u_yz * (1. - cos_) + u_x * sin_], dtype=torch.float32), 
+        #         torch.tensor([u_xz * (1. - cos_) + u_y * sin_, u_yz * (1. - cos_) - u_x * sin_, cos_ + u_zz * (1. - cos_)], dtype=torch.float32)
+        #     ], dim=-1 ## stack those torch tensors ##
+        # )
+        
+        # rot_mtx_numpy = rot_mtx.to_numpy()
+        # rot_mtx_at_rot_mtx = rot_mtx @ rot_mtx.transpose()
+        # print(rot_mtx_at_rot_mtx)
+        return rot_mtx
+
+## joint name = "joint3" ##
+# <joint name="joint3" type="revolute" axis="0.000000 0.000000 -1.000000" pos="4.689700 -4.425000 0.000000" quat="1 0 0 0" frame="WORLD" damping="1e7"/>
+class Joint:
+    def __init__(self, name, joint_type, axis, pos, quat, frame, damping, args) -> None:
+        self.name = name
+        self.type = joint_type
+        self.axis = axis
+        self.pos = pos
+        self.quat = quat
+        self.frame = frame
+        self.damping = damping
+        
+        self.args = args
+        
+        #### TODO: the dimension of the state vector ? ####
+        # self.state = 0. ## parameter
+        self.state = nn.Parameter(
+            torch.zeros((1,), dtype=torch.float32, requires_grad=True).cuda(self.args.th_cuda_idx), requires_grad=True
+        )
+        # self.rot_mtx = np.eye(3, dtypes=np.float32)
+        # self.trans_vec = np.zeros((3,), dtype=np.float32) ## rot m
+        self.rot_mtx = nn.Parameter(torch.eye(n=3, dtype=torch.float32, requires_grad=True).cuda(self.args.th_cuda_idx), requires_grad=True)
+        self.trans_vec  = nn.Parameter(torch.zeros((3,), dtype=torch.float32,  requires_grad=True).cuda(self.args.th_cuda_idx), requires_grad=True)
+        # self.rot_mtx = np.eye(3, dtype=np.float32)
+        # self.trans_vec = np.zeros((3,), dtype=np.float32)
+        
+        self.axis_rot_mtx = torch.tensor(
+            [
+                [1, 0, 0], [0, -1, 0], [0, 0, -1]
+            ], dtype=torch.float32
+        ).cuda(self.args.th_cuda_idx)
+        
+        self.joint_idx = -1
+        
+        self.transformed_joint_pts = self.pos.clone()
+    
+    def print_grads(self, ):
+        print(f"rot_mtx: {self.rot_mtx.grad}")
+        print(f"trans_vec: {self.trans_vec.grad}")
+        
+    def clear_grads(self,):
+        if self.rot_mtx.grad is not None:
+            self.rot_mtx.grad.data = self.rot_mtx.grad.data * 0.
+        if self.trans_vec.grad is not None:
+            self.trans_vec.grad.data = self.trans_vec.grad.data * 0.
+        
+    def compute_transformation(self,):
+        # use the state to transform them # # transform # ## transform the state ##
+        # use the state to transform them # # transform them for the state #
+        if self.type == "revolute":
+            # print(f"computing transformation matrices with axis: {self.axis}, state: {self.state}")
+            # rotation matrix from the axis angle # 
+            rot_mtx = rotation_matrix_from_axis_angle(self.axis, self.state)
+            # rot_mtx(p - p_v) + p_v -> rot_mtx p - rot_mtx p_v + p_v
+            # trans_vec = self.pos - np.matmul(rot_mtx, self.pos.reshape(3, 1)).reshape(3)
+            # self.rot_mtx = np.copy(rot_mtx)
+            # self.trans_vec = np.copy(trans_vec)
+            trans_vec = self.pos - torch.matmul(rot_mtx, self.pos.view(3, 1)).view(3).contiguous()
+            self.rot_mtx = rot_mtx
+            self.trans_vec = trans_vec
+        else:
+            ### TODO: implement transformations for joints in other types ###
+            pass
+    
+    def set_state(self, name_to_state):
+        if self.name in name_to_state:
+            # self.state = name_to_state["name"]
+            self.state = name_to_state[self.name] ## 
+            
+    def set_state_via_vec(self, state_vec): ### transform points via the state vectors here ###
+        if self.joint_idx >= 0:
+            self.state = state_vec[self.joint_idx] ## give the parameter to the parameters ##
+            
+    def set_joint_idx(self, joint_name_to_idx):
+        if self.name in joint_name_to_idx:
+            self.joint_idx = joint_name_to_idx[self.name]
+            
+    
+    def set_args(self, args):
+        self.args = args
+        
+        
+    def compute_transformation_via_state_vals(self, state_vals):
+        if self.joint_idx >= 0:
+            cur_joint_state = state_vals[self.joint_idx]
+        else:
+            cur_joint_state = self.state
+        # use the state to transform them # # transform # ## transform the state ##
+        # use the state to transform them # # transform them for the state #
+        if self.type == "revolute":
+            # print(f"computing transformation matrices with axis: {self.axis}, state: {self.state}")
+            # rotation matrix from the axis angle # 
+            rot_mtx = rotation_matrix_from_axis_angle(self.axis, cur_joint_state)
+            # rot_mtx(p - p_v) + p_v -> rot_mtx p - rot_mtx p_v + p_v
+            # trans_vec = self.pos - np.matmul(rot_mtx, self.pos.reshape(3, 1)).reshape(3)
+            # self.rot_mtx = np.copy(rot_mtx)
+            # self.trans_vec = np.copy(trans_vec)
+            trans_vec = self.pos - torch.matmul(rot_mtx, self.pos.view(3, 1)).view(3).contiguous()
+            self.rot_mtx = rot_mtx
+            self.trans_vec = trans_vec
+        elif self.type == "free2d":
+            cur_joint_state = state_vals # still only for the current scene #
+            # cur_joint_state
+            cur_joint_rot_val = state_vals[2]
+            ### rot_mtx ### ### rot_mtx ###
+            rot_mtx = plane_rotation_matrix_from_angle_xz(cur_joint_rot_val)
+            # rot_mtx = plane_rotation_matrix_from_angle(cur_joint_rot_val) ### 2 x 2 rot matrix # 
+            # cur joint rot val #
+            # rot mtx of the rotation 
+            # xy_val = 
+            # axis_rot_mtx
+            # R_axis^T ( R R_axis (p) + trans (with the y-axis padded) )
+            cur_trans_vec = torch.stack(
+                [state_vals[0], torch.zeros_like(state_vals[0]), state_vals[1]], dim=0
+            )
+            # cur_trans_vec # 
+            rot_mtx = torch.matmul(self.axis_rot_mtx.transpose(1, 0), torch.matmul(rot_mtx, self.axis_rot_mtx))
+            trans_vec = torch.matmul(self.axis_rot_mtx.transpose(1, 0), cur_trans_vec.unsqueeze(-1).contiguous()).squeeze(-1).contiguous() + self.pos
+            
+            self.rot_mtx = rot_mtx
+            self.trans_vec = trans_vec ## rot_mtx and trans_vec # 
+        else:
+            ### TODO: implement transformations for joints in other types ###
+            pass
+        return self.rot_mtx, self.trans_vec
+
+    def transform_joints_via_parent_rot_trans_infos(self, parent_rot_mtx, parent_trans_vec):
+        # 
+        # if self.type == "revolute" or self.type == "free2d":
+        transformed_joint_pts = torch.matmul(parent_rot_mtx, self.pos.view(3 ,1).contiguous()).view(3).contiguous() + parent_trans_vec
+        # else:
+        self.transformed_joint_pts = transformed_joint_pts ### get self transformed joint pts here ###
+        return transformed_joint_pts
+        # if self.joint_idx >= 0:
+        #     cur_joint_state = state_vals[self.joint_idx]
+        # else:
+        #     cur_joint_state = self.state # state #
+        # # use the state to transform them # # transform ### transform the state ##
+        # # use the state to transform them # # transform them for the state # transform for the state #
+        # if self.type == "revolute":
+        #     # print(f"computing transformation matrices with axis: {self.axis}, state: {self.state}")
+        #     # rotation matrix from the axis angle # 
+        #     rot_mtx = rotation_matrix_from_axis_angle(self.axis, cur_joint_state)
+        #     # rot_mtx(p - p_v) + p_v -> rot_mtx p - rot_mtx p_v + p_v
+        #     # trans_vec = self.pos - np.matmul(rot_mtx, self.pos.reshape(3, 1)).reshape(3)
+        #     # self.rot_mtx = np.copy(rot_mtx)
+        #     # self.trans_vec = np.copy(trans_vec)
+        #     trans_vec = self.pos - torch.matmul(rot_mtx, self.pos.view(3, 1)).view(3).contiguous()
+        #     self.rot_mtx = rot_mtx
+        #     self.trans_vec = trans_vec
+        # elif self.type == "free2d":
+        #     cur_joint_state = state_vals # still only for the current scene #
+        #     # cur_joint_state
+        #     cur_joint_rot_val = state_vals[2]
+        #     ### rot_mtx ### ### rot_mtx ###
+        #     rot_mtx = plane_rotation_matrix_from_angle_xz(cur_joint_rot_val)
+        #     # rot_mtx = plane_rotation_matrix_from_angle(cur_joint_rot_val) ### 2 x 2 rot matrix # 
+        #     # cur joint rot val #
+        #     # rot mtx of the rotation 
+        #     # xy_val = 
+        #     # axis_rot_mtx
+        #     # R_axis^T ( R R_axis (p) + trans (with the y-axis padded) )
+        #     cur_trans_vec = torch.stack(
+        #         [state_vals[0], torch.zeros_like(state_vals[0]), state_vals[1]], dim=0
+        #     )
+        #     # cur_trans_vec # 
+        #     rot_mtx = torch.matmul(self.axis_rot_mtx.transpose(1, 0), torch.matmul(rot_mtx, self.axis_rot_mtx))
+        #     trans_vec = torch.matmul(self.axis_rot_mtx.transpose(1, 0), cur_trans_vec.unsqueeze(-1).contiguous()).squeeze(-1).contiguous() + self.pos
+            
+        #     self.rot_mtx = rot_mtx
+        #     self.trans_vec = trans_vec ## rot_mtx and trans_vec # 
+        # else:
+        #     ### TODO: implement transformations for joints in other types ###
+        #     pass
+        # return self.rot_mtx, self.trans_vec
+
+## fixed joint -> translation and rotation ##
+## revolute joint -> can be actuated ##
+## set states and compute the transfromations in a top-to-down manner ##
+
+## trnasform the robot -> a list of qs ##
+## a list of qs ##
+## transform from the root of the robot; pass qs from the root to the leaf node ##
+## visual meshes  or visual meshes from the basic description of robots ##
+## visual meshes; or visual points ##
+## visual meshes -> transform them into the visual density values here ##
+## visual meshes -> transform them into the ## into the visual counterparts ##
+## ## visual meshes -> ## ## ## 
+# <body name="body0" type="mesh"  filename="hand/body0.obj"  pos="0 0 0"  quat="1 0 0 0"  transform_type="OBJ_TO_WORLD"  density="1"  mu="0" rgba="0.700000 0.700000 0.700000 1"/>
+class Body:
+    def __init__(self, name, body_type, filename, pos, quat, transform_type, density, mu, rgba, radius, args) -> None:
+        self.name = name
+        self.body_type = body_type
+        ### for mesh object ###
+        self.filename = filename
+        self.args = args
+
+        self.pos = pos
+        self.quat = quat
+        self.transform_type = transform_type
+        self.density = density
+        self.mu = mu
+        self.rgba = rgba
+        
+        ### for sphere object ###
+        self.radius = radius
+        ## or vertices here #
+        ## pass them to the child and treat them as the parent transformation ##
+        
+        self.visual_pts_ref = None
+        self.visual_faces_ref = None
+        
+        self.visual_pts = None ## visual pts and 
+        
+        self.body_name_to_main_axis = get_body_name_to_main_axis() ### get the body name to main axis here #
+        
+        self.get_visual_counterparts()
+        
+        
+    def update_radius(self,):
+        self.radius.data = self.radius.data - self.radius.grad.data
+        
+        self.radius.grad.data = self.radius.grad.data * 0.
+        
+    
+    def update_xml_file(self,):
+        xml_content_with_flexible_radius = f"""<redmax model="hand">
+<!-- 1) change the damping value here? -->
+<!-- 2) change the center of mass -->
+    <option integrator="BDF2" timestep="0.01" gravity="0. 0. -0.000098"/>
+	<ground pos="0 0 -10" normal="0 0 1"/>
+	<default>
+        <general_primitive_contact kn="1e6" kt="1e3" mu="0.8" damping="3e1" />
+    </default>
+
+    <robot>
+        <link name="link0">
+            <joint name="joint0" type="fixed" pos="0 0 0" quat="1 0 0 0" frame="WORLD"/>
+            <body name="body0" type="mesh"  filename="hand/body0.obj"  pos="0 0 0"  quat="1 0 0 0"  transform_type="OBJ_TO_WORLD"  density="1"  mu="0" rgba="0.700000 0.700000 0.700000 1"/>
+            <link name="link1">
+                <joint name="joint1" type="revolute" axis="0.000000 0.000000 -1.000000" pos="-3.300000 -5.689700 0.000000" quat="1 0 0 0" frame="WORLD" damping="1e4"/>
+                <body name="body1" type="mesh"  filename="hand/body1.obj"  pos="0 0 0"  quat="1 0 0 0"  transform_type="OBJ_TO_WORLD"  density="1"  mu="0" rgba="0.600000 0.600000 0.600000 1"/>
+                <link name="link2">
+                    <joint name="joint2" type="revolute" axis="1.000000 0.000000 0.000000" pos="-3.300000 -7.680000 0.000000" quat="1 0 0 0" frame="WORLD" damping="1e4"/>
+                    <body name="body2" type="mesh"  filename="hand/body2.obj"  pos="0 0 0"  quat="1 0 0 0"  transform_type="OBJ_TO_WORLD"  density="1"  mu="0" rgba="0.500000 0.500000 0.500000 1"/>
+                </link>
+            </link>
+            <link name="link3">
+                <!-- revolute joint -->
+                <joint name="joint3" type="revolute" axis="0.000000 0.000000 -1.000000" pos="4.689700 -4.425000 0.000000" quat="1 0 0 0" frame="WORLD" damping="1e4"/>
+                <body name="body3" type="mesh"  filename="hand/body3.obj"  pos="0 0 0"  quat="1 0 0 0"  transform_type="OBJ_TO_WORLD"  density="1"  mu="0" rgba="0.600000 0.600000 0.600000 1"/>
+                <link name="link4">
+                    <joint name="joint4" type="revolute" axis="0.000000 1.000000 0.000000" pos="6.680000 -4.425000 0.000000" quat="1 0 0 0" frame="WORLD" damping="1e4"/>
+                    <body name="body4" type="mesh"  filename="hand/body4.obj"  pos="0 0 0"  quat="1 0 0 0"  transform_type="OBJ_TO_WORLD"  density="1"  mu="0" rgba="0.500000 0.500000 0.500000 1"/>
+                        <link name="link5">
+                            <joint name="joint5" type="revolute" axis="0.000000 1.000000 0.000000" pos="11.080000 -4.425000 0.000000" quat="1 0 0 0" frame="WORLD" damping="1e4"/>
+                            <body name="body5" type="mesh"  filename="hand/body5.obj"  pos="0 0 0"  quat="1 0 0 0"  transform_type="OBJ_TO_WORLD"  density="1"  mu="0" rgba="0.400000 0.400000 0.400000 1"/>
+                                <link name="link6">
+                                    <joint name="joint6" type="revolute" axis="0.000000 1.000000 0.000000" pos="15.480000 -4.425000 0.000000" quat="1 0 0 0" frame="WORLD" damping="1e4"/>
+                                    <body name="body6" type="mesh"  filename="hand/body6.obj"  pos="0 0 0"  quat="1 0 0 0"  transform_type="OBJ_TO_WORLD"  density="1"  mu="0" rgba="0.300000 0.300000 0.300000 1"/>
+                                </link>
+                        </link>
+                </link>
+            </link>
+            <link name="link7">
+                <joint name="joint7" type="revolute" axis="0.000000 0.000000 -1.000000" pos="4.689700 -1.475000 0.000000" quat="1 0 0 0" frame="WORLD" damping="1e4"/>
+                <body name="body7" type="mesh"  filename="hand/body7.obj"  pos="0 0 0"  quat="1 0 0 0"  transform_type="OBJ_TO_WORLD"  density="1"  mu="0" rgba="0.600000 0.600000 0.600000 1"/>
+                <link name="link8">
+                    <joint name="joint8" type="revolute" axis="0.000000 1.000000 0.000000" pos="6.680000 -1.475000 0.000000" quat="1 0 0 0" frame="WORLD" damping="1e4"/>
+                    <body name="body8" type="mesh"  filename="hand/body8.obj"  pos="0 0 0"  quat="1 0 0 0"  transform_type="OBJ_TO_WORLD"  density="1"  mu="0" rgba="0.500000 0.500000 0.500000 1"/>
+                        <link name="link9">
+                            <joint name="joint9" type="revolute" axis="0.000000 1.000000 0.000000" pos="11.080000 -1.475000 0.000000" quat="1 0 0 0" frame="WORLD" damping="1e4"/>
+                            <body name="body9" type="mesh"  filename="hand/body9.obj"  pos="0 0 0"  quat="1 0 0 0"  transform_type="OBJ_TO_WORLD"  density="1"  mu="0" rgba="0.400000 0.400000 0.400000 1"/>
+                                <link name="link10">
+                                    <joint name="joint10" type="revolute" axis="0.000000 1.000000 0.000000" pos="15.480000 -1.475000 0.000000" quat="1 0 0 0" frame="WORLD" damping="1e4"/>
+                                    <body name="body10" type="mesh"  filename="hand/body10.obj"  pos="0 0 0"  quat="1 0 0 0"  transform_type="OBJ_TO_WORLD"  density="1"  mu="0" rgba="0.300000 0.300000 0.300000 1"/>
+                                </link>
+                        </link>
+                </link>
+            </link>
+            <link name="link11">
+                <joint name="joint11" type="revolute" axis="0.000000 0.000000 -1.000000" pos="4.689700 1.475000 0.000000" quat="1 0 0 0" frame="WORLD" damping="1e4"/>
+                <body name="body11" type="mesh"  filename="hand/body11.obj"  pos="0 0 0"  quat="1 0 0 0"  transform_type="OBJ_TO_WORLD"  density="1"  mu="0" rgba="0.600000 0.600000 0.600000 1"/>
+                <link name="link12">
+                    <joint name="joint12" type="revolute" axis="0.000000 1.000000 0.000000" pos="6.680000 1.475000 0.000000" quat="1 0 0 0" frame="WORLD" damping="1e4"/>
+                    <body name="body12" type="mesh"  filename="hand/body12.obj"  pos="0 0 0"  quat="1 0 0 0"  transform_type="OBJ_TO_WORLD"  density="1"  mu="0" rgba="0.500000 0.500000 0.500000 1"/>
+                        <link name="link13">
+                            <joint name="joint13" type="revolute" axis="0.000000 1.000000 0.000000" pos="11.080000 1.475000 0.000000" quat="1 0 0 0" frame="WORLD" damping="1e4"/>
+                            <body name="body13" type="mesh"  filename="hand/body13.obj"  pos="0 0 0"  quat="1 0 0 0"  transform_type="OBJ_TO_WORLD"  density="1"  mu="0" rgba="0.400000 0.400000 0.400000 1"/>
+                                <link name="link14">
+                                    <joint name="joint14" type="revolute" axis="0.000000 1.000000 0.000000" pos="15.480000 1.475000 0.000000" quat="1 0 0 0" frame="WORLD" damping="1e4"/>
+                                    <body name="body14" type="mesh"  filename="hand/body14.obj"  pos="0 0 0"  quat="1 0 0 0"  transform_type="OBJ_TO_WORLD"  density="1"  mu="0" rgba="0.300000 0.300000 0.300000 1"/>
+                                </link>
+                        </link>
+                </link>
+            </link>
+            <link name="link15">
+                <joint name="joint15" type="revolute" axis="0.000000 0.000000 -1.000000" pos="4.689700 4.425000 0.000000" quat="1 0 0 0" frame="WORLD" damping="1e4"/>
+                <body name="body15" type="mesh"  filename="hand/body15.obj"  pos="0 0 0"  quat="1 0 0 0"  transform_type="OBJ_TO_WORLD"  density="1"  mu="0" rgba="0.600000 0.600000 0.600000 1"/>
+                <link name="link16">
+                    <joint name="joint16" type="revolute" axis="0.000000 1.000000 0.000000" pos="6.680000 4.425000 0.000000" quat="1 0 0 0" frame="WORLD" damping="1e4"/>
+                    <body name="body16" type="mesh"  filename="hand/body16.obj"  pos="0 0 0"  quat="1 0 0 0"  transform_type="OBJ_TO_WORLD"  density="1"  mu="0" rgba="0.500000 0.500000 0.500000 1"/>
+                        <link name="link17">
+                            <joint name="joint17" type="revolute" axis="0.000000 1.000000 0.000000" pos="11.080000 4.425000 0.000000" quat="1 0 0 0" frame="WORLD" damping="1e4"/>
+                            <body name="body17" type="mesh"  filename="hand/body17.obj"  pos="0 0 0"  quat="1 0 0 0"  transform_type="OBJ_TO_WORLD"  density="1"  mu="0" rgba="0.400000 0.400000 0.400000 1"/>
+                        </link>
+                </link>
+            </link>
+        </link>
+    </robot>
+	
+	<robot>
+		<link name="sphere">
+			<joint name="sphere" type="free2d" pos = "10. 0.0 3.5" quat="1 -1 0 0" format="LOCAL" damping="0"/>
+			<body name="sphere" type="sphere" radius="{self.radius[0].detach().cpu().item()}" pos="0 0 0" quat="1 0 0 0" density="0.5" mu="0" texture="resources/textures/sphere.jpg"/>
+		</link>
+	</robot>
+	
+	<contact>
+		<ground_contact body="sphere" kn="1e6" kt="1e3" mu="0.8" damping="3e1"/>
+		<general_primitive_contact general_body="body0" primitive_body="sphere"/>
+		<general_primitive_contact general_body="body1" primitive_body="sphere"/>
+		<general_primitive_contact general_body="body2" primitive_body="sphere"/>
+		<general_primitive_contact general_body="body3" primitive_body="sphere"/>
+		<general_primitive_contact general_body="body4" primitive_body="sphere"/>
+		<general_primitive_contact general_body="body5" primitive_body="sphere"/>
+		<general_primitive_contact general_body="body6" primitive_body="sphere"/>
+		<general_primitive_contact general_body="body7" primitive_body="sphere"/>
+		<general_primitive_contact general_body="body8" primitive_body="sphere"/>
+		<general_primitive_contact general_body="body9" primitive_body="sphere"/>
+		<general_primitive_contact general_body="body10" primitive_body="sphere"/>
+		<general_primitive_contact general_body="body11" primitive_body="sphere"/>
+		<general_primitive_contact general_body="body12" primitive_body="sphere"/>
+		<general_primitive_contact general_body="body13" primitive_body="sphere"/>
+		<general_primitive_contact general_body="body14" primitive_body="sphere"/>
+		<general_primitive_contact general_body="body15" primitive_body="sphere"/>
+		<general_primitive_contact general_body="body16" primitive_body="sphere"/>
+		<general_primitive_contact general_body="body17" primitive_body="sphere"/>
+	</contact>
+
+    <actuator>
+        <motor joint="joint1"  ctrl="force"  ctrl_range="-3e5 3e5"/>
+        <motor joint="joint2"  ctrl="force"  ctrl_range="-3e5 3e5"/>
+        <motor joint="joint3"  ctrl="force"  ctrl_range="-3e5 3e5"/>
+        <motor joint="joint4"  ctrl="force"  ctrl_range="-3e5 3e5"/>
+        <motor joint="joint5"  ctrl="force"  ctrl_range="-3e5 3e5"/>
+        <motor joint="joint6"  ctrl="force"  ctrl_range="-3e5 3e5"/>
+        <motor joint="joint7"  ctrl="force"  ctrl_range="-3e5 3e5"/>
+        <motor joint="joint8"  ctrl="force"  ctrl_range="-3e5 3e5"/>
+        <motor joint="joint9"  ctrl="force"  ctrl_range="-3e5 3e5"/>
+        <motor joint="joint10"  ctrl="force"  ctrl_range="-3e5 3e5"/>
+        <motor joint="joint11"  ctrl="force"  ctrl_range="-3e5 3e5"/>
+        <motor joint="joint12"  ctrl="force"  ctrl_range="-3e5 3e5"/>
+        <motor joint="joint13"  ctrl="force"  ctrl_range="-3e5 3e5"/>
+        <motor joint="joint14"  ctrl="force"  ctrl_range="-3e5 3e5"/>
+        <motor joint="joint15"  ctrl="force"  ctrl_range="-3e5 3e5"/>
+        <motor joint="joint16"  ctrl="force"  ctrl_range="-3e5 3e5"/>
+        <motor joint="joint17"  ctrl="force"  ctrl_range="-3e5 3e5"/>
+    </actuator>
+</redmax>
+"""
+        xml_loading_fn = "/home/xueyi/diffsim/DiffHand/assets/hand_sphere_free_sphere_geo_test.xml"
+        with open(xml_loading_fn, "w") as wf:
+            wf.write(xml_content_with_flexible_radius)
+            wf.close()
+            
+    ### get visual pts colorrs ### ### 
+    def get_visual_pts_colors(self, ):
+        tot_visual_pts_nn = self.visual_pts_ref.size(0)
+        # self.pts_rgba = [torch.from_numpy(self.rgba).float().cuda(self.args.th_cuda_idx) for _ in range(tot_visual_pts_nn)] # total visual pts nn 
+        self.pts_rgba = [torch.tensor(self.rgba.data).cuda(self.args.th_cuda_idx) for _ in range(tot_visual_pts_nn)] # total visual pts nn  skeletong
+        self.pts_rgba = torch.stack(self.pts_rgba, dim=0) # 
+        return self.pts_rgba
+    
+    def get_visual_counterparts(self,):
+        ### TODO: implement this for visual counterparts ### mid line regression and name to body mapping relations --- for each body, how to calculate the midline and other properties?
+        ######## get body type ########## get visual midline of the input mesh and the mesh vertices? ######## # skeleton of the hand -> 21 points ? retarget from this hand to the mano hand and use the mano hand priors?
+        if self.body_type == "sphere":
+            filename = "/home/xueyi/diffsim/DiffHand/examples/save_res/hand_sphere_demo/meshes/18.obj"
+            if not os.path.exists(filename):
+                filename = "/data/xueyi/diffsim/DiffHand/assets/18.obj"
+        else:
+            filename = self.filename
+            rt_asset_path = "/home/xueyi/diffsim/DiffHand/assets" ### assets folder ###
+            if not os.path.exists(rt_asset_path):
+                rt_asset_path = "/data/xueyi/diffsim/DiffHand/assets"
+            filename = os.path.join(rt_asset_path, filename)
+        body_mesh = trimesh.load(filename, process=False) 
+        # verts = np.array(body_mesh.vertices)
+        # faces = np.array(body_mesh.faces, dtype=np.long)
+        
+        # self.visual_pts_ref = np.copy(verts) ## verts ##
+        # self.visual_faces_ref = np.copy(faces) ## faces 
+        # self.visual_pts_ref # 
+        
+        #### body_mesh.vertices ####
+        # verts = torch.tensor(body_mesh.vertices, dtype=torch.float32).cuda(self.args.th_cuda_idx)
+        # faces = torch.tensor(body_mesh.faces, dtype=torch.long).cuda(self.args.th_cuda_idx)
+        #### body_mesh.vertices ####
+        
+        # self.pos = nn.Parameter(
+        #     torch.tensor([0., 0., 0.], dtype=torch.float32,  requires_grad=True).cuda(self.args.th_cuda_idx), requires_grad=True
+        # )
+        
+        self.pos = nn.Parameter(
+            torch.tensor(self.pos.detach().cpu().tolist(), dtype=torch.float32,  requires_grad=True).cuda(self.args.th_cuda_idx), requires_grad=True
+        )
+        
+        ### Step 1 ### -> set the pos to the correct initial pose ###
+        
+        self.radius = nn.Parameter(
+            torch.tensor([self.args.initial_radius], dtype=torch.float32, requires_grad=True).cuda(self.args.th_cuda_idx), requires_grad=True
+        )
+        ### visual pts ref ### ## body_mesh.vertices -> # 
+        self.visual_pts_ref = torch.tensor(body_mesh.vertices, dtype=torch.float32).cuda(self.args.th_cuda_idx)
+        
+        # if self.name == "sphere":
+        #     self.visual_pts_ref = self.visual_pts_ref  / 2. #  the initial radius
+        #     self.visual_pts_ref = self.visual_pts_ref * self.radius ## multiple the initla radius # 
+        
+        # self.visual_pts_ref = nn.Parameter(
+        #     torch.tensor(body_mesh.vertices, dtype=torch.float32, requires_grad=True).cuda(self.args.th_cuda_idx), requires_grad=True
+        # )
+        # self.visual_faces_ref = nn.Parameter(
+        #     torch.tensor(body_mesh.faces, dtype=torch.long, requires_grad=True).cuda(self.args.th_cuda_idx), requires_grad=True
+        # )
+        self.visual_faces_ref = torch.tensor(body_mesh.faces, dtype=torch.long).cuda(self.args.th_cuda_idx)
+        
+        # body_name_to_main_axis
+        # body_name_to_main_axis for the body_name_to_main_axis # 
+        # visual_faces_ref # 
+        # visual_pts_ref #
+        
+        minn_pts, _ = torch.min(self.visual_pts_ref, dim=0) ### get the visual pts minn ###
+        maxx_pts, _ = torch.max(self.visual_pts_ref, dim=0) ### visual pts maxx ###
+        mean_pts = torch.mean(self.visual_pts_ref, dim=0) ### mean_pts of the mean_pts ###
+        
+        if self.name in self.body_name_to_main_axis:
+            cur_main_axis = self.body_name_to_main_axis[self.name] ## get the body name ##
+            
+            if cur_main_axis == -2:
+                main_axis_pts = minn_pts[1] # the main axis pts
+                full_main_axis_pts = torch.tensor([mean_pts[0], main_axis_pts, mean_pts[2]], dtype=torch.float32).cuda(self.args.th_cuda_idx)
+            elif cur_main_axis == 1:
+                main_axis_pts = maxx_pts[0] # the maxx axis pts
+                full_main_axis_pts = torch.tensor([main_axis_pts, mean_pts[1], mean_pts[2]], dtype=torch.float32).cuda(self.args.th_cuda_idx)
+            self.full_main_axis_pts_ref = full_main_axis_pts
+        else:
+            self.full_main_axis_pts_ref = mean_pts.clone() ### get the mean pts ###
+        # mean_pts
+        # main_axis_pts = 
+            
+        
+        # self.visual_pts_ref = verts #
+        # self.visual_faces_ref = faces #
+        # get visual points colors # the color should be an optimizable property # #  # or init visual point colors here ## # or init visual point colors here #
+        # simulatoable assets ## for the
+        
+    def transform_visual_pts_ref(self,):
+        if self.name == "sphere":
+            visual_pts_ref =  self.visual_pts_ref  / 2. #
+            visual_pts_ref = visual_pts_ref * self.radius
+        else:
+            visual_pts_ref = self.visual_pts_ref
+        return visual_pts_ref
+
+    def transform_visual_pts(self, rot_mtx, trans_vec):
+        visual_pts_ref = self.transform_visual_pts_ref()
+        # rot_mtx: 3 x 3 numpy array 
+        # trans_vec: 3 numpy array
+        # print(f"transforming body with rot_mtx: {rot_mtx} and trans_vec: {trans_vec}")
+        # self.visual_pts = np.matmul(rot_mtx, self.visual_pts_ref.T).T + trans_vec.reshape(1, 3) # reshape #
+        # print(f"rot_mtx: {rot_mtx}, trans_vec: {trans_vec}")
+        self.visual_pts = torch.matmul(rot_mtx, visual_pts_ref.transpose(1, 0)).transpose(1, 0) + trans_vec.unsqueeze(0)
+        
+        # full_main_axis_pts -> 
+        self.full_main_axis_pts = torch.matmul(rot_mtx, self.full_main_axis_pts_ref.unsqueeze(-1)).contiguous().squeeze(-1) + trans_vec
+        self.full_main_axis_pts = self.full_main_axis_pts.unsqueeze(0)
+        
+        return self.visual_pts
+    
+    def get_tot_transformed_joints(self, transformed_joints):
+        if self.name in self.body_name_to_main_axis:
+            transformed_joints.append(self.full_main_axis_pts)
+        return transformed_joints
+    
+    def get_nn_pts(self,):
+        self.nn_pts = self.visual_pts_ref.size(0)
+        return self.nn_pts
+    
+    def set_args(self, args):
+        self.args = args
+        
+    def clear_grad(self, ):
+        if self.pos.grad is not None:
+            self.pos.grad.data = self.pos.grad.data  * 0.
+        if self.radius.grad is not None:
+            self.radius.grad.data = self.radius.grad.data  * 0.
+    
+    
+    # get the visual counterparts of the boyd mesh or elements #
+    
+    # xyz attribute ## ## xyz attribute #
+
+# use get_name_to_visual_pts
+# use get_name_to_visual_pts_faces to get the transformed visual pts and faces #
+class Link: 
+    def __init__(self, name, joint: Joint, body: Body, children, args) -> None:
+        
+        self.joint = joint
+        self.body = body
+        self.children = children
+        self.name = name
+        
+        self.args = args
+        
+        # joint # parent_rot_mtx, parent_trans_vec
+        self.parent_rot_mtx = nn.Parameter(torch.eye(n=3, dtype=torch.float32).cuda(self.args.th_cuda_idx), requires_grad=True)
+        self.parent_trans_vec = nn.Parameter(torch.zeros((3,), dtype=torch.float32).cuda(self.args.th_cuda_idx), requires_grad=True)
+        self.curr_rot_mtx = nn.Parameter(torch.eye(n=3, dtype=torch.float32).cuda(self.args.th_cuda_idx), requires_grad=True)
+        self.curr_trans_vec = nn.Parameter(torch.zeros((3,), dtype=torch.float32).cuda(self.args.th_cuda_idx), requires_grad=True)
+        # 
+        self.tot_rot_mtx = nn.Parameter(torch.eye(n=3, dtype=torch.float32).cuda(self.args.th_cuda_idx), requires_grad=True)
+        self.tot_trans_vec = nn.Parameter(torch.zeros((3,), dtype=torch.float32).cuda(self.args.th_cuda_idx), requires_grad=True) ## torch zeros #
+    
+    def print_grads(self, ): ### print grads here ###
+        print(f"parent_rot_mtx: {self.parent_rot_mtx.grad}")
+        print(f"parent_trans_vec: {self.parent_trans_vec.grad}")
+        print(f"curr_rot_mtx: {self.curr_rot_mtx.grad}")
+        print(f"curr_trans_vec: {self.curr_trans_vec.grad}")
+        print(f"tot_rot_mtx: {self.tot_rot_mtx.grad}")
+        print(f"tot_trans_vec: {self.tot_trans_vec.grad}")
+        print(f"Joint")
+        self.joint.print_grads()
+        for cur_link in self.children:
+            cur_link.print_grads()
+        
+    
+    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)
+        # if self.joint_idx >= 0:
+        #     self.state = state_vec[self.joint_idx]
+    
+    ## 
+    def get_tot_transformed_joints(self, transformed_joints):
+        cur_joint_transformed_pts = self.joint.transformed_joint_pts.unsqueeze(0) ### 3 pts 
+        transformed_joints.append(cur_joint_transformed_pts)
+        transformed_joints = self.body.get_tot_transformed_joints(transformed_joints)
+        # if self.joint.name 
+        for cur_link in self.children:
+            transformed_joints = cur_link.get_tot_transformed_joints(transformed_joints)
+        return transformed_joints
+        
+    def compute_transformation_via_state_vecs(self, state_vals, parent_rot_mtx, parent_trans_vec, visual_pts_list):
+        # state vecs and rot mtx # state vecs #####
+        joint_rot_mtx, joint_trans_vec = self.joint.compute_transformation_via_state_vals(state_vals=state_vals)
+
+        self.curr_rot_mtx = joint_rot_mtx
+        self.curr_trans_vec = joint_trans_vec
+        
+        self.joint.transform_joints_via_parent_rot_trans_infos(parent_rot_mtx=parent_rot_mtx, parent_trans_vec=parent_trans_vec) ## get rot and trans mtx and vecs ### 
+
+        tot_parent_rot_mtx = torch.matmul(parent_rot_mtx, joint_rot_mtx)
+        tot_parent_trans_vec = torch.matmul(parent_rot_mtx, joint_trans_vec.unsqueeze(-1)).view(3) + parent_trans_vec
+        
+        self.tot_rot_mtx = tot_parent_rot_mtx
+        self.tot_trans_vec = tot_parent_trans_vec
+        
+        # self.tot_rot_mtx = np.copy(tot_parent_rot_mtx)
+        # self.tot_trans_vec = np.copy(tot_parent_trans_vec)
+        
+        ### visual_pts_list for recording visual pts ###
+        
+        cur_body_visual_pts = self.body.transform_visual_pts(rot_mtx=self.tot_rot_mtx, trans_vec=self.tot_trans_vec)
+        visual_pts_list.append(cur_body_visual_pts)
+        
+        for cur_link in self.children:
+            # cur_link.parent_rot_mtx = np.copy(tot_parent_rot_mtx) ### set children parent rot mtx and the trans vec
+            # cur_link.parent_trans_vec = np.copy(tot_parent_trans_vec) ## 
+            cur_link.parent_rot_mtx = tot_parent_rot_mtx  ### set children parent rot mtx and the trans vec #
+            cur_link.parent_trans_vec = tot_parent_trans_vec  ## 
+            # cur_link.compute_transformation() ## compute self's transformations
+            cur_link.compute_transformation_via_state_vecs(state_vals, tot_parent_rot_mtx, tot_parent_trans_vec, visual_pts_list)
+    
+    def get_visual_pts_rgba_values(self, pts_rgba_vals_list):
+        
+        cur_body_visual_rgba_vals = self.body.get_visual_pts_colors()
+        pts_rgba_vals_list.append(cur_body_visual_rgba_vals)
+        
+        for cur_link in self.children:
+            cur_link.get_visual_pts_rgba_values(pts_rgba_vals_list)
+        
+        
+    
+    def compute_transformation(self,):
+        self.joint.compute_transformation()
+        # self.curr_rot_mtx = np.copy(self.joint.rot_mtx)
+        # self.curr_trans_vec = np.copy(self.joint.trans_vec)
+        
+        self.curr_rot_mtx = self.joint.rot_mtx
+        self.curr_trans_vec = self.joint.trans_vec
+        # rot_p (rot_c p + trans_c) + trans_p # 
+        # rot_p rot_c p + rot_p trans_c + trans_p #
+        #### matmul ####
+        # tot_parent_rot_mtx = np.matmul(self.parent_rot_mtx, self.curr_rot_mtx)
+        # tot_parent_trans_vec = np.matmul(self.parent_rot_mtx, self.curr_trans_vec.reshape(3, 1)).reshape(3) + self.parent_trans_vec
+        
+        tot_parent_rot_mtx = torch.matmul(self.parent_rot_mtx, self.curr_rot_mtx)
+        tot_parent_trans_vec = torch.matmul(self.parent_rot_mtx, self.curr_trans_vec.unsqueeze(-1)).view(3) + self.parent_trans_vec
+        
+        self.tot_rot_mtx = tot_parent_rot_mtx
+        self.tot_trans_vec = tot_parent_trans_vec
+        
+        # self.tot_rot_mtx = np.copy(tot_parent_rot_mtx)
+        # self.tot_trans_vec = np.copy(tot_parent_trans_vec)
+        
+        for cur_link in self.children:
+            # cur_link.parent_rot_mtx = np.copy(tot_parent_rot_mtx) ### set children parent rot mtx and the trans vec
+            # cur_link.parent_trans_vec = np.copy(tot_parent_trans_vec) ## 
+            cur_link.parent_rot_mtx = tot_parent_rot_mtx  ### set children parent rot mtx and the trans vec #
+            cur_link.parent_trans_vec = tot_parent_trans_vec  ## 
+            cur_link.compute_transformation() ## compute self's transformations
+            
+    def get_name_to_visual_pts_faces(self, name_to_visual_pts_faces):
+        # transform_visual_pts # ## rot_mt
+        self.body.transform_visual_pts(rot_mtx=self.tot_rot_mtx, trans_vec=self.tot_trans_vec)
+        name_to_visual_pts_faces[self.body.name] = {"pts": self.body.visual_pts, "faces": self.body.visual_faces_ref}
+        for cur_link in self.children:
+            cur_link.get_name_to_visual_pts_faces(name_to_visual_pts_faces) ## transform the pts faces 
+
+    def get_visual_pts_list(self, visual_pts_list):
+        # transform_visual_pts # ## rot_mt
+        self.body.transform_visual_pts(rot_mtx=self.tot_rot_mtx, trans_vec=self.tot_trans_vec)
+        visual_pts_list.append(self.body.visual_pts) # body template #
+        # name_to_visual_pts_faces[self.body.name] = {"pts": self.body.visual_pts, "faces": self.body.visual_faces_ref}
+        for cur_link in self.children:
+            # cur_link.get_name_to_visual_pts_faces(name_to_visual_pts_faces) ## transform the pts faces 
+            cur_link.get_visual_pts_list(visual_pts_list)
+            
+    
+        
+    def set_joint_idx(self, joint_name_to_idx):
+        self.joint.set_joint_idx(joint_name_to_idx)
+        for cur_link in self.children:
+            cur_link.set_joint_idx(joint_name_to_idx)
+        # if self.name in joint_name_to_idx:
+        #     self.joint_idx = joint_name_to_idx[self.name]
+
+    def get_nn_pts(self,):
+        nn_pts = 0
+        nn_pts += self.body.get_nn_pts()
+        for cur_link in self.children:
+            nn_pts += cur_link.get_nn_pts()
+        self.nn_pts = nn_pts
+        return self.nn_pts
+    
+    def clear_grads(self,):
+        
+        if self.parent_rot_mtx.grad is not None:
+            self.parent_rot_mtx.grad.data = self.parent_rot_mtx.grad.data * 0.
+        if self.parent_trans_vec.grad is not None:
+            self.parent_trans_vec.grad.data = self.parent_trans_vec.grad.data * 0.
+        if self.curr_rot_mtx.grad is not None:
+            self.curr_rot_mtx.grad.data = self.curr_rot_mtx.grad.data * 0.
+        if self.curr_trans_vec.grad is not None:
+            self.curr_trans_vec.grad.data = self.curr_trans_vec.grad.data * 0.
+        if self.tot_rot_mtx.grad is not None:
+            self.tot_rot_mtx.grad.data = self.tot_rot_mtx.grad.data * 0.
+        if self.tot_trans_vec.grad is not None:
+            self.tot_trans_vec.grad.data = self.tot_trans_vec.grad.data * 0.
+        # print(f"parent_rot_mtx: {self.parent_rot_mtx.grad}")
+        # print(f"parent_trans_vec: {self.parent_trans_vec.grad}")
+        # print(f"curr_rot_mtx: {self.curr_rot_mtx.grad}")
+        # print(f"curr_trans_vec: {self.curr_trans_vec.grad}")
+        # print(f"tot_rot_mtx: {self.tot_rot_mtx.grad}")
+        # print(f"tot_trans_vec: {self.tot_trans_vec.grad}")
+        # print(f"Joint")
+        self.joint.clear_grads()
+        self.body.clear_grad()
+        for cur_link in self.children:
+            cur_link.clear_grads()
+    
+    def set_args(self, args):
+        self.args = args
+        for cur_link in self.children:
+            cur_link.set_args(args)
+            
+    
+    
+        
+class Robot: # robot and the robot #
+    def __init__(self, children_links, args) -> None:
+        self.children = children_links
+        self.args = args
+        
+    def set_state(self, name_to_state):
+        for cur_link in self.children:
+            cur_link.set_state(name_to_state)
+    
+    def compute_transformation(self,):
+        for cur_link in self.children:
+            cur_link.compute_transformation()
+    
+    def get_name_to_visual_pts_faces(self, name_to_visual_pts_faces):
+        for cur_link in self.children:
+            cur_link.get_name_to_visual_pts_faces(name_to_visual_pts_faces)
+            
+    def get_visual_pts_list(self, visual_pts_list):
+        for cur_link in self.children:
+            cur_link.get_visual_pts_list(visual_pts_list)
+            
+    def set_joint_idx(self, joint_name_to_idx):
+        for cur_link in self.children:
+            cur_link.set_joint_idx(joint_name_to_idx) ### set joint idx ###
+
+    def set_state_via_vec(self, state_vec): ### set the state vec for the state vec ###
+        for cur_link in self.children: ### set the state vec for the state vec ###
+            cur_link.set_state_via_vec(state_vec)
+        # self.joint.set_state_via_vec(state_vec)
+        # for child_link in self.children:
+        #     child_link.set_state_via_vec(state_vec)
+        
+    # get_tot_transformed_joints
+    def get_tot_transformed_joints(self, transformed_joints):
+        for cur_link in self.children: # 
+            transformed_joints = cur_link.get_tot_transformed_joints(transformed_joints)
+        return transformed_joints
+    
+    def get_nn_pts(self):
+        nn_pts = 0
+        for cur_link in self.children:
+            nn_pts += cur_link.get_nn_pts()
+        self.nn_pts = nn_pts
+        return self.nn_pts
+    
+    def set_args(self, args):
+        self.args = args
+        for cur_link in self.children: ## args ##
+            cur_link.set_args(args)
+    
+    def print_grads(self):
+        for cur_link in self.children:
+            cur_link.print_grads()
+        
+    def clear_grads(self,): ## clear grads ##
+        for cur_link in self.children:
+            cur_link.clear_grads()
+            
+    def compute_transformation_via_state_vecs(self, state_vals, visual_pts_list):
+        # parent_rot_mtx, parent_trans_vec
+        for cur_link in self.children:
+            cur_link.compute_transformation_via_state_vecs(state_vals, cur_link.parent_rot_mtx, cur_link.parent_trans_vec, visual_pts_list)
+        return visual_pts_list
+    
+    # get_visual_pts_rgba_values(self, pts_rgba_vals_list):
+    def get_visual_pts_rgba_values(self, pts_rgba_vals_list):
+        for cur_link in self.children:
+            cur_link.get_visual_pts_rgba_values(pts_rgba_vals_list)
+        return pts_rgba_vals_list ## compute pts rgba vals list ##
+
+def parse_nparray_from_string(strr, args):
+    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.cuda(args.th_cuda_idx), 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
+    
+    
+
+
+def parse_data_from_xml(xml_fn, args):
+    
+    tree = ElementTree()
+    tree.parse(xml_fn)
+    
+    ### get total robots ###
+    robots = tree.findall("./robot")
+    i_robot = 0
+    tot_robots = []
+    for cur_robot in robots:
+        print(f"Getting robot: {i_robot}")
+        i_robot += 1
+        cur_links = cur_robot.findall("./link")
+        # i_link = 0
+        cur_robot_links = []
+        for cur_link in cur_links: ## child of the link ##
+            ### a parse link util -> the child of the link is composed of (the joint; body; and children links (with children or with no child here))
+            # cur_link_name = cur_link.attrib["name"]
+            # print(f"Getting link: {i_link} with name: {cur_link_name}")
+            # i_link += 1 ## 
+            cur_robot_links.append(parse_link_data(cur_link, args=args))
+        cur_robot_obj = Robot(cur_robot_links, args=args)
+        tot_robots.append(cur_robot_obj)
+    
+    
+    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 tot_robots
+
+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 create_zero_states():
+    nn_joints = 17
+    joint_name_to_state = {}
+    for i_j in range(nn_joints):
+        cur_joint_name = f"joint{i_j + 1}"
+        joint_name_to_state[cur_joint_name] = 0.
+    return joint_name_to_state
+
+# [6.96331033e-17 3.54807679e-06 1.74046190e-15 2.66367417e-05
+#  1.22444894e-05 3.38976792e-06 1.46917635e-15 2.66367383e-05
+#  1.22444882e-05 3.38976786e-06 1.97778813e-15 2.66367383e-05
+#  1.22444882e-05 3.38976786e-06 4.76033293e-16 1.26279884e-05
+#  3.51189993e-06 0.00000000e+00 4.89999978e-03 0.00000000e+00]
+
+
+def rotation_matrix_from_axis_angle_np(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 = np.stack(
+            [cos_ + u_xx * (1 - cos_), u_xy * (1. - cos_) + u_z * sin_, u_xz * (1. - cos_) - u_y * sin_], axis=0
+        )
+        # print(f"row_a: {row_a.size()}")
+        row_b = np.stack(
+            [u_xy * (1. - cos_) - u_z * sin_, cos_ + u_yy * (1. - cos_), u_yz * (1. - cos_) + u_x * sin_], axis=0
+        )
+        # print(f"row_b: {row_b.size()}")
+        row_c = np.stack(
+            [u_xz * (1. - cos_) + u_y * sin_, u_yz * (1. - cos_) - u_x * sin_, cos_ + u_zz * (1. - cos_)], axis=0
+        )
+        # print(f"row_c: {row_c.size()}")
+        
+        ### rot_mtx for the rot_mtx ###
+        rot_mtx = np.stack(
+            [row_a, row_b, row_c], axis=-1 ### rot_matrix of he 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 get_camera_to_world_poses(n=10, ):
+    ## sample from the upper half sphere ##
+    # theta and phi for the 
+    theta = np.random.uniform(low=0.0, high=1.0, size=(n,)) * np.pi * 2. # xz palne # 
+    phi = np.random.uniform(low=-1.0, high=0.0, size=(n,)) * np.pi ## [-0.5 \pi, 0.5 \pi] ## negative pi to the original pi 
+    # theta = torch.from_numpy(theta).float().cuda()
+    tot_c2w_matrix = []
+    for i_n in range(n):
+        # y_rot_vec = torch.tensor([0., 1., 0.]).float().cuda(th_cuda_idx)
+        # y_rot_mtx = load_utils.rotation_matrix_from_axis_angle(rot_vec, rot_angle) 
+        
+        
+        z_axis_rot_axis = np.array([0, 0, 1.], dtype=np.float32)
+        z_axis_rot_angle = np.pi - theta[i_n]
+        z_axis_rot_matrix = rotation_matrix_from_axis_angle_np(z_axis_rot_axis, z_axis_rot_angle)
+        rotated_plane_rot_axis_ori = np.array([1, -1, 0], dtype=np.float32)
+        rotated_plane_rot_axis_ori = rotated_plane_rot_axis_ori / np.sqrt(np.sum(rotated_plane_rot_axis_ori ** 2))
+        rotated_plane_rot_axis = np.matmul(z_axis_rot_matrix, rotated_plane_rot_axis_ori)
+        
+        plane_rot_angle = phi[i_n]
+        plane_rot_matrix = rotation_matrix_from_axis_angle_np(rotated_plane_rot_axis, plane_rot_angle)
+        
+        c2w_matrix = np.matmul(plane_rot_matrix, z_axis_rot_matrix)
+        c2w_trans_matrix = np.array(
+            [np.cos(theta[i_n]) * np.sin(phi[i_n]), np.sin(theta[i_n]) * np.sin(phi[i_n]), np.cos(phi[i_n])], dtype=np.float32
+        )
+        c2w_matrix = np.concatenate(
+            [c2w_matrix, c2w_trans_matrix.reshape(3, 1)], axis=-1
+        ) ##c2w matrix
+        tot_c2w_matrix.append(c2w_matrix)
+    tot_c2w_matrix = np.stack(tot_c2w_matrix, axis=0)
+    return tot_c2w_matrix
+        
+
+def get_camera_to_world_poses_th(n=10, th_cuda_idx=0):
+    ## sample from the upper half sphere ##
+    # theta and phi for the 
+    theta = np.random.uniform(low=0.0, high=1.0, size=(n,)) * np.pi * 2. # xz palne # 
+    phi = np.random.uniform(low=-1.0, high=0.0, size=(n,)) * np.pi ## [-0.5 \pi, 0.5 \pi] ## negative pi to the original pi 
+    
+    # n_total = 14
+    # n_xz = 14
+    # n_y = 7
+    # theta = [i_xz * 1.0 / float(n_xz) * np.pi * 2. for i_xz in range(n_xz)]
+    # phi = [i_y * (-1.0) / float(n_y) * np.pi for i_y in range(n_y)]
+    
+    
+    theta = torch.from_numpy(theta).float().cuda(th_cuda_idx)
+    phi = torch.from_numpy(phi).float().cuda(th_cuda_idx)
+    
+    tot_c2w_matrix = []
+    for i_n in range(n): # if use veyr dense views like those 
+        y_rot_angle = theta[i_n]
+        y_rot_vec = torch.tensor([0., 1., 0.]).float().cuda(th_cuda_idx)
+        y_rot_mtx = rotation_matrix_from_axis_angle(y_rot_vec, y_rot_angle) 
+        
+        x_axis = torch.tensor([1., 0., 0.]).float().cuda(th_cuda_idx)
+        y_rot_x_axis = torch.matmul(y_rot_mtx, x_axis.unsqueeze(-1)).squeeze(-1) ### y_rot_x_axis # 
+        
+        x_rot_angle = phi[i_n]
+        x_rot_mtx = rotation_matrix_from_axis_angle(y_rot_x_axis, x_rot_angle)
+        
+        rot_mtx = torch.matmul(x_rot_mtx, y_rot_mtx)
+        xyz_offset = torch.tensor([0., 0., 1.5]).float().cuda(th_cuda_idx) 
+        rot_xyz_offset = torch.matmul(rot_mtx, xyz_offset.unsqueeze(-1)).squeeze(-1).contiguous() + 0.5 ### 3 for the xyz offset 
+        
+        c2w_matrix = torch.cat(
+            [rot_mtx, rot_xyz_offset.unsqueeze(-1)], dim=-1
+        )
+        tot_c2w_matrix.append(c2w_matrix)
+        
+        
+        # z_axis_rot_axis = np.array([0, 0, 1.], dtype=np.float32)
+        # z_axis_rot_angle = np.pi - theta[i_n]
+        # z_axis_rot_matrix = rotation_matrix_from_axis_angle_np(z_axis_rot_axis, z_axis_rot_angle)
+        # rotated_plane_rot_axis_ori = np.array([1, -1, 0], dtype=np.float32)
+        # rotated_plane_rot_axis_ori = rotated_plane_rot_axis_ori / np.sqrt(np.sum(rotated_plane_rot_axis_ori ** 2))
+        # rotated_plane_rot_axis = np.matmul(z_axis_rot_matrix, rotated_plane_rot_axis_ori)
+        
+        # plane_rot_angle = phi[i_n]
+        # plane_rot_matrix = rotation_matrix_from_axis_angle_np(rotated_plane_rot_axis, plane_rot_angle)
+        
+        # c2w_matrix = np.matmul(plane_rot_matrix, z_axis_rot_matrix)
+        # c2w_trans_matrix = np.array(
+        #     [np.cos(theta[i_n]) * np.sin(phi[i_n]), np.sin(theta[i_n]) * np.sin(phi[i_n]), np.cos(phi[i_n])], dtype=np.float32
+        # )
+        # c2w_matrix = np.concatenate(
+        #     [c2w_matrix, c2w_trans_matrix.reshape(3, 1)], axis=-1
+        # ) ##c2w matrix
+        # tot_c2w_matrix.append(c2w_matrix)
+    # tot_c2w_matrix = np.stack(tot_c2w_matrix, axis=0)
+    tot_c2w_matrix = torch.stack(tot_c2w_matrix, dim=0)
+    return tot_c2w_matrix
+
+
+def get_camera_to_world_poses_th_routine_1(n=7, th_cuda_idx=0):
+    ## sample from the upper half sphere ##
+    # theta and phi for the 
+    
+    # theta = np.random.uniform(low=0.0, high=1.0, size=(n,)) * np.pi * 2. # xz palne # 
+    # phi = np.random.uniform(low=-1.0, high=0.0, size=(n,)) * np.pi ## [-0.5 \pi, 0.5 \pi] ## negative pi to the original pi 
+    
+    # n_total = 14
+    n_xz = 2 * n #  14
+    n_y = n # 7
+    theta = [i_xz * 1.0 / float(n_xz) * np.pi * 2. for i_xz in range(n_xz)]
+    phi = [i_y * (-1.0) / float(n_y) * np.pi for i_y in range(n_y)]
+    
+    theta = torch.tensor(theta).float().cuda(th_cuda_idx)
+    phi = torch.tensor(phi).float().cuda(th_cuda_idx)
+    # theta = torch.from_numpy(theta).float().cuda(th_cuda_idx)
+    # phi = torch.from_numpy(phi).float().cuda(th_cuda_idx)
+    
+    tot_c2w_matrix = []
+    
+    for i_theta in range(theta.size(0)):
+        for i_phi in range(phi.size(0)):
+            y_rot_angle = theta[i_theta]
+            y_rot_vec = torch.tensor([0., 1., 0.]).float().cuda(th_cuda_idx)
+            y_rot_mtx = rotation_matrix_from_axis_angle(y_rot_vec, y_rot_angle) 
+            
+            x_axis = torch.tensor([1., 0., 0.]).float().cuda(th_cuda_idx)
+            y_rot_x_axis = torch.matmul(y_rot_mtx, x_axis.unsqueeze(-1)).squeeze(-1) ### y_rot_x_axis # 
+            
+            x_rot_angle = phi[i_phi]
+            x_rot_mtx = rotation_matrix_from_axis_angle(y_rot_x_axis, x_rot_angle)
+            
+            rot_mtx = torch.matmul(x_rot_mtx, y_rot_mtx)
+            xyz_offset = torch.tensor([0., 0., 1.5]).float().cuda(th_cuda_idx) 
+            rot_xyz_offset = torch.matmul(rot_mtx, xyz_offset.unsqueeze(-1)).squeeze(-1).contiguous() + 0.5 ### 3 for the xyz offset 
+            
+            c2w_matrix = torch.cat(
+                [rot_mtx, rot_xyz_offset.unsqueeze(-1)], dim=-1
+            )
+            tot_c2w_matrix.append(c2w_matrix)
+    
+    tot_c2w_matrix = torch.stack(tot_c2w_matrix, dim=0)
+    return tot_c2w_matrix
+
+
+def get_camera_to_world_poses_th_routine_2(n=7, th_cuda_idx=0):
+    ## sample from the upper half sphere ##
+    # theta and phi for the 
+    
+    # theta = np.random.uniform(low=0.0, high=1.0, size=(n,)) * np.pi * 2. # xz palne # 
+    # phi = np.random.uniform(low=-1.0, high=0.0, size=(n,)) * np.pi ## [-0.5 \pi, 0.5 \pi] ## negative pi to the original pi 
+    
+    # n_total = 14
+    n_xz = 2 * n #  14
+    n_y = 2 * n # 7
+    theta = [i_xz * 1.0 / float(n_xz) * np.pi * 2. for i_xz in range(n_xz)]
+    # phi = [i_y * (-1.0) / float(n_y) * np.pi for i_y in range(n_y)]
+    phi = [i_y * (-1.0) / float(n_y) * np.pi * 2. for i_y in range(n_y)]
+    
+    theta = torch.tensor(theta).float().cuda(th_cuda_idx)
+    phi = torch.tensor(phi).float().cuda(th_cuda_idx)
+    # theta = torch.from_numpy(theta).float().cuda(th_cuda_idx)
+    # phi = torch.from_numpy(phi).float().cuda(th_cuda_idx)
+    
+    tot_c2w_matrix = []
+    
+    for i_theta in range(theta.size(0)):
+        for i_phi in range(phi.size(0)):
+            y_rot_angle = theta[i_theta]
+            y_rot_vec = torch.tensor([0., 1., 0.]).float().cuda(th_cuda_idx)
+            y_rot_mtx = rotation_matrix_from_axis_angle(y_rot_vec, y_rot_angle) 
+            
+            x_axis = torch.tensor([1., 0., 0.]).float().cuda(th_cuda_idx)
+            y_rot_x_axis = torch.matmul(y_rot_mtx, x_axis.unsqueeze(-1)).squeeze(-1) ### y_rot_x_axis # 
+            
+            x_rot_angle = phi[i_phi]
+            x_rot_mtx = rotation_matrix_from_axis_angle(y_rot_x_axis, x_rot_angle)
+            
+            rot_mtx = torch.matmul(x_rot_mtx, y_rot_mtx)
+            xyz_offset = torch.tensor([0., 0., 1.5]).float().cuda(th_cuda_idx) 
+            rot_xyz_offset = torch.matmul(rot_mtx, xyz_offset.unsqueeze(-1)).squeeze(-1).contiguous() + 0.5 ### 3 for the xyz offset 
+            
+            c2w_matrix = torch.cat(
+                [rot_mtx, rot_xyz_offset.unsqueeze(-1)], dim=-1
+            )
+            tot_c2w_matrix.append(c2w_matrix)
+    
+    tot_c2w_matrix = torch.stack(tot_c2w_matrix, dim=0)
+    return tot_c2w_matrix
+
+
+
+
+
+
+if __name__=='__main__':
+    xml_fn =   "/home/xueyi/diffsim/DiffHand/assets/hand_sphere.xml"
+    tot_robots = parse_data_from_xml(xml_fn=xml_fn)
+    # tot_robots = 
+    
+    active_optimized_states = """-0.00025872 -0.00025599 -0.00025296 -0.00022881 -0.00024449 -0.0002549 -0.00025296 -0.00022881 -0.00024449 -0.0002549 -0.00025296 -0.00022881 -0.00024449 -0.0002549 -0.00025694 -0.00024656 -0.00025556 0. 0.0049 0."""
+    active_optimized_states = """-1.10617972 -1.10742263 -1.06198363 -1.03212746 -1.05429142 -1.08617289 -1.05868192 -1.01624365 -1.04478191 -1.08260959 -1.06719107 -1.04082455 -1.05995886 -1.08674006 -1.09396691 -1.08965532 -1.10036577 -10.7117466 -3.62511998 1.49450353"""
+    # active_goal_optimized_states = """-1.10617972 -1.10742263 -1.0614858 -1.03189609 -1.05404354 -1.08610468 -1.05863293 -1.0174248 -1.04576456 -1.08297396 -1.06719107 -1.04082455 -1.05995886 -1.08674006 -1.09396691 -1.08965532 -1.10036577 -10.73396897 -3.68095432 1.50679285"""
+    active_optimized_states = """-0.42455298 -0.42570447 -0.40567708 -0.39798589 -0.40953955 -0.42025055 -0.37910662 -0.496165 -0.37664644 -0.41942727 -0.40596508 -0.3982109 -0.40959847 -0.42024905 -0.41835001 -0.41929961 -0.42365131 -1.18756073 -2.90337822 0.4224685"""
+    active_optimized_states = """-0.42442816 -0.42557961 -0.40366201 -0.3977891 -0.40947627 -0.4201424 -0.3799285 -0.3808375 -0.37953552 -0.42039598 -0.4058405 -0.39808804 -0.40947487 -0.42012458 -0.41822534 -0.41917521 -0.4235266 -0.87189658 -1.42093761 0.21977979"""
+    
+    active_robot = tot_robots[0]
+    zero_states = create_zero_states()
+    active_robot.set_state(zero_states)
+    active_robot.compute_transformation()
+    name_to_visual_pts_surfaces = {}
+    active_robot.get_name_to_visual_pts_faces(name_to_visual_pts_surfaces)
+    print(len(name_to_visual_pts_surfaces))
+    
+    sv_res_rt = "/home/xueyi/diffsim/DiffHand/examples/save_res"
+    sv_res_rt = os.path.join(sv_res_rt, "load_utils_test")
+    os.makedirs(sv_res_rt, exist_ok=True)
+    
+    tmp_visual_res_sv_fn = os.path.join(sv_res_rt, f"res_with_zero_states.npy")
+    np.save(tmp_visual_res_sv_fn, name_to_visual_pts_surfaces)
+    print(f"tmp visual res saved to {tmp_visual_res_sv_fn}")
+    
+    
+    optimized_states = get_name_to_state_from_str(active_optimized_states)
+    active_robot.set_state(optimized_states)
+    active_robot.compute_transformation()
+    name_to_visual_pts_surfaces = {}
+    active_robot.get_name_to_visual_pts_faces(name_to_visual_pts_surfaces)
+    print(len(name_to_visual_pts_surfaces))
+    # sv_res_rt = "/home/xueyi/diffsim/DiffHand/examples/save_res"
+    # sv_res_rt = os.path.join(sv_res_rt, "load_utils_test")
+    # os.makedirs(sv_res_rt, exist_ok=True)
+    
+    # tmp_visual_res_sv_fn = os.path.join(sv_res_rt, f"res_with_optimized_states.npy")
+    tmp_visual_res_sv_fn = os.path.join(sv_res_rt, f"active_ngoal_res_with_optimized_states_goal_n3.npy")
+    np.save(tmp_visual_res_sv_fn, name_to_visual_pts_surfaces)
+    print(f"tmp visual res with optimized states saved to {tmp_visual_res_sv_fn}")
+    

models/embedder.py

@@ -0,0 +1,51 @@
+import torch
+import torch.nn as nn
+
+
+# Positional encoding embedding. Code was taken from https://github.com/bmild/nerf.
+class Embedder:
+    def __init__(self, **kwargs):
+        self.kwargs = kwargs
+        self.create_embedding_fn()
+
+    def create_embedding_fn(self):
+        embed_fns = []
+        d = self.kwargs['input_dims']
+        out_dim = 0
+        if self.kwargs['include_input']:
+            embed_fns.append(lambda x: x)
+            out_dim += d
+
+        max_freq = self.kwargs['max_freq_log2']
+        N_freqs = self.kwargs['num_freqs']
+
+        if self.kwargs['log_sampling']:
+            freq_bands = 2. ** torch.linspace(0., max_freq, N_freqs)
+        else:
+            freq_bands = torch.linspace(2.**0., 2.**max_freq, N_freqs)
+
+        for freq in freq_bands:
+            for p_fn in self.kwargs['periodic_fns']:
+                embed_fns.append(lambda x, p_fn=p_fn, freq=freq: p_fn(x * freq))
+                out_dim += d
+
+        self.embed_fns = embed_fns
+        self.out_dim = out_dim
+
+    def embed(self, inputs):
+        return torch.cat([fn(inputs) for fn in self.embed_fns], -1)
+
+
+def get_embedder(multires, input_dims=3):
+    embed_kwargs = {
+        'include_input': True,
+        'input_dims': input_dims,
+        'max_freq_log2': multires-1,
+        'num_freqs': multires,
+        'log_sampling': True,
+        'periodic_fns': [torch.sin, torch.cos],
+    }
+
+    embedder_obj = Embedder(**embed_kwargs)
+    def embed(x, eo=embedder_obj): return eo.embed(x)
+    return embed, embedder_obj.out_dim

models/renderer.py

@@ -0,0 +1,641 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import logging
+import mcubes
+from icecream import ic
+import os
+
+import trimesh
+from pysdf import SDF
+
+from uni_rep.rep_3d.dmtet import marching_tets_tetmesh, create_tetmesh_variables
+
+def create_mt_variable(device):
+    triangle_table = torch.tensor(
+        [
+            [-1, -1, -1, -1, -1, -1],
+            [1, 0, 2, -1, -1, -1],
+            [4, 0, 3, -1, -1, -1],
+            [1, 4, 2, 1, 3, 4],
+            [3, 1, 5, -1, -1, -1],
+            [2, 3, 0, 2, 5, 3],
+            [1, 4, 0, 1, 5, 4],
+            [4, 2, 5, -1, -1, -1],
+            [4, 5, 2, -1, -1, -1],
+            [4, 1, 0, 4, 5, 1],
+            [3, 2, 0, 3, 5, 2],
+            [1, 3, 5, -1, -1, -1],
+            [4, 1, 2, 4, 3, 1],
+            [3, 0, 4, -1, -1, -1],
+            [2, 0, 1, -1, -1, -1],
+            [-1, -1, -1, -1, -1, -1]
+        ], dtype=torch.long, device=device)
+
+    num_triangles_table = torch.tensor([0, 1, 1, 2, 1, 2, 2, 1, 1, 2, 2, 1, 2, 1, 1, 0], dtype=torch.long, device=device)
+    base_tet_edges = torch.tensor([0, 1, 0, 2, 0, 3, 1, 2, 1, 3, 2, 3], dtype=torch.long, device=device)
+    v_id = torch.pow(2, torch.arange(4, dtype=torch.long, device=device))
+    return triangle_table, num_triangles_table, base_tet_edges, v_id
+
+
+
+def  extract_fields_from_tets(bound_min, bound_max, resolution, query_func):
+    # load tet via resolution #
+    # scale them via bounds #
+    # extract the geometry #
+    # /home/xueyi/gen/DeepMetaHandles/data/tets/100_compress.npz # strange #
+    device = bound_min.device
+    # if resolution in [64, 70, 80, 90, 100]:
+    #     tet_fn = f"/home/xueyi/gen/DeepMetaHandles/data/tets/{resolution}_compress.npz"
+    # else:
+    tet_fn = f"/home/xueyi/gen/DeepMetaHandles/data/tets/{100}_compress.npz"
+    tets = np.load(tet_fn)
+    verts = torch.from_numpy(tets['vertices']).float().to(device) # verts positions 
+    indices = torch.from_numpy(tets['tets']).long().to(device) # .to(self.device)
+    # split #
+    # verts; verts; #
+    minn_verts, _ = torch.min(verts, dim=0)
+    maxx_verts, _ = torch.max(verts, dim=0) # (3, ) # exporting the  
+    # scale_verts = maxx_verts - minn_verts
+    scale_bounds = bound_max - bound_min # scale bounds #
+    
+    ### scale the vertices ###
+    scaled_verts = (verts - minn_verts.unsqueeze(0)) / (maxx_verts - minn_verts).unsqueeze(0) ### the maxx and minn verts scales ###
+    
+    # scaled_verts = (verts - minn_verts.unsqueeze(0)) / (maxx_verts - minn_verts).unsqueeze(0) ### the maxx and minn verts scales ###
+
+    scaled_verts = scaled_verts * 2. - 1. # init the sdf filed viathe tet mesh vertices and the sdf values ##
+    # scaled_verts = (scaled_verts * scale_bounds.unsqueeze(0)) + bound_min.unsqueeze(0) ## the scaled verts ###
+    
+    # scaled_verts = scaled_verts - scale_bounds.unsqueeze(0) / 2. # 
+    # scaled_verts = scaled_verts -  bound_min.unsqueeze(0) - scale_bounds.unsqueeze(0) / 2.
+    
+    sdf_values = []
+    N = 64
+    query_bundles = N ** 3 ### N^3
+    query_NNs = scaled_verts.size(0) // query_bundles
+    if query_NNs * query_bundles < scaled_verts.size(0):
+        query_NNs += 1
+    for i_query in range(query_NNs):
+        cur_bundle_st = i_query * query_bundles
+        cur_bundle_ed = (i_query + 1) * query_bundles
+        cur_bundle_ed = min(cur_bundle_ed, scaled_verts.size(0))
+        cur_query_pts = scaled_verts[cur_bundle_st: cur_bundle_ed]
+        cur_query_vals = query_func(cur_query_pts)
+        sdf_values.append(cur_query_vals)
+    sdf_values = torch.cat(sdf_values, dim=0)
+    # print(f"queryed sdf values: {sdf_values.size()}") #
+    
+    GT_sdf_values = np.load("/home/xueyi/diffsim/DiffHand/assets/hand/100_sdf_values.npy", allow_pickle=True)
+    GT_sdf_values = torch.from_numpy(GT_sdf_values).float().to(device)
+    
+    # intrinsic, tet values, pts values, sdf network #
+    triangle_table, num_triangles_table, base_tet_edges, v_id = create_mt_variable(device)
+    tet_table, num_tets_table = create_tetmesh_variables(device)
+    
+    sdf_values = sdf_values.squeeze(-1) # how the rendering # 
+    
+    # print(f"GT_sdf_values: {GT_sdf_values.size()}, sdf_values: {sdf_values.size()}, scaled_verts: {scaled_verts.size()}")
+    # print(f"scaled_verts: {scaled_verts.size()}, ")
+    # pos_nx3, sdf_n, tet_fx4, triangle_table, num_triangles_table, base_tet_edges, v_id,
+    # return_tet_mesh=False, ori_v=None, num_tets_table=None, tet_table=None):
+    # marching_tets_tetmesh ##
+    verts, faces, tet_verts, tets = marching_tets_tetmesh(scaled_verts, sdf_values, indices, triangle_table, num_triangles_table, base_tet_edges, v_id, return_tet_mesh=True, ori_v=scaled_verts, num_tets_table=num_tets_table, tet_table=tet_table)
+    ### use the GT sdf values for the marching tets ###
+    GT_verts,  GT_faces,  GT_tet_verts,  GT_tets = marching_tets_tetmesh(scaled_verts, GT_sdf_values, indices, triangle_table, num_triangles_table, base_tet_edges, v_id, return_tet_mesh=True, ori_v=scaled_verts, num_tets_table=num_tets_table, tet_table=tet_table)
+    
+    # print(f"After tet marching with verts: {verts.size()}, faces: {faces.size()}")
+    return verts, faces, sdf_values, GT_verts, GT_faces  # verts, faces #
+
+def extract_fields(bound_min, bound_max, resolution, query_func):
+    N = 64
+    X = torch.linspace(bound_min[0], bound_max[0], resolution).split(N)
+    Y = torch.linspace(bound_min[1], bound_max[1], resolution).split(N)
+    Z = torch.linspace(bound_min[2], bound_max[2], resolution).split(N)
+
+    u = np.zeros([resolution, resolution, resolution], dtype=np.float32)
+    with torch.no_grad():
+        for xi, xs in enumerate(X):
+            for yi, ys in enumerate(Y):
+                for zi, zs in enumerate(Z):
+                    xx, yy, zz = torch.meshgrid(xs, ys, zs)
+                    pts = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1)
+                    val = query_func(pts).reshape(len(xs), len(ys), len(zs)).detach().cpu().numpy()
+                    u[xi * N: xi * N + len(xs), yi * N: yi * N + len(ys), zi * N: zi * N + len(zs)] = val
+        # should save u here #
+        # save_u_path = os.path.join("/data2/datasets/diffsim/neus/exp/hand_test/womask_sphere_reverse_value/other_saved", "sdf_values.npy")
+        # np.save(save_u_path, u)
+        # print(f"u saved to {save_u_path}")
+    return u
+
+
+def extract_geometry(bound_min, bound_max, resolution, threshold, query_func):
+    print('threshold: {}'.format(threshold))
+    
+    ## using maching cubes ###
+    u = extract_fields(bound_min, bound_max, resolution, query_func)
+    vertices, triangles = mcubes.marching_cubes(u, threshold) # grid sdf and marching cubes #
+    b_max_np = bound_max.detach().cpu().numpy()
+    b_min_np = bound_min.detach().cpu().numpy()
+
+    vertices = vertices / (resolution - 1.0) * (b_max_np - b_min_np)[None, :] + b_min_np[None, :]
+    ### using maching cubes ###
+    
+    ### using marching tets ###
+    # vertices, triangles = extract_fields_from_tets(bound_min, bound_max, resolution, query_func)
+    # vertices = vertices.detach().cpu().numpy()
+    # triangles = triangles.detach().cpu().numpy()
+    ### using marching tets ###
+    
+    # b_max_np = bound_max.detach().cpu().numpy()
+    # b_min_np = bound_min.detach().cpu().numpy()
+
+    # vertices = vertices / (resolution - 1.0) * (b_max_np - b_min_np)[None, :] + b_min_np[None, :]
+    return vertices, triangles
+
+def extract_geometry_tets(bound_min, bound_max, resolution, threshold, query_func):
+    # print('threshold: {}'.format(threshold))
+    
+    ### using maching cubes ###
+    # u = extract_fields(bound_min, bound_max, resolution, query_func)
+    # vertices, triangles = mcubes.marching_cubes(u, threshold) # grid sdf and marching cubes #
+    # b_max_np = bound_max.detach().cpu().numpy()
+    # b_min_np = bound_min.detach().cpu().numpy()
+
+    # vertices = vertices / (resolution - 1.0) * (b_max_np - b_min_np)[None, :] + b_min_np[None, :]
+    ### using maching cubes ###
+    
+    ## 
+    ### using marching tets ### fiels from tets ##
+    vertices, triangles, tet_sdf_values,  GT_verts, GT_faces  = extract_fields_from_tets(bound_min, bound_max, resolution, query_func)
+    # vertices = vertices.detach().cpu().numpy()
+    # triangles = triangles.detach().cpu().numpy()
+    ### using marching tets ###
+    
+    # b_max_np = bound_max.detach().cpu().numpy()
+    # b_min_np = bound_min.detach().cpu().numpy()
+    #
+
+    # vertices = vertices / (resolution - 1.0) * (b_max_np - b_min_np)[None, :] + b_min_np[None, :]
+    return vertices, triangles, tet_sdf_values,  GT_verts, GT_faces 
+
+
+def sample_pdf(bins, weights, n_samples, det=False):
+    # This implementation is from NeRF
+    # Get pdf
+    weights = weights + 1e-5  # prevent nans
+    pdf = weights / torch.sum(weights, -1, keepdim=True)
+    cdf = torch.cumsum(pdf, -1)
+    cdf = torch.cat([torch.zeros_like(cdf[..., :1]), cdf], -1)
+    # Take uniform samples
+    if det:
+        u = torch.linspace(0. + 0.5 / n_samples, 1. - 0.5 / n_samples, steps=n_samples)
+        u = u.expand(list(cdf.shape[:-1]) + [n_samples])
+    else:
+        u = torch.rand(list(cdf.shape[:-1]) + [n_samples])
+
+    # Invert CDF # invert cdf # 
+    u = u.contiguous()
+    inds = torch.searchsorted(cdf, u, right=True)
+    below = torch.max(torch.zeros_like(inds - 1), inds - 1)
+    above = torch.min((cdf.shape[-1] - 1) * torch.ones_like(inds), inds)
+    inds_g = torch.stack([below, above], -1)  # (batch, N_samples, 2)
+
+    matched_shape = [inds_g.shape[0], inds_g.shape[1], cdf.shape[-1]]
+    cdf_g = torch.gather(cdf.unsqueeze(1).expand(matched_shape), 2, inds_g)
+    bins_g = torch.gather(bins.unsqueeze(1).expand(matched_shape), 2, inds_g)
+
+    denom = (cdf_g[..., 1] - cdf_g[..., 0])
+    denom = torch.where(denom < 1e-5, torch.ones_like(denom), denom)
+    t = (u - cdf_g[..., 0]) / denom
+    samples = bins_g[..., 0] + t * (bins_g[..., 1] - bins_g[..., 0])
+
+    return samples
+
+
+def load_GT_vertices(GT_meshes_folder):
+    tot_meshes_fns = os.listdir(GT_meshes_folder)
+    tot_meshes_fns = [fn for fn in tot_meshes_fns if fn.endswith(".obj")]
+    tot_mesh_verts = []
+    tot_mesh_faces = []
+    n_tot_verts = 0
+    for fn in tot_meshes_fns:
+        cur_mesh_fn = os.path.join(GT_meshes_folder, fn)
+        obj_mesh = trimesh.load(cur_mesh_fn, process=False)
+        # obj_mesh.remove_degenerate_faces(height=1e-06)
+
+        verts_obj = np.array(obj_mesh.vertices)
+        faces_obj = np.array(obj_mesh.faces)
+        
+        tot_mesh_verts.append(verts_obj)
+        tot_mesh_faces.append(faces_obj + n_tot_verts)
+        n_tot_verts += verts_obj.shape[0]
+        
+        # tot_mesh_faces.append(faces_obj)
+    tot_mesh_verts = np.concatenate(tot_mesh_verts, axis=0)
+    tot_mesh_faces = np.concatenate(tot_mesh_faces, axis=0)
+    return tot_mesh_verts, tot_mesh_faces
+
+
+class NeuSRenderer:
+    def __init__(self,
+                 nerf,
+                 sdf_network,
+                 deviation_network,
+                 color_network,
+                 n_samples,
+                 n_importance,
+                 n_outside,
+                 up_sample_steps,
+                 perturb):
+        self.nerf = nerf
+        self.sdf_network = sdf_network
+        self.deviation_network = deviation_network
+        self.color_network = color_network
+        self.n_samples = n_samples
+        self.n_importance = n_importance
+        self.n_outside = n_outside
+        self.up_sample_steps = up_sample_steps
+        self.perturb = perturb
+        
+        GT_meshes_folder = "/home/xueyi/diffsim/DiffHand/assets/hand"
+        self.mesh_vertices, self.mesh_faces = load_GT_vertices(GT_meshes_folder=GT_meshes_folder)
+        maxx_pts = 25.
+        minn_pts = -15.
+        self.mesh_vertices = (self.mesh_vertices - minn_pts) / (maxx_pts - minn_pts)
+        f = SDF(self.mesh_vertices, self.mesh_faces)
+        self.gt_sdf = f ## a unite sphere or box
+        
+        self.minn_pts = 0
+        self.maxx_pts = 1.
+        
+        # self.minn_pts = -1.5
+        # self.maxx_pts = 1.5
+        self.bkg_pts = ... # TODO: the bkg pts 
+        self.cur_fr_bkg_pts_defs = ... # TODO: set the cur_bkg_pts_defs for each frame #
+        self.dist_interp_thres = ... # TODO
+        
+    # get the pts and render the pts #
+    # pts and the rendering pts #
+    def deform_pts(self, pts):
+        # pts: nn_batch x nn_samples x 3
+        nnb, nns = pts.size(0), pts.size(1)
+        pts_exp = pts.contiguous().view(nnb * nns, -1).contiguous()
+        dist_pts_to_bkg_pts = torch.sum(
+            (pts_exp.unsqueeze(1) - self.bkg_pts.unsqueeze(0)) ** 2, dim=-1 ## nn_pts_exp x nn_bkg_pts 
+        )
+        
+    
+        
+
+    def render_core_outside(self, rays_o, rays_d, z_vals, sample_dist, nerf, background_rgb=None):
+        """
+        Render background
+        """
+        batch_size, n_samples = z_vals.shape
+
+        # Section length
+        dists = z_vals[..., 1:] - z_vals[..., :-1]
+        dists = torch.cat([dists, torch.Tensor([sample_dist]).expand(dists[..., :1].shape)], -1)
+        mid_z_vals = z_vals + dists * 0.5
+
+        # Section midpoints #
+        pts = rays_o[:, None, :] + rays_d[:, None, :] * mid_z_vals[..., :, None]  # batch_size, n_samples, 3 #
+        
+        # pts = pts.flip((-1,)) * 2 - 1
+        pts = pts * 2 - 1
+
+        dis_to_center = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=True).clip(1.0, 1e10)
+        pts = torch.cat([pts / dis_to_center, 1.0 / dis_to_center], dim=-1)       # batch_size, n_samples, 4 #
+
+        dirs = rays_d[:, None, :].expand(batch_size, n_samples, 3)
+
+        pts = pts.reshape(-1, 3 + int(self.n_outside > 0))
+        dirs = dirs.reshape(-1, 3)
+
+        density, sampled_color = nerf(pts, dirs)
+        sampled_color = torch.sigmoid(sampled_color)
+        alpha = 1.0 - torch.exp(-F.softplus(density.reshape(batch_size, n_samples)) * dists)
+        alpha = alpha.reshape(batch_size, n_samples)
+        weights = alpha * torch.cumprod(torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[:, :-1]
+        sampled_color = sampled_color.reshape(batch_size, n_samples, 3)
+        color = (weights[:, :, None] * sampled_color).sum(dim=1)
+        if background_rgb is not None:
+            color = color + background_rgb * (1.0 - weights.sum(dim=-1, keepdim=True))
+
+        return {
+            'color': color,
+            'sampled_color': sampled_color,
+            'alpha': alpha,
+            'weights': weights,
+        }
+
+    def up_sample(self, rays_o, rays_d, z_vals, sdf, n_importance, inv_s):
+        """
+        Up sampling give a fixed inv_s
+        """
+        batch_size, n_samples = z_vals.shape
+        pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]  # n_rays, n_samples, 3
+        
+        # pts = pts.flip((-1,)) * 2 - 1
+        pts = pts * 2 - 1
+        
+        radius = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=False)
+        inside_sphere = (radius[:, :-1] < 1.0) | (radius[:, 1:] < 1.0)
+        sdf = sdf.reshape(batch_size, n_samples)
+        prev_sdf, next_sdf = sdf[:, :-1], sdf[:, 1:]
+        prev_z_vals, next_z_vals = z_vals[:, :-1], z_vals[:, 1:]
+        mid_sdf = (prev_sdf + next_sdf) * 0.5
+        cos_val = (next_sdf - prev_sdf) / (next_z_vals - prev_z_vals + 1e-5)
+
+        # ----------------------------------------------------------------------------------------------------------
+        # Use min value of [ cos, prev_cos ]
+        # Though it makes the sampling (not rendering) a little bit biased, this strategy can make the sampling more
+        # robust when meeting situations like below:
+        #
+        # SDF
+        # ^
+        # |\          -----x----...
+        # | \        /
+        # |  x      x
+        # |---\----/-------------> 0 level
+        # |    \  /
+        # |     \/
+        # |
+        # ----------------------------------------------------------------------------------------------------------
+        prev_cos_val = torch.cat([torch.zeros([batch_size, 1]), cos_val[:, :-1]], dim=-1)
+        cos_val = torch.stack([prev_cos_val, cos_val], dim=-1)
+        cos_val, _ = torch.min(cos_val, dim=-1, keepdim=False)
+        cos_val = cos_val.clip(-1e3, 0.0) * inside_sphere
+
+        dist = (next_z_vals - prev_z_vals)
+        prev_esti_sdf = mid_sdf - cos_val * dist * 0.5
+        next_esti_sdf = mid_sdf + cos_val * dist * 0.5
+        prev_cdf = torch.sigmoid(prev_esti_sdf * inv_s)
+        next_cdf = torch.sigmoid(next_esti_sdf * inv_s)
+        alpha = (prev_cdf - next_cdf + 1e-5) / (prev_cdf + 1e-5)
+        weights = alpha * torch.cumprod(
+            torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[:, :-1]
+
+        z_samples = sample_pdf(z_vals, weights, n_importance, det=True).detach()
+        return z_samples
+
+    def cat_z_vals(self, rays_o, rays_d, z_vals, new_z_vals, sdf, last=False):
+        batch_size, n_samples = z_vals.shape
+        _, n_importance = new_z_vals.shape
+        pts = rays_o[:, None, :] + rays_d[:, None, :] * new_z_vals[..., :, None]
+        
+        # pts = pts.flip((-1,)) * 2 - 1
+        pts = pts * 2 - 1
+        
+        z_vals = torch.cat([z_vals, new_z_vals], dim=-1)
+        z_vals, index = torch.sort(z_vals, dim=-1)
+
+        if not last:
+            new_sdf = self.sdf_network.sdf(pts.reshape(-1, 3)).reshape(batch_size, n_importance)
+            sdf = torch.cat([sdf, new_sdf], dim=-1)
+            xx = torch.arange(batch_size)[:, None].expand(batch_size, n_samples + n_importance).reshape(-1)
+            index = index.reshape(-1)
+            sdf = sdf[(xx, index)].reshape(batch_size, n_samples + n_importance)
+
+        return z_vals, sdf
+
+    def render_core(self,
+                    rays_o,
+                    rays_d,
+                    z_vals,
+                    sample_dist,
+                    sdf_network,
+                    deviation_network,
+                    color_network,
+                    background_alpha=None,
+                    background_sampled_color=None,
+                    background_rgb=None,
+                    cos_anneal_ratio=0.0):
+        batch_size, n_samples = z_vals.shape
+
+        # Section length
+        dists = z_vals[..., 1:] - z_vals[..., :-1]
+        dists = torch.cat([dists, torch.Tensor([sample_dist]).expand(dists[..., :1].shape)], -1)
+        mid_z_vals = z_vals + dists * 0.5 # z_vals and dists * 0.5 #
+
+        # Section midpoints
+        pts = rays_o[:, None, :] + rays_d[:, None, :] * mid_z_vals[..., :, None]  # n_rays, n_samples, 3
+        dirs = rays_d[:, None, :].expand(pts.shape)
+
+        pts = pts.reshape(-1, 3) # pts, nn_ou
+        dirs = dirs.reshape(-1, 3)
+        
+        pts = (pts - self.minn_pts) / (self.maxx_pts - self.minn_pts)
+        
+        # pts = pts.flip((-1,)) * 2 - 1 # 
+        pts = pts * 2 - 1
+
+        sdf_nn_output = sdf_network(pts)
+        sdf = sdf_nn_output[:, :1]
+        feature_vector = sdf_nn_output[:, 1:]
+
+        gradients = sdf_network.gradient(pts).squeeze()
+        sampled_color = color_network(pts, gradients, dirs, feature_vector).reshape(batch_size, n_samples, 3)
+
+        # deviation network #
+        inv_s = deviation_network(torch.zeros([1, 3]))[:, :1].clip(1e-6, 1e6)           # Single parameter
+        inv_s = inv_s.expand(batch_size * n_samples, 1)
+
+        true_cos = (dirs * gradients).sum(-1, keepdim=True)
+
+        # "cos_anneal_ratio" grows from 0 to 1 in the beginning training iterations. The anneal strategy below makes
+        # the cos value "not dead" at the beginning training iterations, for better convergence.
+        iter_cos = -(F.relu(-true_cos * 0.5 + 0.5) * (1.0 - cos_anneal_ratio) +
+                     F.relu(-true_cos) * cos_anneal_ratio)  # always non-positive
+
+        # Estimate signed distances at section points
+        estimated_next_sdf = sdf + iter_cos * dists.reshape(-1, 1) * 0.5
+        estimated_prev_sdf = sdf - iter_cos * dists.reshape(-1, 1) * 0.5
+
+        prev_cdf = torch.sigmoid(estimated_prev_sdf * inv_s)
+        next_cdf = torch.sigmoid(estimated_next_sdf * inv_s)
+
+        p = prev_cdf - next_cdf
+        c = prev_cdf
+
+        alpha = ((p + 1e-5) / (c + 1e-5)).reshape(batch_size, n_samples).clip(0.0, 1.0)
+
+        pts_norm = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=True).reshape(batch_size, n_samples)
+        inside_sphere = (pts_norm < 1.0).float().detach()
+        relax_inside_sphere = (pts_norm < 1.2).float().detach()
+
+        # Render with background
+        if background_alpha is not None:
+            alpha = alpha * inside_sphere + background_alpha[:, :n_samples] * (1.0 - inside_sphere)
+            alpha = torch.cat([alpha, background_alpha[:, n_samples:]], dim=-1)
+            sampled_color = sampled_color * inside_sphere[:, :, None] +\
+                            background_sampled_color[:, :n_samples] * (1.0 - inside_sphere)[:, :, None]
+            sampled_color = torch.cat([sampled_color, background_sampled_color[:, n_samples:]], dim=1)
+
+        weights = alpha * torch.cumprod(torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[:, :-1]
+        weights_sum = weights.sum(dim=-1, keepdim=True)
+
+        color = (sampled_color * weights[:, :, None]).sum(dim=1)
+        if background_rgb is not None:    # Fixed background, usually black
+            color = color + background_rgb * (1.0 - weights_sum)
+
+        # Eikonal loss
+        gradient_error = (torch.linalg.norm(gradients.reshape(batch_size, n_samples, 3), ord=2,
+                                            dim=-1) - 1.0) ** 2
+        gradient_error = (relax_inside_sphere * gradient_error).sum() / (relax_inside_sphere.sum() + 1e-5)
+
+        return {
+            'color': color,
+            'sdf': sdf,
+            'dists': dists,
+            'gradients': gradients.reshape(batch_size, n_samples, 3),
+            's_val': 1.0 / inv_s,
+            'mid_z_vals': mid_z_vals,
+            'weights': weights,
+            'cdf': c.reshape(batch_size, n_samples),
+            'gradient_error': gradient_error,
+            'inside_sphere': inside_sphere
+        }
+
+    def render(self, rays_o, rays_d, near, far, perturb_overwrite=-1, background_rgb=None, cos_anneal_ratio=0.0, use_gt_sdf=False):
+        batch_size = len(rays_o)
+        sample_dist = 2.0 / self.n_samples   # Assuming the region of interest is a unit sphere
+        z_vals = torch.linspace(0.0, 1.0, self.n_samples)
+        z_vals = near + (far - near) * z_vals[None, :]
+
+        z_vals_outside = None
+        if self.n_outside > 0:
+            z_vals_outside = torch.linspace(1e-3, 1.0 - 1.0 / (self.n_outside + 1.0), self.n_outside)
+
+        n_samples = self.n_samples
+        perturb = self.perturb
+
+        if perturb_overwrite >= 0:
+            perturb = perturb_overwrite
+        if perturb > 0:
+            t_rand = (torch.rand([batch_size, 1]) - 0.5)
+            z_vals = z_vals + t_rand * 2.0 / self.n_samples
+
+            if self.n_outside > 0:
+                mids = .5 * (z_vals_outside[..., 1:] + z_vals_outside[..., :-1])
+                upper = torch.cat([mids, z_vals_outside[..., -1:]], -1)
+                lower = torch.cat([z_vals_outside[..., :1], mids], -1)
+                t_rand = torch.rand([batch_size, z_vals_outside.shape[-1]])
+                z_vals_outside = lower[None, :] + (upper - lower)[None, :] * t_rand
+
+        if self.n_outside > 0:
+            z_vals_outside = far / torch.flip(z_vals_outside, dims=[-1]) + 1.0 / self.n_samples
+
+        background_alpha = None
+        background_sampled_color = None
+
+        # Up sample
+        if self.n_importance > 0:
+            with torch.no_grad():
+                pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]
+                
+                pts = (pts - self.minn_pts) / (self.maxx_pts - self.minn_pts)
+                # sdf = self.sdf_network.sdf(pts.reshape(-1, 3)).reshape(batch_size, self.n_samples)
+                # gt_sdf #
+                
+                # 
+                # pts = ((pts - xyz_min) / (xyz_max - xyz_min)).flip((-1,)) * 2 - 1
+                
+                # pts = pts.flip((-1,)) * 2 - 1
+                pts = pts * 2 - 1
+                
+                pts_exp = pts.reshape(-1, 3)
+                # minn_pts, _ = torch.min(pts_exp, dim=0)
+                # maxx_pts, _ = torch.max(pts_exp, dim=0)
+                # print(f"minn_pts: {minn_pts}, maxx_pts: {maxx_pts}")
+
+                # pts_to_near = pts - near.unsqueeze(1)
+                # maxx_pts = 1.5; minn_pts = -1.5
+                # # maxx_pts = 3; minn_pts = -3
+                # # maxx_pts = 1; minn_pts = -1
+                # pts_exp = (pts_exp - minn_pts) / (maxx_pts - minn_pts)
+                
+                ## render and iamges  ####  
+                if use_gt_sdf:
+                    ### use the GT sdf field ####
+                    # print(f"Using gt sdf :")
+                    sdf = self.gt_sdf(pts_exp.reshape(-1, 3).detach().cpu().numpy())
+                    sdf = torch.from_numpy(sdf).float().cuda()
+                    sdf = sdf.reshape(batch_size, self.n_samples)
+                    ### use the GT sdf field ####
+                else:
+                    #### use the optimized sdf field ####
+                    sdf = self.sdf_network.sdf(pts_exp).reshape(batch_size, self.n_samples)
+                    #### use the optimized sdf field ####
+
+                for i in range(self.up_sample_steps):
+                    new_z_vals = self.up_sample(rays_o,
+                                                rays_d,
+                                                z_vals,
+                                                sdf,
+                                                self.n_importance // self.up_sample_steps,
+                                                64 * 2**i)
+                    z_vals, sdf = self.cat_z_vals(rays_o,
+                                                  rays_d,
+                                                  z_vals,
+                                                  new_z_vals,
+                                                  sdf,
+                                                  last=(i + 1 == self.up_sample_steps))
+
+            n_samples = self.n_samples + self.n_importance
+
+        # Background model
+        if self.n_outside > 0:
+            z_vals_feed = torch.cat([z_vals, z_vals_outside], dim=-1)
+            z_vals_feed, _ = torch.sort(z_vals_feed, dim=-1)
+            ret_outside = self.render_core_outside(rays_o, rays_d, z_vals_feed, sample_dist, self.nerf)
+
+            background_sampled_color = ret_outside['sampled_color']
+            background_alpha = ret_outside['alpha']
+
+        # Render core
+        ret_fine = self.render_core(rays_o, # 
+                                    rays_d,
+                                    z_vals,
+                                    sample_dist,
+                                    self.sdf_network,
+                                    self.deviation_network,
+                                    self.color_network,
+                                    background_rgb=background_rgb,
+                                    background_alpha=background_alpha,
+                                    background_sampled_color=background_sampled_color,
+                                    cos_anneal_ratio=cos_anneal_ratio)
+
+        color_fine = ret_fine['color']
+        weights = ret_fine['weights']
+        weights_sum = weights.sum(dim=-1, keepdim=True)
+        gradients = ret_fine['gradients']
+        s_val = ret_fine['s_val'].reshape(batch_size, n_samples).mean(dim=-1, keepdim=True)
+
+        return {
+            'color_fine': color_fine,
+            's_val': s_val,
+            'cdf_fine': ret_fine['cdf'],
+            'weight_sum': weights_sum,
+            'weight_max': torch.max(weights, dim=-1, keepdim=True)[0],
+            'gradients': gradients,
+            'weights': weights,
+            'gradient_error': ret_fine['gradient_error'],
+            'inside_sphere': ret_fine['inside_sphere']
+        }
+
+    def extract_geometry(self, bound_min, bound_max, resolution, threshold=0.0):
+        return extract_geometry(bound_min, # extract geometry #
+                                bound_max,
+                                resolution=resolution,
+                                threshold=threshold,
+                                query_func=lambda pts: -self.sdf_network.sdf(pts))
+        
+    def extract_geometry_tets(self, bound_min, bound_max, resolution, threshold=0.0):
+        return extract_geometry_tets(bound_min, # extract geometry #
+                                bound_max,
+                                resolution=resolution,
+                                threshold=threshold,
+                                query_func=lambda pts: -self.sdf_network.sdf(pts))

models/renderer_def.py

@@ -0,0 +1,725 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import logging
+import mcubes
+from icecream import ic
+import os
+
+import trimesh
+from pysdf import SDF
+
+from uni_rep.rep_3d.dmtet import marching_tets_tetmesh, create_tetmesh_variables
+
+def create_mt_variable(device):
+    triangle_table = torch.tensor(
+        [
+            [-1, -1, -1, -1, -1, -1],
+            [1, 0, 2, -1, -1, -1],
+            [4, 0, 3, -1, -1, -1],
+            [1, 4, 2, 1, 3, 4],
+            [3, 1, 5, -1, -1, -1],
+            [2, 3, 0, 2, 5, 3],
+            [1, 4, 0, 1, 5, 4],
+            [4, 2, 5, -1, -1, -1],
+            [4, 5, 2, -1, -1, -1],
+            [4, 1, 0, 4, 5, 1],
+            [3, 2, 0, 3, 5, 2],
+            [1, 3, 5, -1, -1, -1],
+            [4, 1, 2, 4, 3, 1],
+            [3, 0, 4, -1, -1, -1],
+            [2, 0, 1, -1, -1, -1],
+            [-1, -1, -1, -1, -1, -1]
+        ], dtype=torch.long, device=device)
+
+    num_triangles_table = torch.tensor([0, 1, 1, 2, 1, 2, 2, 1, 1, 2, 2, 1, 2, 1, 1, 0], dtype=torch.long, device=device)
+    base_tet_edges = torch.tensor([0, 1, 0, 2, 0, 3, 1, 2, 1, 3, 2, 3], dtype=torch.long, device=device)
+    v_id = torch.pow(2, torch.arange(4, dtype=torch.long, device=device))
+    return triangle_table, num_triangles_table, base_tet_edges, v_id
+
+
+
+def  extract_fields_from_tets(bound_min, bound_max, resolution, query_func, def_func=None):
+    # load tet via resolution #
+    # scale them via bounds #
+    # extract the geometry #
+    # /home/xueyi/gen/DeepMetaHandles/data/tets/100_compress.npz # strange #
+    device = bound_min.device
+    # if resolution in [64, 70, 80, 90, 100]:
+    #     tet_fn = f"/home/xueyi/gen/DeepMetaHandles/data/tets/{resolution}_compress.npz"
+    # else:
+    tet_fn = f"/home/xueyi/gen/DeepMetaHandles/data/tets/{100}_compress.npz"
+    tets = np.load(tet_fn)
+    verts = torch.from_numpy(tets['vertices']).float().to(device) # verts positions 
+    indices = torch.from_numpy(tets['tets']).long().to(device) # .to(self.device)
+    # split #
+    # verts; verts; #
+    minn_verts, _ = torch.min(verts, dim=0)
+    maxx_verts, _ = torch.max(verts, dim=0) # (3, ) # exporting the  
+    # scale_verts = maxx_verts - minn_verts
+    scale_bounds = bound_max - bound_min # scale bounds #
+    
+    ### scale the vertices ###
+    scaled_verts = (verts - minn_verts.unsqueeze(0)) / (maxx_verts - minn_verts).unsqueeze(0) ### the maxx and minn verts scales ###
+    
+    # scaled_verts = (verts - minn_verts.unsqueeze(0)) / (maxx_verts - minn_verts).unsqueeze(0) ### the maxx and minn verts scales ###
+
+    scaled_verts = scaled_verts * 2. - 1. # init the sdf filed viathe tet mesh vertices and the sdf values ##
+    # scaled_verts = (scaled_verts * scale_bounds.unsqueeze(0)) + bound_min.unsqueeze(0) ## the scaled verts ###
+    
+    # scaled_verts = scaled_verts - scale_bounds.unsqueeze(0) / 2. # 
+    # scaled_verts = scaled_verts -  bound_min.unsqueeze(0) - scale_bounds.unsqueeze(0) / 2.
+    
+    sdf_values = []
+    N = 64
+    query_bundles = N ** 3 ### N^3
+    query_NNs = scaled_verts.size(0) // query_bundles
+    if query_NNs * query_bundles < scaled_verts.size(0):
+        query_NNs += 1
+    for i_query in range(query_NNs):
+        cur_bundle_st = i_query * query_bundles
+        cur_bundle_ed = (i_query + 1) * query_bundles
+        cur_bundle_ed = min(cur_bundle_ed, scaled_verts.size(0))
+        cur_query_pts = scaled_verts[cur_bundle_st: cur_bundle_ed]
+        if def_func is not None:
+            cur_query_pts = def_func(cur_query_pts)
+        cur_query_vals = query_func(cur_query_pts)
+        sdf_values.append(cur_query_vals)
+    sdf_values = torch.cat(sdf_values, dim=0)
+    # print(f"queryed sdf values: {sdf_values.size()}") #
+    
+    GT_sdf_values = np.load("/home/xueyi/diffsim/DiffHand/assets/hand/100_sdf_values.npy", allow_pickle=True)
+    GT_sdf_values = torch.from_numpy(GT_sdf_values).float().to(device)
+    
+    # intrinsic, tet values, pts values, sdf network #
+    triangle_table, num_triangles_table, base_tet_edges, v_id = create_mt_variable(device)
+    tet_table, num_tets_table = create_tetmesh_variables(device)
+    
+    sdf_values = sdf_values.squeeze(-1) # how the rendering # 
+    
+    # print(f"GT_sdf_values: {GT_sdf_values.size()}, sdf_values: {sdf_values.size()}, scaled_verts: {scaled_verts.size()}")
+    # print(f"scaled_verts: {scaled_verts.size()}, ")
+    # pos_nx3, sdf_n, tet_fx4, triangle_table, num_triangles_table, base_tet_edges, v_id,
+    # return_tet_mesh=False, ori_v=None, num_tets_table=None, tet_table=None):
+    # marching_tets_tetmesh ##
+    verts, faces, tet_verts, tets = marching_tets_tetmesh(scaled_verts, sdf_values, indices, triangle_table, num_triangles_table, base_tet_edges, v_id, return_tet_mesh=True, ori_v=scaled_verts, num_tets_table=num_tets_table, tet_table=tet_table)
+    ### use the GT sdf values for the marching tets ###
+    GT_verts,  GT_faces,  GT_tet_verts,  GT_tets = marching_tets_tetmesh(scaled_verts, GT_sdf_values, indices, triangle_table, num_triangles_table, base_tet_edges, v_id, return_tet_mesh=True, ori_v=scaled_verts, num_tets_table=num_tets_table, tet_table=tet_table)
+    
+    # print(f"After tet marching with verts: {verts.size()}, faces: {faces.size()}")
+    return verts, faces, sdf_values, GT_verts, GT_faces  # verts, faces #
+
+def extract_fields(bound_min, bound_max, resolution, query_func):
+    N = 64
+    X = torch.linspace(bound_min[0], bound_max[0], resolution).split(N)
+    Y = torch.linspace(bound_min[1], bound_max[1], resolution).split(N)
+    Z = torch.linspace(bound_min[2], bound_max[2], resolution).split(N)
+
+    u = np.zeros([resolution, resolution, resolution], dtype=np.float32)
+    with torch.no_grad():
+        for xi, xs in enumerate(X):
+            for yi, ys in enumerate(Y):
+                for zi, zs in enumerate(Z):
+                    xx, yy, zz = torch.meshgrid(xs, ys, zs)
+                    pts = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1)
+                    val = query_func(pts).reshape(len(xs), len(ys), len(zs)).detach().cpu().numpy()
+                    u[xi * N: xi * N + len(xs), yi * N: yi * N + len(ys), zi * N: zi * N + len(zs)] = val
+        # should save u here #
+        # save_u_path = os.path.join("/data2/datasets/diffsim/neus/exp/hand_test/womask_sphere_reverse_value/other_saved", "sdf_values.npy")
+        # np.save(save_u_path, u)
+        # print(f"u saved to {save_u_path}")
+    return u
+
+
+def extract_geometry(bound_min, bound_max, resolution, threshold, query_func):
+    print('threshold: {}'.format(threshold))
+    
+    ## using maching cubes ###
+    u = extract_fields(bound_min, bound_max, resolution, query_func)
+    vertices, triangles = mcubes.marching_cubes(u, threshold) # grid sdf and marching cubes #
+    b_max_np = bound_max.detach().cpu().numpy()
+    b_min_np = bound_min.detach().cpu().numpy()
+
+    vertices = vertices / (resolution - 1.0) * (b_max_np - b_min_np)[None, :] + b_min_np[None, :]
+    ### using maching cubes ###
+    
+    ### using marching tets ###
+    # vertices, triangles = extract_fields_from_tets(bound_min, bound_max, resolution, query_func)
+    # vertices = vertices.detach().cpu().numpy()
+    # triangles = triangles.detach().cpu().numpy()
+    ### using marching tets ###
+    
+    # b_max_np = bound_max.detach().cpu().numpy()
+    # b_min_np = bound_min.detach().cpu().numpy()
+
+    # vertices = vertices / (resolution - 1.0) * (b_max_np - b_min_np)[None, :] + b_min_np[None, :]
+    return vertices, triangles
+
+def extract_geometry_tets(bound_min, bound_max, resolution, threshold, query_func, def_func=None):
+    # print('threshold: {}'.format(threshold))
+    
+    ### using maching cubes ###
+    # u = extract_fields(bound_min, bound_max, resolution, query_func)
+    # vertices, triangles = mcubes.marching_cubes(u, threshold) # grid sdf and marching cubes #
+    # b_max_np = bound_max.detach().cpu().numpy()
+    # b_min_np = bound_min.detach().cpu().numpy()
+
+    # vertices = vertices / (resolution - 1.0) * (b_max_np - b_min_np)[None, :] + b_min_np[None, :]
+    ### using maching cubes ###
+    
+    ## 
+    ### using marching tets ### fiels from tets ##
+    vertices, triangles, tet_sdf_values,  GT_verts, GT_faces  = extract_fields_from_tets(bound_min, bound_max, resolution, query_func, def_func=def_func)
+    # vertices = vertices.detach().cpu().numpy()
+    # triangles = triangles.detach().cpu().numpy()
+    ### using marching tets ###
+    
+    # b_max_np = bound_max.detach().cpu().numpy()
+    # b_min_np = bound_min.detach().cpu().numpy()
+    #
+
+    # vertices = vertices / (resolution - 1.0) * (b_max_np - b_min_np)[None, :] + b_min_np[None, :]
+    return vertices, triangles, tet_sdf_values,  GT_verts, GT_faces 
+
+
+def sample_pdf(bins, weights, n_samples, det=False):
+    # This implementation is from NeRF
+    # Get pdf
+    weights = weights + 1e-5  # prevent nans
+    pdf = weights / torch.sum(weights, -1, keepdim=True)
+    cdf = torch.cumsum(pdf, -1)
+    cdf = torch.cat([torch.zeros_like(cdf[..., :1]), cdf], -1)
+    # Take uniform samples
+    if det:
+        u = torch.linspace(0. + 0.5 / n_samples, 1. - 0.5 / n_samples, steps=n_samples)
+        u = u.expand(list(cdf.shape[:-1]) + [n_samples])
+    else:
+        u = torch.rand(list(cdf.shape[:-1]) + [n_samples])
+
+    # Invert CDF # invert cdf # 
+    u = u.contiguous()
+    inds = torch.searchsorted(cdf, u, right=True)
+    below = torch.max(torch.zeros_like(inds - 1), inds - 1)
+    above = torch.min((cdf.shape[-1] - 1) * torch.ones_like(inds), inds)
+    inds_g = torch.stack([below, above], -1)  # (batch, N_samples, 2)
+
+    matched_shape = [inds_g.shape[0], inds_g.shape[1], cdf.shape[-1]]
+    cdf_g = torch.gather(cdf.unsqueeze(1).expand(matched_shape), 2, inds_g)
+    bins_g = torch.gather(bins.unsqueeze(1).expand(matched_shape), 2, inds_g)
+
+    denom = (cdf_g[..., 1] - cdf_g[..., 0])
+    denom = torch.where(denom < 1e-5, torch.ones_like(denom), denom)
+    t = (u - cdf_g[..., 0]) / denom
+    samples = bins_g[..., 0] + t * (bins_g[..., 1] - bins_g[..., 0])
+
+    return samples
+
+
+def load_GT_vertices(GT_meshes_folder):
+    tot_meshes_fns = os.listdir(GT_meshes_folder)
+    tot_meshes_fns = [fn for fn in tot_meshes_fns if fn.endswith(".obj")]
+    tot_mesh_verts = []
+    tot_mesh_faces = []
+    n_tot_verts = 0
+    for fn in tot_meshes_fns:
+        cur_mesh_fn = os.path.join(GT_meshes_folder, fn)
+        obj_mesh = trimesh.load(cur_mesh_fn, process=False)
+        # obj_mesh.remove_degenerate_faces(height=1e-06)
+
+        verts_obj = np.array(obj_mesh.vertices)
+        faces_obj = np.array(obj_mesh.faces)
+        
+        tot_mesh_verts.append(verts_obj)
+        tot_mesh_faces.append(faces_obj + n_tot_verts)
+        n_tot_verts += verts_obj.shape[0]
+        
+        # tot_mesh_faces.append(faces_obj)
+    tot_mesh_verts = np.concatenate(tot_mesh_verts, axis=0)
+    tot_mesh_faces = np.concatenate(tot_mesh_faces, axis=0)
+    return tot_mesh_verts, tot_mesh_faces
+
+
+class NeuSRenderer:
+    def __init__(self,
+                 nerf,
+                 sdf_network,
+                 deviation_network,
+                 color_network,
+                 n_samples,
+                 n_importance,
+                 n_outside,
+                 up_sample_steps,
+                 perturb):
+        self.nerf = nerf
+        self.sdf_network = sdf_network
+        self.deviation_network = deviation_network
+        self.color_network = color_network
+        self.n_samples = n_samples
+        self.n_importance = n_importance
+        self.n_outside = n_outside
+        self.up_sample_steps = up_sample_steps
+        self.perturb = perturb
+        
+        GT_meshes_folder = "/home/xueyi/diffsim/DiffHand/assets/hand"
+        self.mesh_vertices, self.mesh_faces = load_GT_vertices(GT_meshes_folder=GT_meshes_folder)
+        maxx_pts = 25.
+        minn_pts = -15.
+        self.mesh_vertices = (self.mesh_vertices - minn_pts) / (maxx_pts - minn_pts)
+        f = SDF(self.mesh_vertices, self.mesh_faces)
+        self.gt_sdf = f ## a unite sphere or box
+        
+        self.minn_pts = 0
+        self.maxx_pts = 1.
+        
+        # self.minn_pts = -1.5 # gorudn-truth states with the deformation -> update the sdf value fiedl 
+        # self.maxx_pts = 1.5
+        self.bkg_pts = ... # TODO: the bkg pts # bkg_pts; # bkg_pts_defs #
+        self.cur_fr_bkg_pts_defs = ... # TODO: set the cur_bkg_pts_defs for each frame #
+        self.dist_interp_thres = ... # TODO: set the cur_bkg_pts_defs #
+        
+        self.bending_network = ... # TODO: add the bending network # 
+        self.use_bending_network = ... # TODO: set the property #
+        self.use_delta_bending = ... # TODO
+        # use bending network # 
+        
+        
+    # get the pts and render the pts #
+    # pts and the rendering pts #
+    def deform_pts(self, pts, pts_ts=0):
+        
+        if self.use_bending_network:
+            if len(pts.size()) == 3:
+                nnb, nns = pts.size(0), pts.size(1)
+                pts_exp = pts.contiguous().view(nnb * nns, -1).contiguous()
+            else:
+                pts_exp = pts
+            # pts_ts #
+            if self.use_delta_bending:
+                # if pts_ts >= 5:
+                #     pts_exp = self.bending_network(pts_exp, input_pts_ts=pts_ts)
+                #     for cur_pts_ts in range(4, -1, -1):
+                #         # print(f"using delta bending with pts_ts: {cur_pts_ts}")
+                #         pts_exp = self.bending_network(pts_exp, input_pts_ts=cur_pts_ts)
+                # else:
+                #     for cur_pts_ts in range(pts_ts, -1, -1):
+                #         # print(f"using delta bending with pts_ts: {cur_pts_ts}")
+                #         pts_exp = self.bending_network(pts_exp, input_pts_ts=cur_pts_ts)
+                for cur_pts_ts in range(pts_ts, -1, -1):
+                    # print(f"using delta bending with pts_ts: {cur_pts_ts}")
+                    pts_exp = self.bending_network(pts_exp, input_pts_ts=cur_pts_ts)
+            else:
+                pts_exp = self.bending_network(pts_exp, input_pts_ts=pts_ts)
+            if len(pts.size()) == 3:
+                pts = pts_exp.contiguous().view(nnb, nns, -1).contiguous()
+            else:
+                pts = pts_exp
+            return pts
+            
+        # pts: nn_batch x nn_samples x 3
+        if len(pts.size()) == 3:
+            nnb, nns = pts.size(0), pts.size(1)
+            pts_exp = pts.contiguous().view(nnb * nns, -1).contiguous()
+        else:
+            pts_exp = pts
+        # print(f"prior to deforming: {pts.size()}")
+        
+        dist_pts_to_bkg_pts = torch.sum(
+            (pts_exp.unsqueeze(1) - self.bkg_pts.unsqueeze(0)) ** 2, dim=-1 ## nn_pts_exp x nn_bkg_pts
+        )
+        dist_mask = dist_pts_to_bkg_pts <= self.dist_interp_thres # 
+        dist_mask_float = dist_mask.float()
+        
+        # dist_mask_float #
+        cur_fr_bkg_def_exp = self.cur_fr_bkg_pts_defs.unsqueeze(0).repeat(pts_exp.size(0), 1, 1).contiguous()
+        cur_fr_pts_def = torch.sum(
+            cur_fr_bkg_def_exp * dist_mask_float.unsqueeze(-1), dim=1
+        )
+        dist_mask_float_summ = torch.sum(
+            dist_mask_float, dim=1
+        )
+        dist_mask_float_summ = torch.clamp(dist_mask_float_summ, min=1)
+        cur_fr_pts_def = cur_fr_pts_def / dist_mask_float_summ.unsqueeze(-1) # bkg pts deformation #
+        pts_exp = pts_exp - cur_fr_pts_def
+        if len(pts.size()) == 3:
+            pts = pts_exp.contiguous().view(nnb, nns, -1).contiguous()
+        else:
+            pts = pts_exp
+        return pts # 
+    
+    
+        
+
+    def render_core_outside(self, rays_o, rays_d, z_vals, sample_dist, nerf, background_rgb=None, pts_ts=0):
+        """
+        Render background
+        """
+        batch_size, n_samples = z_vals.shape
+
+        # Section length
+        dists = z_vals[..., 1:] - z_vals[..., :-1]
+        dists = torch.cat([dists, torch.Tensor([sample_dist]).expand(dists[..., :1].shape)], -1)
+        mid_z_vals = z_vals + dists * 0.5
+
+        # Section midpoints #
+        pts = rays_o[:, None, :] + rays_d[:, None, :] * mid_z_vals[..., :, None]  # batch_size, n_samples, 3 #
+        
+        # pts = pts.flip((-1,)) * 2 - 1
+        pts = pts * 2 - 1
+        
+        pts = self.deform_pts(pts=pts, pts_ts=pts_ts)
+
+        dis_to_center = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=True).clip(1.0, 1e10)
+        pts = torch.cat([pts / dis_to_center, 1.0 / dis_to_center], dim=-1)       # batch_size, n_samples, 4 #
+
+        dirs = rays_d[:, None, :].expand(batch_size, n_samples, 3)
+
+        pts = pts.reshape(-1, 3 + int(self.n_outside > 0))
+        dirs = dirs.reshape(-1, 3)
+
+        density, sampled_color = nerf(pts, dirs)
+        sampled_color = torch.sigmoid(sampled_color)
+        alpha = 1.0 - torch.exp(-F.softplus(density.reshape(batch_size, n_samples)) * dists)
+        alpha = alpha.reshape(batch_size, n_samples)
+        weights = alpha * torch.cumprod(torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[:, :-1]
+        sampled_color = sampled_color.reshape(batch_size, n_samples, 3)
+        color = (weights[:, :, None] * sampled_color).sum(dim=1)
+        if background_rgb is not None:
+            color = color + background_rgb * (1.0 - weights.sum(dim=-1, keepdim=True))
+
+        return {
+            'color': color,
+            'sampled_color': sampled_color,
+            'alpha': alpha,
+            'weights': weights,
+        }
+
+    def up_sample(self, rays_o, rays_d, z_vals, sdf, n_importance, inv_s, pts_ts=0):
+        """
+        Up sampling give a fixed inv_s
+        """
+        batch_size, n_samples = z_vals.shape
+        pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]  # n_rays, n_samples, 3
+        
+        # pts = pts.flip((-1,)) * 2 - 1
+        pts = pts * 2 - 1
+        
+        pts = self.deform_pts(pts=pts, pts_ts=pts_ts)
+        
+        radius = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=False)
+        inside_sphere = (radius[:, :-1] < 1.0) | (radius[:, 1:] < 1.0)
+        sdf = sdf.reshape(batch_size, n_samples)
+        prev_sdf, next_sdf = sdf[:, :-1], sdf[:, 1:]
+        prev_z_vals, next_z_vals = z_vals[:, :-1], z_vals[:, 1:]
+        mid_sdf = (prev_sdf + next_sdf) * 0.5
+        cos_val = (next_sdf - prev_sdf) / (next_z_vals - prev_z_vals + 1e-5)
+
+        # ----------------------------------------------------------------------------------------------------------
+        # Use min value of [ cos, prev_cos ]
+        # Though it makes the sampling (not rendering) a little bit biased, this strategy can make the sampling more
+        # robust when meeting situations like below:
+        #
+        # SDF
+        # ^
+        # |\          -----x----...
+        # | \        /
+        # |  x      x
+        # |---\----/-------------> 0 level
+        # |    \  /
+        # |     \/
+        # |
+        # ----------------------------------------------------------------------------------------------------------
+        prev_cos_val = torch.cat([torch.zeros([batch_size, 1]), cos_val[:, :-1]], dim=-1)
+        cos_val = torch.stack([prev_cos_val, cos_val], dim=-1)
+        cos_val, _ = torch.min(cos_val, dim=-1, keepdim=False)
+        cos_val = cos_val.clip(-1e3, 0.0) * inside_sphere
+
+        dist = (next_z_vals - prev_z_vals)
+        prev_esti_sdf = mid_sdf - cos_val * dist * 0.5
+        next_esti_sdf = mid_sdf + cos_val * dist * 0.5
+        prev_cdf = torch.sigmoid(prev_esti_sdf * inv_s)
+        next_cdf = torch.sigmoid(next_esti_sdf * inv_s)
+        alpha = (prev_cdf - next_cdf + 1e-5) / (prev_cdf + 1e-5)
+        weights = alpha * torch.cumprod(
+            torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[:, :-1]
+
+        z_samples = sample_pdf(z_vals, weights, n_importance, det=True).detach()
+        return z_samples
+
+    def cat_z_vals(self, rays_o, rays_d, z_vals, new_z_vals, sdf, last=False, pts_ts=0):
+        batch_size, n_samples = z_vals.shape
+        _, n_importance = new_z_vals.shape
+        pts = rays_o[:, None, :] + rays_d[:, None, :] * new_z_vals[..., :, None]
+        
+        # pts = pts.flip((-1,)) * 2 - 1
+        pts = pts * 2 - 1
+        
+        pts = self.deform_pts(pts=pts, pts_ts=pts_ts)
+        
+        z_vals = torch.cat([z_vals, new_z_vals], dim=-1)
+        z_vals, index = torch.sort(z_vals, dim=-1)
+
+        if not last:
+            new_sdf = self.sdf_network.sdf(pts.reshape(-1, 3)).reshape(batch_size, n_importance)
+            sdf = torch.cat([sdf, new_sdf], dim=-1)
+            xx = torch.arange(batch_size)[:, None].expand(batch_size, n_samples + n_importance).reshape(-1)
+            index = index.reshape(-1)
+            sdf = sdf[(xx, index)].reshape(batch_size, n_samples + n_importance)
+
+        return z_vals, sdf
+
+    def render_core(self,
+                    rays_o,
+                    rays_d,
+                    z_vals,
+                    sample_dist,
+                    sdf_network,
+                    deviation_network,
+                    color_network,
+                    background_alpha=None,
+                    background_sampled_color=None,
+                    background_rgb=None,
+                    cos_anneal_ratio=0.0,
+                    pts_ts=0):
+        batch_size, n_samples = z_vals.shape
+
+        # Section length
+        dists = z_vals[..., 1:] - z_vals[..., :-1]
+        dists = torch.cat([dists, torch.Tensor([sample_dist]).expand(dists[..., :1].shape)], -1)
+        mid_z_vals = z_vals + dists * 0.5 # z_vals and dists * 0.5 #
+
+        # Section midpoints
+        pts = rays_o[:, None, :] + rays_d[:, None, :] * mid_z_vals[..., :, None]  # n_rays, n_samples, 3
+        dirs = rays_d[:, None, :].expand(pts.shape)
+
+        pts = pts.reshape(-1, 3) # pts, nn_ou
+        dirs = dirs.reshape(-1, 3)
+        
+        pts = (pts - self.minn_pts) / (self.maxx_pts - self.minn_pts)
+        
+        # pts = pts.flip((-1,)) * 2 - 1
+        pts = pts * 2 - 1
+        
+        
+        pts = self.deform_pts(pts=pts, pts_ts=pts_ts)
+
+        sdf_nn_output = sdf_network(pts)
+        sdf = sdf_nn_output[:, :1]
+        feature_vector = sdf_nn_output[:, 1:]
+
+        gradients = sdf_network.gradient(pts).squeeze()
+        sampled_color = color_network(pts, gradients, dirs, feature_vector).reshape(batch_size, n_samples, 3)
+
+        # deviation network #
+        inv_s = deviation_network(torch.zeros([1, 3]))[:, :1].clip(1e-6, 1e6)           # Single parameter
+        inv_s = inv_s.expand(batch_size * n_samples, 1)
+
+        true_cos = (dirs * gradients).sum(-1, keepdim=True)
+
+        # "cos_anneal_ratio" grows from 0 to 1 in the beginning training iterations. The anneal strategy below makes
+        # the cos value "not dead" at the beginning training iterations, for better convergence.
+        iter_cos = -(F.relu(-true_cos * 0.5 + 0.5) * (1.0 - cos_anneal_ratio) +
+                     F.relu(-true_cos) * cos_anneal_ratio)  # always non-positive
+
+        # Estimate signed distances at section points
+        estimated_next_sdf = sdf + iter_cos * dists.reshape(-1, 1) * 0.5
+        estimated_prev_sdf = sdf - iter_cos * dists.reshape(-1, 1) * 0.5
+
+        prev_cdf = torch.sigmoid(estimated_prev_sdf * inv_s)
+        next_cdf = torch.sigmoid(estimated_next_sdf * inv_s)
+
+        p = prev_cdf - next_cdf
+        c = prev_cdf
+
+        alpha = ((p + 1e-5) / (c + 1e-5)).reshape(batch_size, n_samples).clip(0.0, 1.0)
+
+        pts_norm = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=True).reshape(batch_size, n_samples)
+        inside_sphere = (pts_norm < 1.0).float().detach()
+        relax_inside_sphere = (pts_norm < 1.2).float().detach()
+
+        # Render with background
+        if background_alpha is not None:
+            alpha = alpha * inside_sphere + background_alpha[:, :n_samples] * (1.0 - inside_sphere)
+            alpha = torch.cat([alpha, background_alpha[:, n_samples:]], dim=-1)
+            sampled_color = sampled_color * inside_sphere[:, :, None] +\
+                            background_sampled_color[:, :n_samples] * (1.0 - inside_sphere)[:, :, None]
+            sampled_color = torch.cat([sampled_color, background_sampled_color[:, n_samples:]], dim=1)
+
+        weights = alpha * torch.cumprod(torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[:, :-1]
+        weights_sum = weights.sum(dim=-1, keepdim=True)
+
+        color = (sampled_color * weights[:, :, None]).sum(dim=1)
+        if background_rgb is not None:    # Fixed background, usually black
+            color = color + background_rgb * (1.0 - weights_sum)
+
+        # Eikonal loss
+        gradient_error = (torch.linalg.norm(gradients.reshape(batch_size, n_samples, 3), ord=2,
+                                            dim=-1) - 1.0) ** 2
+        gradient_error = (relax_inside_sphere * gradient_error).sum() / (relax_inside_sphere.sum() + 1e-5)
+
+        return {
+            'color': color,
+            'sdf': sdf,
+            'dists': dists,
+            'gradients': gradients.reshape(batch_size, n_samples, 3),
+            's_val': 1.0 / inv_s,
+            'mid_z_vals': mid_z_vals,
+            'weights': weights,
+            'cdf': c.reshape(batch_size, n_samples),
+            'gradient_error': gradient_error,
+            'inside_sphere': inside_sphere
+        }
+
+    def render(self, rays_o, rays_d, near, far, pts_ts=0, perturb_overwrite=-1, background_rgb=None, cos_anneal_ratio=0.0, use_gt_sdf=False):
+        batch_size = len(rays_o)
+        sample_dist = 2.0 / self.n_samples   # Assuming the region of interest is a unit sphere
+        z_vals = torch.linspace(0.0, 1.0, self.n_samples)
+        z_vals = near + (far - near) * z_vals[None, :]
+
+        z_vals_outside = None
+        if self.n_outside > 0:
+            z_vals_outside = torch.linspace(1e-3, 1.0 - 1.0 / (self.n_outside + 1.0), self.n_outside)
+
+        n_samples = self.n_samples
+        perturb = self.perturb
+
+        if perturb_overwrite >= 0:
+            perturb = perturb_overwrite
+        if perturb > 0:
+            t_rand = (torch.rand([batch_size, 1]) - 0.5)
+            z_vals = z_vals + t_rand * 2.0 / self.n_samples
+
+            if self.n_outside > 0:
+                mids = .5 * (z_vals_outside[..., 1:] + z_vals_outside[..., :-1])
+                upper = torch.cat([mids, z_vals_outside[..., -1:]], -1)
+                lower = torch.cat([z_vals_outside[..., :1], mids], -1)
+                t_rand = torch.rand([batch_size, z_vals_outside.shape[-1]])
+                z_vals_outside = lower[None, :] + (upper - lower)[None, :] * t_rand
+
+        if self.n_outside > 0:
+            z_vals_outside = far / torch.flip(z_vals_outside, dims=[-1]) + 1.0 / self.n_samples
+
+        background_alpha = None
+        background_sampled_color = None
+
+        # Up sample
+        if self.n_importance > 0:
+            with torch.no_grad():
+                pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]
+                
+                pts = (pts - self.minn_pts) / (self.maxx_pts - self.minn_pts)
+                # sdf = self.sdf_network.sdf(pts.reshape(-1, 3)).reshape(batch_size, self.n_samples)
+                # gt_sdf #
+                
+                # 
+                # pts = ((pts - xyz_min) / (xyz_max - xyz_min)).flip((-1,)) * 2 - 1
+                
+                # pts = pts.flip((-1,)) * 2 - 1
+                pts = pts * 2 - 1
+                
+                pts = self.deform_pts(pts=pts, pts_ts=pts_ts)
+                
+                pts_exp = pts.reshape(-1, 3)
+                # minn_pts, _ = torch.min(pts_exp, dim=0)
+                # maxx_pts, _ = torch.max(pts_exp, dim=0) # deformation field (not a rigid one) -> the meshes #
+                # print(f"minn_pts: {minn_pts}, maxx_pts: {maxx_pts}")
+
+                # pts_to_near = pts - near.unsqueeze(1)
+                # maxx_pts = 1.5; minn_pts = -1.5
+                # # maxx_pts = 3; minn_pts = -3
+                # # maxx_pts = 1; minn_pts = -1
+                # pts_exp = (pts_exp - minn_pts) / (maxx_pts - minn_pts)
+                
+                ## render and iamges  ####  
+                if use_gt_sdf:
+                    ### use the GT sdf field ####
+                    # print(f"Using gt sdf :")
+                    sdf = self.gt_sdf(pts_exp.reshape(-1, 3).detach().cpu().numpy())
+                    sdf = torch.from_numpy(sdf).float().cuda()
+                    sdf = sdf.reshape(batch_size, self.n_samples)
+                    ### use the GT sdf field ####
+                else:
+                    #### use the optimized sdf field ####
+                    sdf = self.sdf_network.sdf(pts_exp).reshape(batch_size, self.n_samples)
+                    #### use the optimized sdf field ####
+
+                for i in range(self.up_sample_steps):
+                    new_z_vals = self.up_sample(rays_o,
+                                                rays_d,
+                                                z_vals,
+                                                sdf,
+                                                self.n_importance // self.up_sample_steps,
+                                                64 * 2**i,
+                                                pts_ts=pts_ts)
+                    z_vals, sdf = self.cat_z_vals(rays_o,
+                                                  rays_d,
+                                                  z_vals,
+                                                  new_z_vals,
+                                                  sdf,
+                                                  last=(i + 1 == self.up_sample_steps),
+                                                  pts_ts=pts_ts)
+
+            n_samples = self.n_samples + self.n_importance
+
+        # Background model
+        if self.n_outside > 0:
+            z_vals_feed = torch.cat([z_vals, z_vals_outside], dim=-1)
+            z_vals_feed, _ = torch.sort(z_vals_feed, dim=-1)
+            ret_outside = self.render_core_outside(rays_o, rays_d, z_vals_feed, sample_dist, self.nerf, pts_ts=pts_ts)
+
+            background_sampled_color = ret_outside['sampled_color']
+            background_alpha = ret_outside['alpha']
+
+        # Render core
+        ret_fine = self.render_core(rays_o, # 
+                                    rays_d,
+                                    z_vals,
+                                    sample_dist,
+                                    self.sdf_network,
+                                    self.deviation_network,
+                                    self.color_network,
+                                    background_rgb=background_rgb,
+                                    background_alpha=background_alpha,
+                                    background_sampled_color=background_sampled_color,
+                                    cos_anneal_ratio=cos_anneal_ratio,
+                                    pts_ts=pts_ts)
+
+        color_fine = ret_fine['color']
+        weights = ret_fine['weights']
+        weights_sum = weights.sum(dim=-1, keepdim=True)
+        gradients = ret_fine['gradients']
+        s_val = ret_fine['s_val'].reshape(batch_size, n_samples).mean(dim=-1, keepdim=True)
+
+        return {
+            'color_fine': color_fine,
+            's_val': s_val,
+            'cdf_fine': ret_fine['cdf'],
+            'weight_sum': weights_sum,
+            'weight_max': torch.max(weights, dim=-1, keepdim=True)[0],
+            'gradients': gradients,
+            'weights': weights,
+            'gradient_error': ret_fine['gradient_error'],
+            'inside_sphere': ret_fine['inside_sphere']
+        }
+
+    def extract_geometry(self, bound_min, bound_max, resolution, threshold=0.0):
+        return extract_geometry(bound_min, # extract geometry #
+                                bound_max,
+                                resolution=resolution,
+                                threshold=threshold,
+                                query_func=lambda pts: -self.sdf_network.sdf(pts))
+        
+    def extract_geometry_tets(self, bound_min, bound_max, resolution, pts_ts=0, threshold=0.0, wdef=False):
+        if wdef:
+            return extract_geometry_tets(bound_min, # extract geometry #
+                                bound_max,
+                                resolution=resolution,
+                                threshold=threshold,
+                                query_func=lambda pts: -self.sdf_network.sdf(pts),
+                                def_func=lambda pts: self.deform_pts(pts, pts_ts=pts_ts))
+        else:
+            return extract_geometry_tets(bound_min, # extract geometry #
+                                    bound_max,
+                                    resolution=resolution,
+                                    threshold=threshold,
+                                    query_func=lambda pts: -self.sdf_network.sdf(pts))

models/renderer_def_multi_objs.py

@@ -0,0 +1,1088 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import logging
+import mcubes
+from icecream import ic
+import os
+
+import trimesh
+from pysdf import SDF
+
+import models.fields as fields
+
+from uni_rep.rep_3d.dmtet import marching_tets_tetmesh, create_tetmesh_variables
+
+def batched_index_select(values, indices, dim = 1):
+  value_dims = values.shape[(dim + 1):]
+  values_shape, indices_shape = map(lambda t: list(t.shape), (values, indices))
+  indices = indices[(..., *((None,) * len(value_dims)))]
+  indices = indices.expand(*((-1,) * len(indices_shape)), *value_dims)
+  value_expand_len = len(indices_shape) - (dim + 1)
+  values = values[(*((slice(None),) * dim), *((None,) * value_expand_len), ...)]
+
+  value_expand_shape = [-1] * len(values.shape)
+  expand_slice = slice(dim, (dim + value_expand_len))
+  value_expand_shape[expand_slice] = indices.shape[expand_slice]
+  values = values.expand(*value_expand_shape)
+
+  dim += value_expand_len
+  return values.gather(dim, indices)
+
+
+def create_mt_variable(device):
+    triangle_table = torch.tensor(
+        [
+            [-1, -1, -1, -1, -1, -1],
+            [1, 0, 2, -1, -1, -1],
+            [4, 0, 3, -1, -1, -1],
+            [1, 4, 2, 1, 3, 4],
+            [3, 1, 5, -1, -1, -1],
+            [2, 3, 0, 2, 5, 3],
+            [1, 4, 0, 1, 5, 4],
+            [4, 2, 5, -1, -1, -1],
+            [4, 5, 2, -1, -1, -1],
+            [4, 1, 0, 4, 5, 1],
+            [3, 2, 0, 3, 5, 2],
+            [1, 3, 5, -1, -1, -1],
+            [4, 1, 2, 4, 3, 1],
+            [3, 0, 4, -1, -1, -1],
+            [2, 0, 1, -1, -1, -1],
+            [-1, -1, -1, -1, -1, -1]
+        ], dtype=torch.long, device=device)
+
+    num_triangles_table = torch.tensor([0, 1, 1, 2, 1, 2, 2, 1, 1, 2, 2, 1, 2, 1, 1, 0], dtype=torch.long, device=device)
+    base_tet_edges = torch.tensor([0, 1, 0, 2, 0, 3, 1, 2, 1, 3, 2, 3], dtype=torch.long, device=device)
+    v_id = torch.pow(2, torch.arange(4, dtype=torch.long, device=device))
+    return triangle_table, num_triangles_table, base_tet_edges, v_id
+
+
+
+def  extract_fields_from_tets(bound_min, bound_max, resolution, query_func, def_func=None):
+    # load tet via resolution #
+    # scale them via bounds #
+    # extract the geometry #
+    # /home/xueyi/gen/DeepMetaHandles/data/tets/100_compress.npz # strange #
+    device = bound_min.device
+    # if resolution in [64, 70, 80, 90, 100]:
+    #     tet_fn = f"/home/xueyi/gen/DeepMetaHandles/data/tets/{resolution}_compress.npz"
+    # else:
+    tet_fn = f"/home/xueyi/gen/DeepMetaHandles/data/tets/{100}_compress.npz"
+    tets = np.load(tet_fn)
+    verts = torch.from_numpy(tets['vertices']).float().to(device) # verts positions 
+    indices = torch.from_numpy(tets['tets']).long().to(device) # .to(self.device)
+    # split #
+    # verts; verts; #
+    minn_verts, _ = torch.min(verts, dim=0)
+    maxx_verts, _ = torch.max(verts, dim=0) # (3, ) # exporting the  
+    # scale_verts = maxx_verts - minn_verts
+    scale_bounds = bound_max - bound_min # scale bounds #
+    
+    ### scale the vertices ###
+    scaled_verts = (verts - minn_verts.unsqueeze(0)) / (maxx_verts - minn_verts).unsqueeze(0) ### the maxx and minn verts scales ###
+    
+    # scaled_verts = (verts - minn_verts.unsqueeze(0)) / (maxx_verts - minn_verts).unsqueeze(0) ### the maxx and minn verts scales ###
+
+    scaled_verts = scaled_verts * 2. - 1. # init the sdf filed viathe tet mesh vertices and the sdf values ##
+    # scaled_verts = (scaled_verts * scale_bounds.unsqueeze(0)) + bound_min.unsqueeze(0) ## the scaled verts ###
+    
+    # scaled_verts = scaled_verts - scale_bounds.unsqueeze(0) / 2. # 
+    # scaled_verts = scaled_verts -  bound_min.unsqueeze(0) - scale_bounds.unsqueeze(0) / 2.
+    
+    sdf_values = []
+    N = 64
+    query_bundles = N ** 3 ### N^3
+    query_NNs = scaled_verts.size(0) // query_bundles
+    if query_NNs * query_bundles < scaled_verts.size(0):
+        query_NNs += 1
+    for i_query in range(query_NNs):
+        cur_bundle_st = i_query * query_bundles
+        cur_bundle_ed = (i_query + 1) * query_bundles
+        cur_bundle_ed = min(cur_bundle_ed, scaled_verts.size(0))
+        cur_query_pts = scaled_verts[cur_bundle_st: cur_bundle_ed]
+        if def_func is not None:
+            cur_query_pts = def_func(cur_query_pts)
+        cur_query_vals = query_func(cur_query_pts)
+        sdf_values.append(cur_query_vals)
+    sdf_values = torch.cat(sdf_values, dim=0)
+    # print(f"queryed sdf values: {sdf_values.size()}") #
+    
+    GT_sdf_values = np.load("/home/xueyi/diffsim/DiffHand/assets/hand/100_sdf_values.npy", allow_pickle=True)
+    GT_sdf_values = torch.from_numpy(GT_sdf_values).float().to(device)
+    
+    # intrinsic, tet values, pts values, sdf network #
+    triangle_table, num_triangles_table, base_tet_edges, v_id = create_mt_variable(device)
+    tet_table, num_tets_table = create_tetmesh_variables(device)
+    
+    sdf_values = sdf_values.squeeze(-1) # how the rendering # 
+    
+    # print(f"GT_sdf_values: {GT_sdf_values.size()}, sdf_values: {sdf_values.size()}, scaled_verts: {scaled_verts.size()}")
+    # print(f"scaled_verts: {scaled_verts.size()}, ")
+    # pos_nx3, sdf_n, tet_fx4, triangle_table, num_triangles_table, base_tet_edges, v_id,
+    # return_tet_mesh=False, ori_v=None, num_tets_table=None, tet_table=None):
+    # marching_tets_tetmesh ##
+    verts, faces, tet_verts, tets = marching_tets_tetmesh(scaled_verts, sdf_values, indices, triangle_table, num_triangles_table, base_tet_edges, v_id, return_tet_mesh=True, ori_v=scaled_verts, num_tets_table=num_tets_table, tet_table=tet_table)
+    ### use the GT sdf values for the marching tets ###
+    GT_verts,  GT_faces,  GT_tet_verts,  GT_tets = marching_tets_tetmesh(scaled_verts, GT_sdf_values, indices, triangle_table, num_triangles_table, base_tet_edges, v_id, return_tet_mesh=True, ori_v=scaled_verts, num_tets_table=num_tets_table, tet_table=tet_table)
+    
+    # print(f"After tet marching with verts: {verts.size()}, faces: {faces.size()}")
+    return verts, faces, sdf_values, GT_verts, GT_faces  # verts, faces #
+
+def extract_fields(bound_min, bound_max, resolution, query_func):
+    N = 64
+    X = torch.linspace(bound_min[0], bound_max[0], resolution).split(N)
+    Y = torch.linspace(bound_min[1], bound_max[1], resolution).split(N)
+    Z = torch.linspace(bound_min[2], bound_max[2], resolution).split(N)
+
+    u = np.zeros([resolution, resolution, resolution], dtype=np.float32)
+    with torch.no_grad():
+        for xi, xs in enumerate(X):
+            for yi, ys in enumerate(Y):
+                for zi, zs in enumerate(Z):
+                    xx, yy, zz = torch.meshgrid(xs, ys, zs)
+                    pts = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1)
+                    val = query_func(pts).reshape(len(xs), len(ys), len(zs)).detach().cpu().numpy()
+                    u[xi * N: xi * N + len(xs), yi * N: yi * N + len(ys), zi * N: zi * N + len(zs)] = val
+        # should save u here #
+        # save_u_path = os.path.join("/data2/datasets/diffsim/neus/exp/hand_test/womask_sphere_reverse_value/other_saved", "sdf_values.npy")
+        # np.save(save_u_path, u)
+        # print(f"u saved to {save_u_path}")
+    return u
+
+
+def extract_geometry(bound_min, bound_max, resolution, threshold, query_func):
+    print('threshold: {}'.format(threshold))
+    
+    ## using maching cubes ###
+    u = extract_fields(bound_min, bound_max, resolution, query_func)
+    vertices, triangles = mcubes.marching_cubes(u, threshold) # grid sdf and marching cubes #
+    b_max_np = bound_max.detach().cpu().numpy()
+    b_min_np = bound_min.detach().cpu().numpy()
+
+    vertices = vertices / (resolution - 1.0) * (b_max_np - b_min_np)[None, :] + b_min_np[None, :]
+    ### using maching cubes ###
+    
+    ### using marching tets ###
+    # vertices, triangles = extract_fields_from_tets(bound_min, bound_max, resolution, query_func)
+    # vertices = vertices.detach().cpu().numpy()
+    # triangles = triangles.detach().cpu().numpy()
+    ### using marching tets ###
+    
+    # b_max_np = bound_max.detach().cpu().numpy()
+    # b_min_np = bound_min.detach().cpu().numpy()
+
+    # vertices = vertices / (resolution - 1.0) * (b_max_np - b_min_np)[None, :] + b_min_np[None, :]
+    return vertices, triangles
+
+def extract_geometry_tets(bound_min, bound_max, resolution, threshold, query_func, def_func=None):
+    # print('threshold: {}'.format(threshold))
+    
+    ### using maching cubes ###
+    # u = extract_fields(bound_min, bound_max, resolution, query_func)
+    # vertices, triangles = mcubes.marching_cubes(u, threshold) # grid sdf and marching cubes #
+    # b_max_np = bound_max.detach().cpu().numpy()
+    # b_min_np = bound_min.detach().cpu().numpy()
+
+    # vertices = vertices / (resolution - 1.0) * (b_max_np - b_min_np)[None, :] + b_min_np[None, :]
+    ### using maching cubes ###
+    
+    ## 
+    ### using marching tets ### fiels from tets ##
+    vertices, triangles, tet_sdf_values,  GT_verts, GT_faces  = extract_fields_from_tets(bound_min, bound_max, resolution, query_func, def_func=def_func)
+    # vertices = vertices.detach().cpu().numpy()
+    # triangles = triangles.detach().cpu().numpy()
+    ### using marching tets ###
+    
+    # b_max_np = bound_max.detach().cpu().numpy()
+    # b_min_np = bound_min.detach().cpu().numpy()
+    #
+
+    # vertices = vertices / (resolution - 1.0) * (b_max_np - b_min_np)[None, :] + b_min_np[None, :]
+    return vertices, triangles, tet_sdf_values,  GT_verts, GT_faces 
+
+
+def sample_pdf(bins, weights, n_samples, det=False):
+    # This implementation is from NeRF
+    # Get pdf
+    weights = weights + 1e-5  # prevent nans
+    pdf = weights / torch.sum(weights, -1, keepdim=True)
+    cdf = torch.cumsum(pdf, -1)
+    cdf = torch.cat([torch.zeros_like(cdf[..., :1]), cdf], -1)
+    # Take uniform samples
+    if det:
+        u = torch.linspace(0. + 0.5 / n_samples, 1. - 0.5 / n_samples, steps=n_samples)
+        u = u.expand(list(cdf.shape[:-1]) + [n_samples])
+    else:
+        u = torch.rand(list(cdf.shape[:-1]) + [n_samples])
+
+    # Invert CDF # invert cdf # 
+    u = u.contiguous()
+    inds = torch.searchsorted(cdf, u, right=True)
+    below = torch.max(torch.zeros_like(inds - 1), inds - 1)
+    above = torch.min((cdf.shape[-1] - 1) * torch.ones_like(inds), inds)
+    inds_g = torch.stack([below, above], -1)  # (batch, N_samples, 2)
+
+    matched_shape = [inds_g.shape[0], inds_g.shape[1], cdf.shape[-1]]
+    cdf_g = torch.gather(cdf.unsqueeze(1).expand(matched_shape), 2, inds_g)
+    bins_g = torch.gather(bins.unsqueeze(1).expand(matched_shape), 2, inds_g)
+
+    denom = (cdf_g[..., 1] - cdf_g[..., 0])
+    denom = torch.where(denom < 1e-5, torch.ones_like(denom), denom)
+    t = (u - cdf_g[..., 0]) / denom
+    samples = bins_g[..., 0] + t * (bins_g[..., 1] - bins_g[..., 0])
+
+    return samples
+
+
+def load_GT_vertices(GT_meshes_folder):
+    tot_meshes_fns = os.listdir(GT_meshes_folder)
+    tot_meshes_fns = [fn for fn in tot_meshes_fns if fn.endswith(".obj")]
+    tot_mesh_verts = []
+    tot_mesh_faces = []
+    n_tot_verts = 0
+    for fn in tot_meshes_fns:
+        cur_mesh_fn = os.path.join(GT_meshes_folder, fn)
+        obj_mesh = trimesh.load(cur_mesh_fn, process=False)
+        # obj_mesh.remove_degenerate_faces(height=1e-06)
+
+        verts_obj = np.array(obj_mesh.vertices)
+        faces_obj = np.array(obj_mesh.faces)
+        
+        tot_mesh_verts.append(verts_obj)
+        tot_mesh_faces.append(faces_obj + n_tot_verts)
+        n_tot_verts += verts_obj.shape[0]
+        
+        # tot_mesh_faces.append(faces_obj)
+    tot_mesh_verts = np.concatenate(tot_mesh_verts, axis=0)
+    tot_mesh_faces = np.concatenate(tot_mesh_faces, axis=0)
+    return tot_mesh_verts, tot_mesh_faces
+
+
+class NeuSRenderer:
+    def __init__(self,
+                 nerf,
+                 sdf_network,
+                 deviation_network,
+                 color_network,
+                 n_samples,
+                 n_importance,
+                 n_outside,
+                 up_sample_steps,
+                 perturb):
+        self.nerf = nerf # multiple sdf networks and deviation networks and xxx #
+        self.sdf_network = sdf_network
+        self.deviation_network = deviation_network
+        self.color_network = color_network
+        self.n_samples = n_samples
+        self.n_importance = n_importance
+        self.n_outside = n_outside
+        self.up_sample_steps = up_sample_steps
+        self.perturb = perturb
+        
+        GT_meshes_folder = "/home/xueyi/diffsim/DiffHand/assets/hand"
+        self.mesh_vertices, self.mesh_faces = load_GT_vertices(GT_meshes_folder=GT_meshes_folder)
+        maxx_pts = 25.
+        minn_pts = -15.
+        self.mesh_vertices = (self.mesh_vertices - minn_pts) / (maxx_pts - minn_pts)
+        f = SDF(self.mesh_vertices, self.mesh_faces)
+        self.gt_sdf = f ## a unite sphere or box
+        
+        self.minn_pts = 0
+        self.maxx_pts = 1.
+        
+        # self.minn_pts = -1.5 # gorudn-truth states with the deformation -> update the sdf value fiedl 
+        # self.maxx_pts = 1.5 # 
+        self.bkg_pts = ... # TODO: the bkg pts # bkg_pts; # bkg_pts_defs #
+        self.cur_fr_bkg_pts_defs = ... # TODO: set the cur_bkg_pts_defs for each frame #
+        self.dist_interp_thres = ... # TODO: set the cur_bkg_pts_defs #
+        
+        self.bending_network = ... # TODO: add the bending network # 
+        self.use_bending_network = ... # TODO: set the property #
+        self.use_delta_bending = ... # TODO
+        self.prev_sdf_network = ... # TODO
+        self.use_selector = False
+        # use bending network # 
+        # two bending netwrok 
+        # two sdf networks
+        
+        
+    # get the pts and render the pts #
+    # pts and the rendering pts #
+    def deform_pts(self, pts, pts_ts=0): # deform pts # 
+        
+        if self.use_bending_network:
+            if len(pts.size()) == 3:
+                nnb, nns = pts.size(0), pts.size(1)
+                pts_exp = pts.contiguous().view(nnb * nns, -1).contiguous()
+            else:
+                pts_exp = pts
+            # pts_ts #
+            if self.use_delta_bending:
+                
+                if isinstance(self.bending_network, list):
+                    pts_offsets = []
+                    for i_obj, cur_bending_network in enumerate(self.bending_network):
+                        if isinstance(cur_bending_network, fields.BendingNetwork):
+                            for cur_pts_ts in range(pts_ts, -1, -1):
+                                cur_pts_exp = cur_bending_network(pts_exp if cur_pts_ts == pts_ts else cur_pts_exp, input_pts_ts=cur_pts_ts) 
+                        elif isinstance(cur_bending_network, fields.BendingNetworkRigidTrans):
+                            cur_pts_exp = cur_bending_network(pts_exp, input_pts_ts=cur_pts_ts)
+                        else:
+                            raise ValueError('Encountered with unexpected bending network class...')
+                        pts_offsets.append(cur_pts_exp - pts_exp)
+                    pts_offsets = torch.stack(pts_offsets, dim=0)
+                    pts_offsets = torch.sum(pts_offsets, dim=0)
+                    pts_exp = pts_exp + pts_offsets
+                # for cur_pts_ts in range(pts_ts, -1, -1):
+                #     if isinstance(self.bending_network, list): # pts ts # 
+                #         for i_obj, cur_bending_network in enumerate(self.bending_network):
+                #             pts_exp = cur_bending_network(pts_exp, input_pts_ts=cur_pts_ts) 
+                #     else:
+                #         pts_exp = self.bending_network(pts_exp, input_pts_ts=cur_pts_ts)
+            else:
+                if isinstance(self.bending_network, list): # prev sdf network #
+                    pts_offsets = []
+                    for i_obj, cur_bending_network in enumerate(self.bending_network):
+                        bended_pts_exp = cur_bending_network(pts_exp, input_pts_ts=pts_ts)
+                        pts_offsets.append(bended_pts_exp - pts_exp)
+                    pts_offsets = torch.stack(pts_offsets, dim=0)
+                    pts_offsets = torch.sum(pts_offsets, dim=0)
+                    pts_exp = pts_exp + pts_offsets
+                else:
+                    pts_exp = self.bending_network(pts_exp, input_pts_ts=pts_ts)
+            if len(pts.size()) == 3:
+                pts = pts_exp.contiguous().view(nnb, nns, -1).contiguous()
+            else:
+                pts = pts_exp
+            return pts
+            
+        # pts: nn_batch x nn_samples x 3
+        if len(pts.size()) == 3:
+            nnb, nns = pts.size(0), pts.size(1)
+            pts_exp = pts.contiguous().view(nnb * nns, -1).contiguous()
+        else:
+            pts_exp = pts
+        # print(f"prior to deforming: {pts.size()}")
+        
+        dist_pts_to_bkg_pts = torch.sum(
+            (pts_exp.unsqueeze(1) - self.bkg_pts.unsqueeze(0)) ** 2, dim=-1 ## nn_pts_exp x nn_bkg_pts
+        )
+        dist_mask = dist_pts_to_bkg_pts <= self.dist_interp_thres # 
+        dist_mask_float = dist_mask.float()
+        
+        # dist_mask_float #
+        cur_fr_bkg_def_exp = self.cur_fr_bkg_pts_defs.unsqueeze(0).repeat(pts_exp.size(0), 1, 1).contiguous()
+        cur_fr_pts_def = torch.sum(
+            cur_fr_bkg_def_exp * dist_mask_float.unsqueeze(-1), dim=1
+        )
+        dist_mask_float_summ = torch.sum(
+            dist_mask_float, dim=1
+        )
+        dist_mask_float_summ = torch.clamp(dist_mask_float_summ, min=1)
+        cur_fr_pts_def = cur_fr_pts_def / dist_mask_float_summ.unsqueeze(-1) # bkg pts deformation #
+        pts_exp = pts_exp - cur_fr_pts_def
+        if len(pts.size()) == 3:
+            pts = pts_exp.contiguous().view(nnb, nns, -1).contiguous()
+        else:
+            pts = pts_exp
+        return pts # 
+    
+    
+    def deform_pts_with_selector(self, pts, pts_ts=0): # deform pts # 
+        
+        if self.use_bending_network:
+            if len(pts.size()) == 3:
+                nnb, nns = pts.size(0), pts.size(1)
+                pts_exp = pts.contiguous().view(nnb * nns, -1).contiguous()
+            else:
+                pts_exp = pts
+            # pts_ts #
+            if self.use_delta_bending:
+                if isinstance(self.bending_network, list):
+                    bended_pts = []
+                    queries_sdfs_selector = []
+                    for i_obj, cur_bending_network in enumerate(self.bending_network):
+                        if cur_bending_network.use_opt_rigid_translations:
+                            bended_pts_exp = cur_bending_network(pts_exp, input_pts_ts=pts_ts)
+                        else: 
+                            # bended_pts_exp = pts_exp.clone()
+                            for cur_pts_ts in range(pts_ts, -1, -1):
+                                bended_pts_exp = cur_bending_network(pts_exp if cur_pts_ts == pts_ts else bended_pts_exp, input_pts_ts=cur_pts_ts)
+                        _, cur_bended_pts_selecotr = self.query_pts_sdf_fn_for_selector(bended_pts_exp)
+                        bended_pts.append(bended_pts_exp)
+                        queries_sdfs_selector.append(cur_bended_pts_selecotr)
+                    bended_pts = torch.stack(bended_pts, dim=1) # nn_pts x 2 x 3 for bended pts #
+                    queries_sdfs_selector = torch.stack(queries_sdfs_selector, dim=1) # nn_pts x 2
+                    # queries_sdfs_selector = (queries_sdfs_selector.sum(dim=1) > 0.5).float().long()
+                    sdf_selector = queries_sdfs_selector[:, -1]
+                    # sdf_selector = queries_sdfs_selector
+                    # delta_sdf, sdf_selector = self.query_pts_sdf_fn_for_selector(pts_exp)
+                    bended_pts = batched_index_select(values=bended_pts, indices=sdf_selector.unsqueeze(1), dim=1).squeeze(1) # nn_pts x 3 #
+                    # print(f"bended_pts: {bended_pts.size()}, pts_exp: {pts_exp.size()}")
+                    pts_exp = bended_pts.squeeze(1)
+                            
+                
+                # for cur_pts_ts in range(pts_ts, -1, -1):
+                #     if isinstance(self.bending_network, list):
+                #         for i_obj, cur_bending_network in enumerate(self.bending_network):
+                #             pts_exp = cur_bending_network(pts_exp, input_pts_ts=cur_pts_ts) 
+                #     else:
+                        
+                #         pts_exp = self.bending_network(pts_exp, input_pts_ts=cur_pts_ts)
+            else:
+                if isinstance(self.bending_network, list): # prev sdf network #
+                    # pts_offsets = []
+                    bended_pts = []
+                    queries_sdfs_selector = []
+                    for i_obj, cur_bending_network in enumerate(self.bending_network):
+                        bended_pts_exp = cur_bending_network(pts_exp, input_pts_ts=pts_ts)
+                        # pts_offsets.append(bended_pts_exp - pts_exp)
+                        _, cur_bended_pts_selecotr = self.query_pts_sdf_fn_for_selector(bended_pts_exp)
+                        bended_pts.append(bended_pts_exp)
+                        queries_sdfs_selector.append(cur_bended_pts_selecotr)
+                    bended_pts = torch.stack(bended_pts, dim=1) # nn_pts x 2 x 3 for bended pts #
+                    queries_sdfs_selector = torch.stack(queries_sdfs_selector, dim=1) # nn_pts x 2
+                    # queries_sdfs_selector = (queries_sdfs_selector.sum(dim=1) > 0.5).float().long()
+                    sdf_selector = queries_sdfs_selector[:, -1]
+                    # sdf_selector = queries_sdfs_selector
+                    
+                    
+                    # delta_sdf, sdf_selector = self.query_pts_sdf_fn_for_selector(pts_exp)
+                    bended_pts = batched_index_select(values=bended_pts, indices=sdf_selector.unsqueeze(1), dim=1).squeeze(1) # nn_pts x 3 #
+                    # print(f"bended_pts: {bended_pts.size()}, pts_exp: {pts_exp.size()}")
+                    pts_exp = bended_pts.squeeze(1)
+                    
+                    # pts_offsets = torch.stack(pts_offsets, dim=0)
+                    # pts_offsets = torch.sum(pts_offsets, dim=0)
+                    # pts_exp = pts_exp + pts_offsets
+                else:
+                    pts_exp = self.bending_network(pts_exp, input_pts_ts=pts_ts)
+            if len(pts.size()) == 3:
+                pts = pts_exp.contiguous().view(nnb, nns, -1).contiguous()
+            else:
+                pts = pts_exp
+            return pts
+            
+        # pts: nn_batch x nn_samples x 3
+        if len(pts.size()) == 3:
+            nnb, nns = pts.size(0), pts.size(1)
+            pts_exp = pts.contiguous().view(nnb * nns, -1).contiguous()
+        else:
+            pts_exp = pts
+        # print(f"prior to deforming: {pts.size()}")
+        
+        dist_pts_to_bkg_pts = torch.sum(
+            (pts_exp.unsqueeze(1) - self.bkg_pts.unsqueeze(0)) ** 2, dim=-1 ## nn_pts_exp x nn_bkg_pts
+        )
+        dist_mask = dist_pts_to_bkg_pts <= self.dist_interp_thres # 
+        dist_mask_float = dist_mask.float()
+        
+        # dist_mask_float #
+        cur_fr_bkg_def_exp = self.cur_fr_bkg_pts_defs.unsqueeze(0).repeat(pts_exp.size(0), 1, 1).contiguous()
+        cur_fr_pts_def = torch.sum(
+            cur_fr_bkg_def_exp * dist_mask_float.unsqueeze(-1), dim=1
+        )
+        dist_mask_float_summ = torch.sum(
+            dist_mask_float, dim=1
+        )
+        dist_mask_float_summ = torch.clamp(dist_mask_float_summ, min=1)
+        cur_fr_pts_def = cur_fr_pts_def / dist_mask_float_summ.unsqueeze(-1) # bkg pts deformation #
+        pts_exp = pts_exp - cur_fr_pts_def
+        if len(pts.size()) == 3:
+            pts = pts_exp.contiguous().view(nnb, nns, -1).contiguous()
+        else:
+            pts = pts_exp
+        return pts # 
+    
+    
+    def deform_pts_passive(self, pts, pts_ts=0):
+        
+        if self.use_bending_network:
+            if len(pts.size()) == 3:
+                nnb, nns = pts.size(0), pts.size(1)
+                pts_exp = pts.contiguous().view(nnb * nns, -1).contiguous()
+            else:
+                pts_exp = pts
+            # pts_ts #
+            if self.use_delta_bending:
+                for cur_pts_ts in range(pts_ts, -1, -1):
+                    if isinstance(self.bending_network, list):
+                        for i_obj, cur_bending_network in enumerate(self.bending_network):
+                            pts_exp = cur_bending_network(pts_exp, input_pts_ts=cur_pts_ts) 
+                    else:
+                        pts_exp = self.bending_network(pts_exp, input_pts_ts=cur_pts_ts)
+            else:
+                # if isinstance(self.bending_network, list):
+                #     pts_offsets = []
+                #     for i_obj, cur_bending_network in enumerate(self.bending_network):
+                #         bended_pts_exp = cur_bending_network(pts_exp, input_pts_ts=pts_ts)
+                #         pts_offsets.append(bended_pts_exp - pts_exp)
+                #     pts_offsets = torch.stack(pts_offsets, dim=0)
+                #     pts_offsets = torch.sum(pts_offsets, dim=0)
+                #     pts_exp = pts_exp + pts_offsets
+                # else:
+                pts_exp = self.bending_network[-1](pts_exp, input_pts_ts=pts_ts)
+            if len(pts.size()) == 3:
+                pts = pts_exp.contiguous().view(nnb, nns, -1).contiguous()
+            else:
+                pts = pts_exp
+            return pts
+            
+        # pts: nn_batch x nn_samples x 3
+        if len(pts.size()) == 3:
+            nnb, nns = pts.size(0), pts.size(1)
+            pts_exp = pts.contiguous().view(nnb * nns, -1).contiguous()
+        else:
+            pts_exp = pts
+        # print(f"prior to deforming: {pts.size()}")
+        
+        dist_pts_to_bkg_pts = torch.sum(
+            (pts_exp.unsqueeze(1) - self.bkg_pts.unsqueeze(0)) ** 2, dim=-1 ## nn_pts_exp x nn_bkg_pts
+        )
+        dist_mask = dist_pts_to_bkg_pts <= self.dist_interp_thres # 
+        dist_mask_float = dist_mask.float()
+        
+        # dist_mask_float #
+        cur_fr_bkg_def_exp = self.cur_fr_bkg_pts_defs.unsqueeze(0).repeat(pts_exp.size(0), 1, 1).contiguous()
+        cur_fr_pts_def = torch.sum(
+            cur_fr_bkg_def_exp * dist_mask_float.unsqueeze(-1), dim=1
+        )
+        dist_mask_float_summ = torch.sum(
+            dist_mask_float, dim=1
+        )
+        dist_mask_float_summ = torch.clamp(dist_mask_float_summ, min=1)
+        cur_fr_pts_def = cur_fr_pts_def / dist_mask_float_summ.unsqueeze(-1) # bkg pts deformation #
+        pts_exp = pts_exp - cur_fr_pts_def
+        if len(pts.size()) == 3:
+            pts = pts_exp.contiguous().view(nnb, nns, -1).contiguous()
+        else:
+            pts = pts_exp
+        return pts # 
+    
+    
+    def query_pts_sdf_fn_for_selector(self, pts):
+        # for negative 
+        # 1) inside the current mesh but outside the previous mesh ---> negative sdf for this field but positive for another field
+        # 2) negative in thie field and also negative in the previous field ---> 
+        # 2) for positive values of this current field ---> 
+        cur_sdf = self.sdf_network.sdf(pts)
+        prev_sdf = self.prev_sdf_network.sdf(pts)
+        neg_neg = ((cur_sdf < 0.).float() + (prev_sdf < 0.).float()) > 1.5
+        neg_pos = ((cur_sdf < 0.).float() + (prev_sdf >= 0.).float()) > 1.5
+        
+        neg_weq_pos = ((cur_sdf <= 0.).float() + (prev_sdf > 0.).float()) > 1.5
+        
+        pos_neg = ((cur_sdf >= 0.).float() + (prev_sdf < 0.).float()) > 1.5
+        pos_pos = ((cur_sdf >= 0.).float() + (prev_sdf >= 0.).float()) > 1.5
+        res_sdf = torch.zeros_like(cur_sdf)
+        res_sdf[neg_neg] = 1. # 
+        res_sdf[neg_pos] = cur_sdf[neg_pos]
+        res_sdf[pos_neg] = cur_sdf[pos_neg]
+        
+        # inside the residual mesh -> must be neg and pos
+        res_sdf_selector = torch.zeros_like(cur_sdf).long() # 
+        # res_sdf_selector[neg_pos] = 1 # is the residual mesh
+        res_sdf_selector[neg_weq_pos] = 1 
+        # res_sdf_selector[]
+        
+        cat_cur_prev_sdf = torch.stack(
+            [cur_sdf, prev_sdf], dim=-1
+        )
+        minn_cur_prev_sdf, _ = torch.min(cat_cur_prev_sdf, dim=-1)
+        res_sdf[pos_pos] = minn_cur_prev_sdf[pos_pos]
+        
+        return res_sdf, res_sdf_selector
+    
+    def query_func_sdf(self, pts):
+        if isinstance(self.sdf_network, list):
+            tot_sdf_values = []
+            for i_obj, cur_sdf_network in enumerate(self.sdf_network):
+                cur_sdf_values = cur_sdf_network.sdf(pts)
+                tot_sdf_values.append(cur_sdf_values)
+            tot_sdf_values = torch.stack(tot_sdf_values, dim=-1)
+            tot_sdf_values, _ = torch.min(tot_sdf_values, dim=-1) # totsdf values #
+            sdf = tot_sdf_values
+        else:
+            sdf = self.sdf_network.sdf(pts)
+        return sdf
+    
+    def query_func_sdf_passive(self, pts):
+        # if isinstance(self.sdf_network, list):
+        #     tot_sdf_values = []
+        #     for i_obj, cur_sdf_network in enumerate(self.sdf_network):
+        #         cur_sdf_values = cur_sdf_network.sdf(pts)
+        #         tot_sdf_values.append(cur_sdf_values)
+        #     tot_sdf_values = torch.stack(tot_sdf_values, dim=-1)
+        #     tot_sdf_values, _ = torch.min(tot_sdf_values, dim=-1) # totsdf values #
+        #     sdf = tot_sdf_values
+        # else:
+        sdf = self.sdf_network[-1].sdf(pts)
+            
+        return sdf
+        
+
+    def render_core_outside(self, rays_o, rays_d, z_vals, sample_dist, nerf, background_rgb=None, pts_ts=0):
+        """
+        Render background
+        """
+        batch_size, n_samples = z_vals.shape
+
+        # Section length
+        dists = z_vals[..., 1:] - z_vals[..., :-1]
+        dists = torch.cat([dists, torch.Tensor([sample_dist]).expand(dists[..., :1].shape)], -1)
+        mid_z_vals = z_vals + dists * 0.5
+
+        # Section midpoints #
+        pts = rays_o[:, None, :] + rays_d[:, None, :] * mid_z_vals[..., :, None]  # batch_size, n_samples, 3 #
+        
+        # pts = pts.flip((-1,)) * 2 - 1
+        pts = pts * 2 - 1
+        
+        if self.use_selector:
+            pts = self.deform_pts_with_selector(pts=pts, pts_ts=pts_ts)
+        else:
+            pts = self.deform_pts(pts=pts, pts_ts=pts_ts)
+
+        dis_to_center = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=True).clip(1.0, 1e10)
+        pts = torch.cat([pts / dis_to_center, 1.0 / dis_to_center], dim=-1)       # batch_size, n_samples, 4 #
+
+        dirs = rays_d[:, None, :].expand(batch_size, n_samples, 3)
+
+        pts = pts.reshape(-1, 3 + int(self.n_outside > 0))
+        dirs = dirs.reshape(-1, 3)
+
+        density, sampled_color = nerf(pts, dirs)
+        sampled_color = torch.sigmoid(sampled_color)
+        alpha = 1.0 - torch.exp(-F.softplus(density.reshape(batch_size, n_samples)) * dists)
+        alpha = alpha.reshape(batch_size, n_samples)
+        weights = alpha * torch.cumprod(torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[:, :-1]
+        sampled_color = sampled_color.reshape(batch_size, n_samples, 3)
+        color = (weights[:, :, None] * sampled_color).sum(dim=1)
+        if background_rgb is not None:
+            color = color + background_rgb * (1.0 - weights.sum(dim=-1, keepdim=True))
+
+        return {
+            'color': color,
+            'sampled_color': sampled_color,
+            'alpha': alpha,
+            'weights': weights,
+        }
+
+    def up_sample(self, rays_o, rays_d, z_vals, sdf, n_importance, inv_s, pts_ts=0):
+        """
+        Up sampling give a fixed inv_s
+        """
+        batch_size, n_samples = z_vals.shape
+        pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]  # n_rays, n_samples, 3
+        
+        # pts = pts.flip((-1,)) * 2 - 1
+        pts = pts * 2 - 1
+        
+        if self.use_selector:
+            pts = self.deform_pts_with_selector(pts=pts, pts_ts=pts_ts)
+        else:
+            pts = self.deform_pts(pts=pts, pts_ts=pts_ts)
+        
+        radius = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=False)
+        inside_sphere = (radius[:, :-1] < 1.0) | (radius[:, 1:] < 1.0)
+        sdf = sdf.reshape(batch_size, n_samples)
+        prev_sdf, next_sdf = sdf[:, :-1], sdf[:, 1:]
+        prev_z_vals, next_z_vals = z_vals[:, :-1], z_vals[:, 1:]
+        mid_sdf = (prev_sdf + next_sdf) * 0.5
+        cos_val = (next_sdf - prev_sdf) / (next_z_vals - prev_z_vals + 1e-5)
+
+        # ----------------------------------------------------------------------------------------------------------
+        # Use min value of [ cos, prev_cos ]
+        # Though it makes the sampling (not rendering) a little bit biased, this strategy can make the sampling more
+        # robust when meeting situations like below:
+        #
+        # SDF
+        # ^
+        # |\          -----x----...
+        # | \        /
+        # |  x      x
+        # |---\----/-------------> 0 level
+        # |    \  /
+        # |     \/
+        # |
+        # ----------------------------------------------------------------------------------------------------------
+        prev_cos_val = torch.cat([torch.zeros([batch_size, 1]), cos_val[:, :-1]], dim=-1)
+        cos_val = torch.stack([prev_cos_val, cos_val], dim=-1)
+        cos_val, _ = torch.min(cos_val, dim=-1, keepdim=False)
+        cos_val = cos_val.clip(-1e3, 0.0) * inside_sphere
+
+        dist = (next_z_vals - prev_z_vals)
+        prev_esti_sdf = mid_sdf - cos_val * dist * 0.5
+        next_esti_sdf = mid_sdf + cos_val * dist * 0.5
+        prev_cdf = torch.sigmoid(prev_esti_sdf * inv_s)
+        next_cdf = torch.sigmoid(next_esti_sdf * inv_s)
+        alpha = (prev_cdf - next_cdf + 1e-5) / (prev_cdf + 1e-5)
+        weights = alpha * torch.cumprod(
+            torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[:, :-1]
+
+        z_samples = sample_pdf(z_vals, weights, n_importance, det=True).detach()
+        return z_samples
+
+    def cat_z_vals(self, rays_o, rays_d, z_vals, new_z_vals, sdf, last=False, pts_ts=0):
+        batch_size, n_samples = z_vals.shape
+        _, n_importance = new_z_vals.shape
+        pts = rays_o[:, None, :] + rays_d[:, None, :] * new_z_vals[..., :, None]
+        
+        # pts = pts.flip((-1,)) * 2 - 1
+        pts = pts * 2 - 1
+        
+        if self.use_selector:
+            pts = self.deform_pts_with_selector(pts=pts, pts_ts=pts_ts)
+        else:
+            pts = self.deform_pts(pts=pts, pts_ts=pts_ts)
+        
+        z_vals = torch.cat([z_vals, new_z_vals], dim=-1)
+        z_vals, index = torch.sort(z_vals, dim=-1)
+
+        if not last:
+            if isinstance(self.sdf_network, list):
+                tot_new_sdf = []
+                for i_obj, cur_sdf_network in enumerate(self.sdf_network):
+                    cur_new_sdf = cur_sdf_network.sdf(pts.reshape(-1, 3)).reshape(batch_size, n_importance)
+                    tot_new_sdf.append(cur_new_sdf)
+                tot_new_sdf = torch.stack(tot_new_sdf, dim=-1)
+                new_sdf, _ = torch.min(tot_new_sdf, dim=-1) # 
+            else:
+                new_sdf = self.sdf_network.sdf(pts.reshape(-1, 3)).reshape(batch_size, n_importance)
+            sdf = torch.cat([sdf, new_sdf], dim=-1)
+            xx = torch.arange(batch_size)[:, None].expand(batch_size, n_samples + n_importance).reshape(-1)
+            index = index.reshape(-1)
+            sdf = sdf[(xx, index)].reshape(batch_size, n_samples + n_importance)
+
+        return z_vals, sdf
+    
+    
+
+    def render_core(self,
+                    rays_o,
+                    rays_d,
+                    z_vals,
+                    sample_dist,
+                    sdf_network,
+                    deviation_network,
+                    color_network,
+                    background_alpha=None,
+                    background_sampled_color=None,
+                    background_rgb=None,
+                    cos_anneal_ratio=0.0,
+                    pts_ts=0):
+        batch_size, n_samples = z_vals.shape
+
+        # Section length
+        dists = z_vals[..., 1:] - z_vals[..., :-1]
+        dists = torch.cat([dists, torch.Tensor([sample_dist]).expand(dists[..., :1].shape)], -1)
+        mid_z_vals = z_vals + dists * 0.5 # z_vals and dists * 0.5 #
+
+        # Section midpoints
+        pts = rays_o[:, None, :] + rays_d[:, None, :] * mid_z_vals[..., :, None]  # n_rays, n_samples, 3
+        dirs = rays_d[:, None, :].expand(pts.shape)
+
+        pts = pts.reshape(-1, 3) # pts, nn_ou
+        dirs = dirs.reshape(-1, 3)
+        
+        pts = (pts - self.minn_pts) / (self.maxx_pts - self.minn_pts)
+        
+        # pts = pts.flip((-1,)) * 2 - 1
+        pts = pts * 2 - 1
+        
+        if self.use_selector:
+            pts = self.deform_pts_with_selector(pts=pts, pts_ts=pts_ts)
+        else:
+            pts = self.deform_pts(pts=pts, pts_ts=pts_ts)
+        
+        if isinstance(sdf_network, list):
+            tot_sdf = []
+            tot_feature_vector = []
+            tot_obj_sel = []
+            tot_gradients = []
+            for i_obj, cur_sdf_network in enumerate(sdf_network):
+                cur_sdf_nn_output = cur_sdf_network(pts)
+                cur_sdf, cur_feature_vector = cur_sdf_nn_output[:, :1], cur_sdf_nn_output[:, 1:]
+                tot_sdf.append(cur_sdf)
+                tot_feature_vector.append(cur_feature_vector)
+                
+                gradients = cur_sdf_network.gradient(pts).squeeze()
+                tot_gradients.append(gradients)
+            tot_sdf = torch.stack(tot_sdf, dim=-1)
+            sdf, obj_sel = torch.min(tot_sdf, dim=-1)
+            feature_vector = torch.stack(tot_feature_vector, dim=1)
+            
+            # batched_index_select
+            # print(f"before sel: {feature_vector.size()}, obj_sel: {obj_sel.size()}")
+            feature_vector = batched_index_select(values=feature_vector, indices=obj_sel, dim=1).squeeze(1)
+            
+            
+            # feature_vector = feature_vector[obj_sel.unsqueeze(-1), :].squeeze(1)
+            # print(f"after sel: {feature_vector.size()}")
+            tot_gradients = torch.stack(tot_gradients, dim=1)
+            # gradients = tot_gradients[obj_sel.unsqueeze(-1)].squeeze(1)
+            gradients = batched_index_select(values=tot_gradients, indices=obj_sel, dim=1).squeeze(1)
+            # print(f"gradients: {gradients.size()}, tot_gradients: {tot_gradients.size()}")
+            
+        else:
+            sdf_nn_output = sdf_network(pts)
+            sdf = sdf_nn_output[:, :1]
+            feature_vector = sdf_nn_output[:, 1:]
+            gradients = sdf_network.gradient(pts).squeeze()
+            
+        sampled_color = color_network(pts, gradients, dirs, feature_vector).reshape(batch_size, n_samples, 3)
+
+        # deviation network #
+        inv_s = deviation_network(torch.zeros([1, 3]))[:, :1].clip(1e-6, 1e6)           # Single parameter
+        inv_s = inv_s.expand(batch_size * n_samples, 1)
+
+        true_cos = (dirs * gradients).sum(-1, keepdim=True)
+
+        # "cos_anneal_ratio" grows from 0 to 1 in the beginning training iterations. The anneal strategy below makes
+        # the cos value "not dead" at the beginning training iterations, for better convergence.
+        iter_cos = -(F.relu(-true_cos * 0.5 + 0.5) * (1.0 - cos_anneal_ratio) +
+                     F.relu(-true_cos) * cos_anneal_ratio)  # always non-positive
+
+        # Estimate signed distances at section points
+        estimated_next_sdf = sdf + iter_cos * dists.reshape(-1, 1) * 0.5
+        estimated_prev_sdf = sdf - iter_cos * dists.reshape(-1, 1) * 0.5
+
+        prev_cdf = torch.sigmoid(estimated_prev_sdf * inv_s)
+        next_cdf = torch.sigmoid(estimated_next_sdf * inv_s)
+
+        p = prev_cdf - next_cdf
+        c = prev_cdf
+
+        alpha = ((p + 1e-5) / (c + 1e-5)).reshape(batch_size, n_samples).clip(0.0, 1.0)
+
+        pts_norm = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=True).reshape(batch_size, n_samples)
+        inside_sphere = (pts_norm < 1.0).float().detach()
+        relax_inside_sphere = (pts_norm < 1.2).float().detach()
+
+        # Render with background
+        if background_alpha is not None:
+            alpha = alpha * inside_sphere + background_alpha[:, :n_samples] * (1.0 - inside_sphere)
+            alpha = torch.cat([alpha, background_alpha[:, n_samples:]], dim=-1)
+            sampled_color = sampled_color * inside_sphere[:, :, None] +\
+                            background_sampled_color[:, :n_samples] * (1.0 - inside_sphere)[:, :, None]
+            sampled_color = torch.cat([sampled_color, background_sampled_color[:, n_samples:]], dim=1)
+
+        weights = alpha * torch.cumprod(torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[:, :-1]
+        weights_sum = weights.sum(dim=-1, keepdim=True)
+
+        color = (sampled_color * weights[:, :, None]).sum(dim=1)
+        if background_rgb is not None:    # Fixed background, usually black
+            color = color + background_rgb * (1.0 - weights_sum)
+
+        # Eikonal loss
+        gradient_error = (torch.linalg.norm(gradients.reshape(batch_size, n_samples, 3), ord=2,
+                                            dim=-1) - 1.0) ** 2
+        gradient_error = (relax_inside_sphere * gradient_error).sum() / (relax_inside_sphere.sum() + 1e-5)
+
+        return {
+            'color': color,
+            'sdf': sdf,
+            'dists': dists,
+            'gradients': gradients.reshape(batch_size, n_samples, 3),
+            's_val': 1.0 / inv_s,
+            'mid_z_vals': mid_z_vals,
+            'weights': weights,
+            'cdf': c.reshape(batch_size, n_samples),
+            'gradient_error': gradient_error,
+            'inside_sphere': inside_sphere
+        }
+
+    def render(self, rays_o, rays_d, near, far, pts_ts=0, perturb_overwrite=-1, background_rgb=None, cos_anneal_ratio=0.0, use_gt_sdf=False):
+        batch_size = len(rays_o)
+        sample_dist = 2.0 / self.n_samples  # in a unit sphere # # Assuming the region of interest is a unit sphere
+        z_vals = torch.linspace(0.0, 1.0, self.n_samples) # linspace #
+        z_vals = near + (far - near) * z_vals[None, :]
+
+        z_vals_outside = None
+        if self.n_outside > 0:
+            z_vals_outside = torch.linspace(1e-3, 1.0 - 1.0 / (self.n_outside + 1.0), self.n_outside)
+
+        n_samples = self.n_samples
+        perturb = self.perturb
+
+        if perturb_overwrite >= 0:
+            perturb = perturb_overwrite
+        if perturb > 0:
+            t_rand = (torch.rand([batch_size, 1]) - 0.5)
+            z_vals = z_vals + t_rand * 2.0 / self.n_samples
+
+            if self.n_outside > 0: # z values output # n_outside #
+                mids = .5 * (z_vals_outside[..., 1:] + z_vals_outside[..., :-1])
+                upper = torch.cat([mids, z_vals_outside[..., -1:]], -1)
+                lower = torch.cat([z_vals_outside[..., :1], mids], -1)
+                t_rand = torch.rand([batch_size, z_vals_outside.shape[-1]])
+                z_vals_outside = lower[None, :] + (upper - lower)[None, :] * t_rand
+
+        if self.n_outside > 0:
+            z_vals_outside = far / torch.flip(z_vals_outside, dims=[-1]) + 1.0 / self.n_samples
+
+        background_alpha = None
+        background_sampled_color = None
+
+        # Up sample
+        if self.n_importance > 0:
+            with torch.no_grad():
+                pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]
+                
+                pts = (pts - self.minn_pts) / (self.maxx_pts - self.minn_pts)
+                # sdf = self.sdf_network.sdf(pts.reshape(-1, 3)).reshape(batch_size, self.n_samples)
+                # gt_sdf #
+                
+                # 
+                # pts = ((pts - xyz_min) / (xyz_max - xyz_min)).flip((-1,)) * 2 - 1
+                
+                # pts = pts.flip((-1,)) * 2 - 1
+                pts = pts * 2 - 1                             
+                
+                if self.use_selector:
+                    pts = self.deform_pts_with_selector(pts=pts, pts_ts=pts_ts)
+                else:
+                    pts = self.deform_pts(pts=pts, pts_ts=pts_ts) # give nthe pts
+                
+                pts_exp = pts.reshape(-1, 3)
+                # minn_pts, _ = torch.min(pts_exp, dim=0)
+                # maxx_pts, _ = torch.max(pts_exp, dim=0) # deformation field (not a rigid one) -> the meshes #
+                # print(f"minn_pts: {minn_pts}, maxx_pts: {maxx_pts}")
+
+                # pts_to_near = pts - near.unsqueeze(1)
+                # maxx_pts = 1.5; minn_pts = -1.5
+                # # maxx_pts = 3; minn_pts = -3
+                # # maxx_pts = 1; minn_pts = -1
+                # pts_exp = (pts_exp - minn_pts) / (maxx_pts - minn_pts)
+                
+                ## render and iamges  ####  
+                if use_gt_sdf:
+                    ### use the GT sdf field ####
+                    # print(f"Using gt sdf :")
+                    sdf = self.gt_sdf(pts_exp.reshape(-1, 3).detach().cpu().numpy())
+                    sdf = torch.from_numpy(sdf).float().cuda()
+                    sdf = sdf.reshape(batch_size, self.n_samples)
+                    ### use the GT sdf field ####
+                else:
+                    # pts_exp: (bsz x nn_s) x 3 -> (sdf_network) -> (bsz x nn_s)
+                    #### use the optimized sdf field ####
+                    
+                    # sdf = self.sdf_network.sdf(pts_exp).reshape(batch_size, self.n_samples)
+                    
+                    if isinstance(self.sdf_network, list):
+                        tot_sdf_values = []
+                        for i_obj, cur_sdf_network in enumerate(self.sdf_network):
+                            cur_sdf_values = cur_sdf_network.sdf(pts_exp).reshape(batch_size, self.n_samples)
+                            tot_sdf_values.append(cur_sdf_values)
+                        tot_sdf_values = torch.stack(tot_sdf_values, dim=-1)
+                        tot_sdf_values, _ = torch.min(tot_sdf_values, dim=-1) # totsdf values #
+                        sdf = tot_sdf_values
+                    else:
+                        sdf = self.sdf_network.sdf(pts_exp).reshape(batch_size, self.n_samples)
+                    
+                    #### use the optimized sdf field ####
+
+                for i in range(self.up_sample_steps):
+                    new_z_vals = self.up_sample(rays_o,
+                                                rays_d,
+                                                z_vals,
+                                                sdf,
+                                                self.n_importance // self.up_sample_steps,
+                                                64 * 2**i,
+                                                pts_ts=pts_ts)
+                    z_vals, sdf = self.cat_z_vals(rays_o,
+                                                  rays_d,
+                                                  z_vals,
+                                                  new_z_vals,
+                                                  sdf,
+                                                  last=(i + 1 == self.up_sample_steps),
+                                                  pts_ts=pts_ts)
+
+            n_samples = self.n_samples + self.n_importance
+
+        # Background model
+        if self.n_outside > 0:
+            z_vals_feed = torch.cat([z_vals, z_vals_outside], dim=-1)
+            z_vals_feed, _ = torch.sort(z_vals_feed, dim=-1)
+            ret_outside = self.render_core_outside(rays_o, rays_d, z_vals_feed, sample_dist, self.nerf, pts_ts=pts_ts)
+
+            background_sampled_color = ret_outside['sampled_color']
+            background_alpha = ret_outside['alpha']
+
+        # Render core
+        ret_fine = self.render_core(rays_o, # 
+                                    rays_d,
+                                    z_vals,
+                                    sample_dist,
+                                    self.sdf_network,
+                                    self.deviation_network,
+                                    self.color_network,
+                                    background_rgb=background_rgb,
+                                    background_alpha=background_alpha,
+                                    background_sampled_color=background_sampled_color,
+                                    cos_anneal_ratio=cos_anneal_ratio,
+                                    pts_ts=pts_ts)
+
+        color_fine = ret_fine['color']
+        weights = ret_fine['weights']
+        weights_sum = weights.sum(dim=-1, keepdim=True)
+        gradients = ret_fine['gradients']
+        s_val = ret_fine['s_val'].reshape(batch_size, n_samples).mean(dim=-1, keepdim=True)
+
+        return {
+            'color_fine': color_fine,
+            's_val': s_val,
+            'cdf_fine': ret_fine['cdf'],
+            'weight_sum': weights_sum,
+            'weight_max': torch.max(weights, dim=-1, keepdim=True)[0],
+            'gradients': gradients,
+            'weights': weights,
+            'gradient_error': ret_fine['gradient_error'],
+            'inside_sphere': ret_fine['inside_sphere']
+        }
+
+    def extract_geometry(self, bound_min, bound_max, resolution, threshold=0.0):
+        return extract_geometry(bound_min, # extract geometry #
+                                bound_max,
+                                resolution=resolution,
+                                threshold=threshold,
+                                # query_func=lambda pts: -self.sdf_network.sdf(pts),
+                                query_func=lambda pts: -self.query_func_sdf(pts)
+                                )
+        
+    # if self.deform_pts_with_selector:
+    #         pts = self.deform_pts_with_selector(pts=pts, pts_ts=pts_ts)
+    def extract_geometry_tets(self, bound_min, bound_max, resolution, pts_ts=0, threshold=0.0, wdef=False):
+        if wdef:
+            return extract_geometry_tets(bound_min, # extract geometry #
+                                bound_max,
+                                resolution=resolution,
+                                threshold=threshold,
+                                query_func=lambda pts: -self.query_func_sdf(pts), # lambda pts: -self.sdf_network.sdf(pts),
+                                def_func=lambda pts: self.deform_pts(pts, pts_ts=pts_ts) if not self.use_selector else self.deform_pts_with_selector(pts=pts, pts_ts=pts_ts))
+        else:
+            return extract_geometry_tets(bound_min, # extract geometry #
+                                    bound_max,
+                                    resolution=resolution,
+                                    threshold=threshold,
+                                    # query_func=lambda pts: -self.sdf_network.sdf(pts)
+                                     query_func=lambda pts: -self.query_func_sdf(pts), # lambda pts: -self.sdf_network.sdf(pts),
+                                    )
+            
+    def extract_geometry_tets_passive(self, bound_min, bound_max, resolution, pts_ts=0, threshold=0.0, wdef=False):
+        if wdef:
+            return extract_geometry_tets(bound_min, # extract geometry #
+                                bound_max,
+                                resolution=resolution,
+                                threshold=threshold,
+                                query_func=lambda pts: -self.query_func_sdf_passive(pts), # lambda pts: -self.sdf_network.sdf(pts),
+                                def_func=lambda pts: self.deform_pts_passive(pts, pts_ts=pts_ts))
+        else:
+            return extract_geometry_tets(bound_min, # extract geometry #
+                                    bound_max,
+                                    resolution=resolution,
+                                    threshold=threshold,
+                                    # query_func=lambda pts: -self.sdf_network.sdf(pts)
+                                     query_func=lambda pts: -self.query_func_sdf(pts), # lambda pts: -self.sdf_network.sdf(pts),
+                                    )

models/renderer_def_multi_objs_compositional.py

@@ -0,0 +1,1510 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import logging
+import mcubes
+from icecream import ic
+import os
+
+import trimesh
+from pysdf import SDF
+
+import models.fields as fields
+
+from uni_rep.rep_3d.dmtet import marching_tets_tetmesh, create_tetmesh_variables
+
+def batched_index_select(values, indices, dim = 1):
+  value_dims = values.shape[(dim + 1):]
+  values_shape, indices_shape = map(lambda t: list(t.shape), (values, indices))
+  indices = indices[(..., *((None,) * len(value_dims)))]
+  indices = indices.expand(*((-1,) * len(indices_shape)), *value_dims)
+  value_expand_len = len(indices_shape) - (dim + 1)
+  values = values[(*((slice(None),) * dim), *((None,) * value_expand_len), ...)]
+
+  value_expand_shape = [-1] * len(values.shape)
+  expand_slice = slice(dim, (dim + value_expand_len))
+  value_expand_shape[expand_slice] = indices.shape[expand_slice]
+  values = values.expand(*value_expand_shape)
+
+  dim += value_expand_len
+  return values.gather(dim, indices)
+
+
+def create_mt_variable(device):
+    triangle_table = torch.tensor(
+        [
+            [-1, -1, -1, -1, -1, -1],
+            [1, 0, 2, -1, -1, -1],
+            [4, 0, 3, -1, -1, -1],
+            [1, 4, 2, 1, 3, 4],
+            [3, 1, 5, -1, -1, -1],
+            [2, 3, 0, 2, 5, 3],
+            [1, 4, 0, 1, 5, 4],
+            [4, 2, 5, -1, -1, -1],
+            [4, 5, 2, -1, -1, -1],
+            [4, 1, 0, 4, 5, 1],
+            [3, 2, 0, 3, 5, 2],
+            [1, 3, 5, -1, -1, -1],
+            [4, 1, 2, 4, 3, 1],
+            [3, 0, 4, -1, -1, -1],
+            [2, 0, 1, -1, -1, -1],
+            [-1, -1, -1, -1, -1, -1]
+        ], dtype=torch.long, device=device)
+
+    num_triangles_table = torch.tensor([0, 1, 1, 2, 1, 2, 2, 1, 1, 2, 2, 1, 2, 1, 1, 0], dtype=torch.long, device=device)
+    base_tet_edges = torch.tensor([0, 1, 0, 2, 0, 3, 1, 2, 1, 3, 2, 3], dtype=torch.long, device=device)
+    v_id = torch.pow(2, torch.arange(4, dtype=torch.long, device=device))
+    return triangle_table, num_triangles_table, base_tet_edges, v_id
+
+
+
+def  extract_fields_from_tets(bound_min, bound_max, resolution, query_func, def_func=None):
+    # load tet via resolution #
+    # scale them via bounds #
+    # extract the geometry #
+    # /home/xueyi/gen/DeepMetaHandles/data/tets/100_compress.npz # strange #
+    device = bound_min.device
+    # if resolution in [64, 70, 80, 90, 100]:
+    #     tet_fn = f"/home/xueyi/gen/DeepMetaHandles/data/tets/{resolution}_compress.npz"
+    # else:
+    # tet_fn = f"/home/xueyi/gen/DeepMetaHandles/data/tets/{100}_compress.npz"
+    tet_fn = f"/home/xueyi/gen/DeepMetaHandles/data/tets/{100}_compress.npz"
+    if not os.path.exists(tet_fn):
+        tet_fn = f"/data/xueyi/NeuS/data/tets/{100}_compress.npz"
+    tets = np.load(tet_fn)
+    verts = torch.from_numpy(tets['vertices']).float().to(device) # verts positions 
+    indices = torch.from_numpy(tets['tets']).long().to(device) # .to(self.device)
+    # split #
+    # verts; verts; #
+    minn_verts, _ = torch.min(verts, dim=0)
+    maxx_verts, _ = torch.max(verts, dim=0) # (3, ) # exporting the  
+    # scale_verts = maxx_verts - minn_verts
+    scale_bounds = bound_max - bound_min # scale bounds #
+    
+    ### scale the vertices ###
+    scaled_verts = (verts - minn_verts.unsqueeze(0)) / (maxx_verts - minn_verts).unsqueeze(0) ### the maxx and minn verts scales ###
+    
+    # scaled_verts = (verts - minn_verts.unsqueeze(0)) / (maxx_verts - minn_verts).unsqueeze(0) ### the maxx and minn verts scales ###
+
+    scaled_verts = scaled_verts * 2. - 1. # init the sdf filed viathe tet mesh vertices and the sdf values ##
+    # scaled_verts = (scaled_verts * scale_bounds.unsqueeze(0)) + bound_min.unsqueeze(0) ## the scaled verts ###
+    
+    # scaled_verts = scaled_verts - scale_bounds.unsqueeze(0) / 2. # 
+    # scaled_verts = scaled_verts -  bound_min.unsqueeze(0) - scale_bounds.unsqueeze(0) / 2.
+    
+    sdf_values = []
+    N = 64
+    query_bundles = N ** 3 ### N^3
+    query_NNs = scaled_verts.size(0) // query_bundles
+    if query_NNs * query_bundles < scaled_verts.size(0):
+        query_NNs += 1
+    for i_query in range(query_NNs):
+        cur_bundle_st = i_query * query_bundles
+        cur_bundle_ed = (i_query + 1) * query_bundles
+        cur_bundle_ed = min(cur_bundle_ed, scaled_verts.size(0))
+        cur_query_pts = scaled_verts[cur_bundle_st: cur_bundle_ed]
+        if def_func is not None:
+            cur_query_pts = def_func(cur_query_pts)
+        cur_query_vals = query_func(cur_query_pts)
+        sdf_values.append(cur_query_vals)
+    sdf_values = torch.cat(sdf_values, dim=0)
+    # print(f"queryed sdf values: {sdf_values.size()}") #
+    
+    # GT_sdf_values = np.load("/home/xueyi/diffsim/DiffHand/assets/hand/100_sdf_values.npy", allow_pickle=True)
+    gt_sdf_fn = "/home/xueyi/diffsim/DiffHand/assets/hand/100_sdf_values.npy"
+    if not os.path.exists(gt_sdf_fn):
+        gt_sdf_fn = "/data/xueyi/NeuS/data/100_sdf_values.npy"
+    GT_sdf_values = np.load(gt_sdf_fn, allow_pickle=True)
+    GT_sdf_values = torch.from_numpy(GT_sdf_values).float().to(device)
+    
+    # intrinsic, tet values, pts values, sdf network #
+    triangle_table, num_triangles_table, base_tet_edges, v_id = create_mt_variable(device)
+    tet_table, num_tets_table = create_tetmesh_variables(device)
+    
+    sdf_values = sdf_values.squeeze(-1) # how the rendering # 
+    
+    # print(f"GT_sdf_values: {GT_sdf_values.size()}, sdf_values: {sdf_values.size()}, scaled_verts: {scaled_verts.size()}")
+    # print(f"scaled_verts: {scaled_verts.size()}, ")
+    # pos_nx3, sdf_n, tet_fx4, triangle_table, num_triangles_table, base_tet_edges, v_id,
+    # return_tet_mesh=False, ori_v=None, num_tets_table=None, tet_table=None):
+    # marching_tets_tetmesh ##
+    verts, faces, tet_verts, tets = marching_tets_tetmesh(scaled_verts, sdf_values, indices, triangle_table, num_triangles_table, base_tet_edges, v_id, return_tet_mesh=True, ori_v=scaled_verts, num_tets_table=num_tets_table, tet_table=tet_table)
+    ### use the GT sdf values for the marching tets ###
+    GT_verts,  GT_faces,  GT_tet_verts,  GT_tets = marching_tets_tetmesh(scaled_verts, GT_sdf_values, indices, triangle_table, num_triangles_table, base_tet_edges, v_id, return_tet_mesh=True, ori_v=scaled_verts, num_tets_table=num_tets_table, tet_table=tet_table)
+    
+    # print(f"After tet marching with verts: {verts.size()}, faces: {faces.size()}")
+    return verts, faces, sdf_values, GT_verts, GT_faces  # verts, faces #
+
+
+def  extract_fields_from_tets_selector(bound_min, bound_max, resolution, query_func, def_func=None):
+    # load tet via resolution #
+    # scale them via bounds #
+    # extract the geometry #
+    # /home/xueyi/gen/DeepMetaHandles/data/tets/100_compress.npz # strange #
+    device = bound_min.device
+    # if resolution in [64, 70, 80, 90, 100]:
+    #     tet_fn = f"/home/xueyi/gen/DeepMetaHandles/data/tets/{resolution}_compress.npz"
+    # else:
+    # tet_fn = f"/home/xueyi/gen/DeepMetaHandles/data/tets/{100}_compress.npz"
+    tet_fn = f"/home/xueyi/gen/DeepMetaHandles/data/tets/{100}_compress.npz"
+    if not os.path.exists(tet_fn):
+        tet_fn = f"/data/xueyi/NeuS/data/tets/{100}_compress.npz"
+    tets = np.load(tet_fn)
+    verts = torch.from_numpy(tets['vertices']).float().to(device) # verts positions 
+    indices = torch.from_numpy(tets['tets']).long().to(device) # .to(self.device)
+    # split #
+    # verts; verts; #
+    minn_verts, _ = torch.min(verts, dim=0)
+    maxx_verts, _ = torch.max(verts, dim=0) # (3, ) # exporting the  
+    # scale_verts = maxx_verts - minn_verts
+    scale_bounds = bound_max - bound_min # scale bounds #
+    
+    ### scale the vertices ###
+    scaled_verts = (verts - minn_verts.unsqueeze(0)) / (maxx_verts - minn_verts).unsqueeze(0) ### the maxx and minn verts scales ###
+    
+    # scaled_verts = (verts - minn_verts.unsqueeze(0)) / (maxx_verts - minn_verts).unsqueeze(0) ### the maxx and minn verts scales ###
+
+    scaled_verts = scaled_verts * 2. - 1. # init the sdf filed viathe tet mesh vertices and the sdf values ##
+    # scaled_verts = (scaled_verts * scale_bounds.unsqueeze(0)) + bound_min.unsqueeze(0) ## the scaled verts ###
+    
+    # scaled_verts = scaled_verts - scale_bounds.unsqueeze(0) / 2. # 
+    # scaled_verts = scaled_verts -  bound_min.unsqueeze(0) - scale_bounds.unsqueeze(0) / 2.
+    
+    sdf_values = []
+    N = 64
+    query_bundles = N ** 3 ### N^3
+    query_NNs = scaled_verts.size(0) // query_bundles
+    if query_NNs * query_bundles < scaled_verts.size(0):
+        query_NNs += 1
+    for i_query in range(query_NNs):
+        cur_bundle_st = i_query * query_bundles
+        cur_bundle_ed = (i_query + 1) * query_bundles
+        cur_bundle_ed = min(cur_bundle_ed, scaled_verts.size(0))
+        cur_query_pts = scaled_verts[cur_bundle_st: cur_bundle_ed]
+        if def_func is not None:
+            cur_query_pts, _ = def_func(cur_query_pts)
+        cur_query_vals = query_func(cur_query_pts)
+        sdf_values.append(cur_query_vals)
+    sdf_values = torch.cat(sdf_values, dim=0)
+    # print(f"queryed sdf values: {sdf_values.size()}") #
+    
+    # GT_sdf_values = np.load("/home/xueyi/diffsim/DiffHand/assets/hand/100_sdf_values.npy", allow_pickle=True)
+    gt_sdf_fn = "/home/xueyi/diffsim/DiffHand/assets/hand/100_sdf_values.npy"
+    if not os.path.exists(gt_sdf_fn):
+        gt_sdf_fn = "/data/xueyi/NeuS/data/100_sdf_values.npy"
+    GT_sdf_values = np.load(gt_sdf_fn, allow_pickle=True)
+    GT_sdf_values = torch.from_numpy(GT_sdf_values).float().to(device)
+    
+    # intrinsic, tet values, pts values, sdf network #
+    triangle_table, num_triangles_table, base_tet_edges, v_id = create_mt_variable(device)
+    tet_table, num_tets_table = create_tetmesh_variables(device)
+    
+    sdf_values = sdf_values.squeeze(-1) # how the rendering # 
+    
+    # print(f"GT_sdf_values: {GT_sdf_values.size()}, sdf_values: {sdf_values.size()}, scaled_verts: {scaled_verts.size()}")
+    # print(f"scaled_verts: {scaled_verts.size()}, ")
+    # pos_nx3, sdf_n, tet_fx4, triangle_table, num_triangles_table, base_tet_edges, v_id,
+    # return_tet_mesh=False, ori_v=None, num_tets_table=None, tet_table=None):
+    # marching_tets_tetmesh ##
+    verts, faces, tet_verts, tets = marching_tets_tetmesh(scaled_verts, sdf_values, indices, triangle_table, num_triangles_table, base_tet_edges, v_id, return_tet_mesh=True, ori_v=scaled_verts, num_tets_table=num_tets_table, tet_table=tet_table)
+    ### use the GT sdf values for the marching tets ###
+    GT_verts,  GT_faces,  GT_tet_verts,  GT_tets = marching_tets_tetmesh(scaled_verts, GT_sdf_values, indices, triangle_table, num_triangles_table, base_tet_edges, v_id, return_tet_mesh=True, ori_v=scaled_verts, num_tets_table=num_tets_table, tet_table=tet_table)
+    
+    # print(f"After tet marching with verts: {verts.size()}, faces: {faces.size()}")
+    return verts, faces, sdf_values, GT_verts, GT_faces  # verts, faces #
+
+
+def extract_fields(bound_min, bound_max, resolution, query_func):
+    N = 64
+    X = torch.linspace(bound_min[0], bound_max[0], resolution).split(N)
+    Y = torch.linspace(bound_min[1], bound_max[1], resolution).split(N)
+    Z = torch.linspace(bound_min[2], bound_max[2], resolution).split(N)
+
+    u = np.zeros([resolution, resolution, resolution], dtype=np.float32)
+    with torch.no_grad():
+        for xi, xs in enumerate(X):
+            for yi, ys in enumerate(Y):
+                for zi, zs in enumerate(Z):
+                    xx, yy, zz = torch.meshgrid(xs, ys, zs)
+                    pts = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1)
+                    val = query_func(pts).reshape(len(xs), len(ys), len(zs)).detach().cpu().numpy()
+                    u[xi * N: xi * N + len(xs), yi * N: yi * N + len(ys), zi * N: zi * N + len(zs)] = val
+        # should save u here #
+        # save_u_path = os.path.join("/data2/datasets/diffsim/neus/exp/hand_test/womask_sphere_reverse_value/other_saved", "sdf_values.npy")
+        # np.save(save_u_path, u) # 
+        # print(f"u saved to {save_u_path}")
+    return u
+
+
+def extract_geometry(bound_min, bound_max, resolution, threshold, query_func):
+    print('threshold: {}'.format(threshold))
+    
+    ## using maching cubes ###
+    u = extract_fields(bound_min, bound_max, resolution, query_func)
+    vertices, triangles = mcubes.marching_cubes(u, threshold) # grid sdf and marching cubes #
+    b_max_np = bound_max.detach().cpu().numpy()
+    b_min_np = bound_min.detach().cpu().numpy()
+
+    vertices = vertices / (resolution - 1.0) * (b_max_np - b_min_np)[None, :] + b_min_np[None, :]
+    ### using maching cubes ###
+    
+    ### using marching tets ###
+    # vertices, triangles = extract_fields_from_tets(bound_min, bound_max, resolution, query_func)
+    # vertices = vertices.detach().cpu().numpy()
+    # triangles = triangles.detach().cpu().numpy()
+    ### using marching tets ###
+    
+    # b_max_np = bound_max.detach().cpu().numpy()
+    # b_min_np = bound_min.detach().cpu().numpy()
+
+    # vertices = vertices / (resolution - 1.0) * (b_max_np - b_min_np)[None, :] + b_min_np[None, :]
+    return vertices, triangles
+
+def extract_geometry_tets(bound_min, bound_max, resolution, threshold, query_func, def_func=None, selector=False):
+    # print('threshold: {}'.format(threshold))
+    
+    ### using maching cubes ###
+    # u = extract_fields(bound_min, bound_max, resolution, query_func)
+    # vertices, triangles = mcubes.marching_cubes(u, threshold) # grid sdf and marching cubes #
+    # b_max_np = bound_max.detach().cpu().numpy()
+    # b_min_np = bound_min.detach().cpu().numpy()
+
+    # vertices = vertices / (resolution - 1.0) * (b_max_np - b_min_np)[None, :] + b_min_np[None, :]
+    ### using maching cubes ###
+    
+    ## 
+    ### using marching tets ### fiels from tets ##
+    if selector:
+        vertices, triangles, tet_sdf_values,  GT_verts, GT_faces  = extract_fields_from_tets_selector(bound_min, bound_max, resolution, query_func, def_func=def_func)
+    else:
+        vertices, triangles, tet_sdf_values,  GT_verts, GT_faces  = extract_fields_from_tets(bound_min, bound_max, resolution, query_func, def_func=def_func)
+    # vertices = vertices.detach().cpu().numpy()
+    # triangles = triangles.detach().cpu().numpy()
+    ### using marching tets ###
+    
+    # b_max_np = bound_max.detach().cpu().numpy()
+    # b_min_np = bound_min.detach().cpu().numpy()
+    #
+
+    # vertices = vertices / (resolution - 1.0) * (b_max_np - b_min_np)[None, :] + b_min_np[None, :]
+    return vertices, triangles, tet_sdf_values,  GT_verts, GT_faces 
+
+
+### sample pdfs ###
+def sample_pdf(bins, weights, n_samples, det=False):
+    # This implementation is from NeRF
+    # Get pdf
+    weights = weights + 1e-5  # prevent nans
+    pdf = weights / torch.sum(weights, -1, keepdim=True)
+    cdf = torch.cumsum(pdf, -1)
+    cdf = torch.cat([torch.zeros_like(cdf[..., :1]), cdf], -1)
+    # Take uniform samples
+    if det:
+        u = torch.linspace(0. + 0.5 / n_samples, 1. - 0.5 / n_samples, steps=n_samples)
+        u = u.expand(list(cdf.shape[:-1]) + [n_samples])
+    else:
+        u = torch.rand(list(cdf.shape[:-1]) + [n_samples])
+
+    # Invert CDF # invert cdf # 
+    u = u.contiguous()
+    inds = torch.searchsorted(cdf, u, right=True)
+    below = torch.max(torch.zeros_like(inds - 1), inds - 1)
+    above = torch.min((cdf.shape[-1] - 1) * torch.ones_like(inds), inds)
+    inds_g = torch.stack([below, above], -1)  # (batch, N_samples, 2)
+
+    matched_shape = [inds_g.shape[0], inds_g.shape[1], cdf.shape[-1]]
+    cdf_g = torch.gather(cdf.unsqueeze(1).expand(matched_shape), 2, inds_g)
+    bins_g = torch.gather(bins.unsqueeze(1).expand(matched_shape), 2, inds_g)
+
+    denom = (cdf_g[..., 1] - cdf_g[..., 0])
+    denom = torch.where(denom < 1e-5, torch.ones_like(denom), denom)
+    t = (u - cdf_g[..., 0]) / denom
+    samples = bins_g[..., 0] + t * (bins_g[..., 1] - bins_g[..., 0])
+
+    return samples
+
+
+def load_GT_vertices(GT_meshes_folder):
+    tot_meshes_fns = os.listdir(GT_meshes_folder)
+    tot_meshes_fns = [fn for fn in tot_meshes_fns if fn.endswith(".obj")]
+    tot_mesh_verts = []
+    tot_mesh_faces = []
+    n_tot_verts = 0
+    for fn in tot_meshes_fns:
+        cur_mesh_fn = os.path.join(GT_meshes_folder, fn)
+        obj_mesh = trimesh.load(cur_mesh_fn, process=False)
+        # obj_mesh.remove_degenerate_faces(height=1e-06)
+
+        verts_obj = np.array(obj_mesh.vertices)
+        faces_obj = np.array(obj_mesh.faces)
+        
+        tot_mesh_verts.append(verts_obj)
+        tot_mesh_faces.append(faces_obj + n_tot_verts)
+        n_tot_verts += verts_obj.shape[0]
+        
+        # tot_mesh_faces.append(faces_obj)
+    tot_mesh_verts = np.concatenate(tot_mesh_verts, axis=0)
+    tot_mesh_faces = np.concatenate(tot_mesh_faces, axis=0)
+    return tot_mesh_verts, tot_mesh_faces
+
+
+class NeuSRenderer:
+    def __init__(self,
+                 nerf,
+                 sdf_network,
+                 deviation_network,
+                 color_network,
+                 n_samples,
+                 n_importance,
+                 n_outside,
+                 up_sample_steps,
+                 perturb):
+        self.nerf = nerf # 
+        self.sdf_network = sdf_network
+        self.deviation_network = deviation_network
+        self.color_network = color_network
+        self.n_samples = n_samples
+        self.n_importance = n_importance
+        self.n_outside = n_outside
+        self.up_sample_steps = up_sample_steps
+        self.perturb = perturb
+        
+        GT_meshes_folder = "/home/xueyi/diffsim/DiffHand/assets/hand"
+        if not os.path.exists(GT_meshes_folder):
+            GT_meshes_folder = "/data/xueyi/diffsim/DiffHand/assets/hand"
+        self.mesh_vertices, self.mesh_faces = load_GT_vertices(GT_meshes_folder=GT_meshes_folder)
+        maxx_pts = 25.
+        minn_pts = -15.
+        self.mesh_vertices = (self.mesh_vertices - minn_pts) / (maxx_pts - minn_pts)
+        f = SDF(self.mesh_vertices, self.mesh_faces)
+        self.gt_sdf = f ## a unite sphere or box
+        
+        self.minn_pts = 0
+        self.maxx_pts = 1.
+        
+        # self.minn_pts = -1.5 #
+        # self.maxx_pts = 1.5 #
+        self.bkg_pts = ... # TODO
+        self.cur_fr_bkg_pts_defs = ... # TODO: set the cur_bkg_pts_defs for each frame #
+        self.dist_interp_thres = ... # TODO: set the cur_bkg_pts_defs #
+        
+        self.bending_network = ... # TODO: add the bending network # 
+        self.use_bending_network = ... # TODO: set the property #
+        self.use_delta_bending = ... # TODO
+        self.prev_sdf_network = ... # TODO
+        self.use_selector = False
+        self.timestep_to_passive_mesh = ... # TODO
+        self.timestep_to_active_mesh = ... # TODO
+
+        
+        
+    def deform_pts(self, pts, pts_ts=0, update_tot_def=True): # deform pts # 
+        
+        if self.use_bending_network:
+            if len(pts.size()) == 3:
+                nnb, nns = pts.size(0), pts.size(1)
+                pts_exp = pts.contiguous().view(nnb * nns, -1).contiguous()
+            else:
+                pts_exp = pts
+            # pts_ts #
+            if self.use_delta_bending:
+                
+                if isinstance(self.bending_network, list):
+                    pts_offsets = []
+                    for i_obj, cur_bending_network in enumerate(self.bending_network):
+                        if isinstance(cur_bending_network, fields.BendingNetwork):
+                            for cur_pts_ts in range(pts_ts, -1, -1):
+                                cur_pts_exp = cur_bending_network(pts_exp if cur_pts_ts == pts_ts else cur_pts_exp, input_pts_ts=cur_pts_ts) 
+                        elif isinstance(cur_bending_network, fields.BendingNetworkRigidTrans):
+                            cur_pts_exp = cur_bending_network(pts_exp, input_pts_ts=cur_pts_ts)
+                        else:
+                            raise ValueError('Encountered with unexpected bending network class...')
+                        pts_offsets.append(cur_pts_exp - pts_exp)
+                    pts_offsets = torch.stack(pts_offsets, dim=0)
+                    pts_offsets = torch.sum(pts_offsets, dim=0)
+                    pts_exp = pts_exp + pts_offsets
+                # for cur_pts_ts in range(pts_ts, -1, -1):
+                #     if isinstance(self.bending_network, list): # pts ts # 
+                #         for i_obj, cur_bending_network in enumerate(self.bending_network):
+                #             pts_exp = cur_bending_network(pts_exp, input_pts_ts=cur_pts_ts) 
+                #     else:
+                #         pts_exp = self.bending_network(pts_exp, input_pts_ts=cur_pts_ts)
+            else:
+                if isinstance(self.bending_network, list): # prev sdf network #
+                    pts_offsets = []
+                    for i_obj, cur_bending_network in enumerate(self.bending_network):
+                        bended_pts_exp = cur_bending_network(pts_exp, input_pts_ts=pts_ts)
+                        pts_offsets.append(bended_pts_exp - pts_exp)
+                    pts_offsets = torch.stack(pts_offsets, dim=0)
+                    pts_offsets = torch.sum(pts_offsets, dim=0)
+                    pts_exp = pts_exp + pts_offsets
+                else:
+                    pts_exp = self.bending_network(pts_exp, input_pts_ts=pts_ts)
+            if len(pts.size()) == 3:
+                pts = pts_exp.contiguous().view(nnb, nns, -1).contiguous()
+            else:
+                pts = pts_exp
+            return pts
+            
+        # pts: nn_batch x nn_samples x 3
+        if len(pts.size()) == 3:
+            nnb, nns = pts.size(0), pts.size(1)
+            pts_exp = pts.contiguous().view(nnb * nns, -1).contiguous()
+        else:
+            pts_exp = pts
+        # print(f"prior to deforming: {pts.size()}")
+        
+        dist_pts_to_bkg_pts = torch.sum(
+            (pts_exp.unsqueeze(1) - self.bkg_pts.unsqueeze(0)) ** 2, dim=-1 ## nn_pts_exp x nn_bkg_pts
+        )
+        dist_mask = dist_pts_to_bkg_pts <= self.dist_interp_thres # 
+        dist_mask_float = dist_mask.float()
+        
+        # dist_mask_float #
+        cur_fr_bkg_def_exp = self.cur_fr_bkg_pts_defs.unsqueeze(0).repeat(pts_exp.size(0), 1, 1).contiguous()
+        cur_fr_pts_def = torch.sum(
+            cur_fr_bkg_def_exp * dist_mask_float.unsqueeze(-1), dim=1
+        )
+        dist_mask_float_summ = torch.sum(
+            dist_mask_float, dim=1
+        )
+        dist_mask_float_summ = torch.clamp(dist_mask_float_summ, min=1)
+        cur_fr_pts_def = cur_fr_pts_def / dist_mask_float_summ.unsqueeze(-1) # bkg pts deformation #
+        pts_exp = pts_exp - cur_fr_pts_def
+        if len(pts.size()) == 3:
+            pts = pts_exp.contiguous().view(nnb, nns, -1).contiguous()
+        else:
+            pts = pts_exp
+        return pts # 
+    
+    
+    def deform_pts_with_selector(self, pts, pts_ts=0, update_tot_def=True): # deform pts # 
+        
+        if self.use_bending_network:
+            if len(pts.size()) == 3:
+                nnb, nns = pts.size(0), pts.size(1)
+                pts_exp = pts.contiguous().view(nnb * nns, -1).contiguous()
+            else:
+                pts_exp = pts
+            # pts_ts #
+            if self.use_delta_bending:
+                if isinstance(self.bending_network, list):
+                    bended_pts = []
+                    queries_sdfs_selector = []
+                    for i_obj, cur_bending_network in enumerate(self.bending_network):
+                        if cur_bending_network.use_opt_rigid_translations:
+                            bended_pts_exp = cur_bending_network(pts_exp, input_pts_ts=pts_ts)
+                        else: 
+                            # bended_pts_exp = pts_exp.clone()
+                            if i_obj == 1 and pts_ts == 0:
+                                bended_pts_exp = pts_exp
+                            elif i_obj == 1:
+                                for cur_pts_ts in range(pts_ts, 0, -1): ### before 0 ###
+                                    if isinstance(cur_bending_network, fields.BendingNetwork):
+                                        bended_pts_exp = cur_bending_network(pts_exp if cur_pts_ts == pts_ts else bended_pts_exp, input_pts_ts=cur_pts_ts)
+                                    elif isinstance(cur_bending_network, fields.BendingNetworkForceForward) or isinstance(cur_bending_network, fields.BendingNetworkRigidTransForward):
+                                        # input_pts, input_pts_ts, timestep_to_passive_mesh, act_sdf_net=None, details=None, special_loss_return=False
+                                        bended_pts_exp = cur_bending_network(pts_exp if cur_pts_ts == pts_ts else bended_pts_exp, input_pts_ts=cur_pts_ts, timestep_to_passive_mesh=self.timestep_to_passive_mesh)
+                                    elif isinstance(cur_bending_network, fields.BendingNetworkForceFieldForward):
+                                        # input_pts, input_pts_ts, timestep_to_passive_mesh, passive_sdf_net, details=None, special_loss_return=False
+                                        bended_pts_exp = cur_bending_network(pts_exp if cur_pts_ts == pts_ts else bended_pts_exp, input_pts_ts=cur_pts_ts, timestep_to_passive_mesh=self.timestep_to_passive_mesh, passive_sdf_net=self.sdf_network[1])
+                                    elif isinstance(cur_bending_network, fields.BendingNetworkActiveForceFieldForward):
+                                        # active_bending_net, active_sdf_net, 
+                                        bended_pts_exp = cur_bending_network(pts_exp if cur_pts_ts == pts_ts else bended_pts_exp, input_pts_ts=cur_pts_ts, timestep_to_passive_mesh=self.timestep_to_passive_mesh, passive_sdf_net=self.sdf_network[1], active_bending_net=self.bending_network[0], active_sdf_net=self.sdf_network[0])
+                                    elif isinstance(cur_bending_network, fields.BendingNetworkActiveForceFieldForwardV2):
+                                        # active_bending_net, active_sdf_net, 
+                                        bended_pts_exp = cur_bending_network(pts_exp if cur_pts_ts == pts_ts else bended_pts_exp, input_pts_ts=cur_pts_ts, timestep_to_passive_mesh=self.timestep_to_passive_mesh, passive_sdf_net=self.sdf_network[1], active_bending_net=self.bending_network[0], active_sdf_net=self.sdf_network[0])
+                                    elif isinstance(cur_bending_network, fields.BendingNetworkActiveForceFieldForwardV3):
+                                        # active_bending_net, active_sdf_net, 
+                                        bended_pts_exp = cur_bending_network(pts_exp if cur_pts_ts == pts_ts else bended_pts_exp, input_pts_ts=cur_pts_ts,  timestep_to_passive_mesh=self.timestep_to_passive_mesh, passive_sdf_net=self.sdf_network[1], active_bending_net=self.bending_network[0], active_sdf_net=self.sdf_network[0])
+                                    elif isinstance(cur_bending_network, fields.BendingNetworkActiveForceFieldForwardV4):
+                                        # active_bending_net, active_sdf_net, 
+                                        bended_pts_exp = cur_bending_network(pts_exp if cur_pts_ts == pts_ts else bended_pts_exp, input_pts_ts=cur_pts_ts, timestep_to_active_mesh=self.timestep_to_active_mesh,  timestep_to_passive_mesh=self.timestep_to_passive_mesh, passive_sdf_net=self.sdf_network[1], active_bending_net=self.bending_network[0], active_sdf_net=self.sdf_network[0])
+                                    elif isinstance(cur_bending_network, fields.BendingNetworkActiveForceFieldForwardV5):
+                                        # active_bending_net, active_sdf_net, 
+                                        bended_pts_exp = cur_bending_network(pts_exp if cur_pts_ts == pts_ts else bended_pts_exp, input_pts_ts=cur_pts_ts, timestep_to_active_mesh=self.timestep_to_active_mesh,  timestep_to_passive_mesh=self.timestep_to_passive_mesh, passive_sdf_net=self.sdf_network[1], active_bending_net=self.bending_network[0], active_sdf_net=self.sdf_network[0])
+                                    elif isinstance(cur_bending_network, fields.BendingNetworkActiveForceFieldForwardV6):
+                                        # active_bending_net, active_sdf_net, 
+                                        bended_pts_exp = cur_bending_network(pts_exp if cur_pts_ts == pts_ts else bended_pts_exp, input_pts_ts=cur_pts_ts, timestep_to_active_mesh=self.timestep_to_active_mesh,  timestep_to_passive_mesh=self.timestep_to_passive_mesh, passive_sdf_net=self.sdf_network[1], active_bending_net=self.bending_network[0], active_sdf_net=self.sdf_network[0])
+                                    elif isinstance(cur_bending_network, fields.BendingNetworkActiveForceFieldForwardV7):
+                                        # active_bending_net, active_sdf_net, 
+                                        bended_pts_exp = cur_bending_network(pts_exp if cur_pts_ts == pts_ts else bended_pts_exp, input_pts_ts=cur_pts_ts, timestep_to_active_mesh=self.timestep_to_active_mesh,  timestep_to_passive_mesh=self.timestep_to_passive_mesh, passive_sdf_net=self.sdf_network[1], active_bending_net=self.bending_network[0], active_sdf_net=self.sdf_network[0])
+                                    elif isinstance(cur_bending_network, fields.BendingNetworkActiveForceFieldForwardV8):
+                                        # active_bending_net, active_sdf_net, 
+                                        bended_pts_exp = cur_bending_network(pts_exp if cur_pts_ts == pts_ts else bended_pts_exp, input_pts_ts=cur_pts_ts, timestep_to_active_mesh=self.timestep_to_active_mesh,  timestep_to_passive_mesh=self.timestep_to_passive_mesh, passive_sdf_net=self.sdf_network[1], active_bending_net=self.bending_network[0], active_sdf_net=self.sdf_network[0], update_tot_def=update_tot_def)
+                                    elif isinstance(cur_bending_network, fields.BendingNetworkActiveForceFieldForwardV9):
+                                        # active_bending_net, active_sdf_net, 
+                                        bended_pts_exp = cur_bending_network(pts_exp if cur_pts_ts == pts_ts else bended_pts_exp, input_pts_ts=cur_pts_ts, timestep_to_active_mesh=self.timestep_to_active_mesh,  timestep_to_passive_mesh=self.timestep_to_passive_mesh, passive_sdf_net=self.sdf_network[1], active_bending_net=self.bending_network[0], active_sdf_net=self.sdf_network[0], update_tot_def=update_tot_def)
+                                    else:
+                                        raise ValueError(f"Unrecognized bending network type: {type(cur_bending_network)}")
+                            else:
+                                for cur_pts_ts in range(pts_ts, -1, -1):
+                                    bended_pts_exp = cur_bending_network(pts_exp if cur_pts_ts == pts_ts else bended_pts_exp, input_pts_ts=cur_pts_ts)
+                        # _, cur_bended_pts_selecotr = self.query_pts_sdf_fn_for_selector(bended_pts_exp)
+                        _, cur_bended_pts_selecotr = self.query_pts_sdf_fn_for_selector_ndelta(bended_pts_exp, i_net=i_obj)
+                        bended_pts.append(bended_pts_exp)
+                        queries_sdfs_selector.append(cur_bended_pts_selecotr)
+                    bended_pts = torch.stack(bended_pts, dim=1) # nn_pts x 2 x 3 for bended pts # # bended_pts # 
+                    queries_sdfs_selector = torch.stack(queries_sdfs_selector, dim=1) # nn_pts x 2
+                    # queries_sdfs_selector = (queries_sdfs_selector.sum(dim=1) > 0.5).float().long()
+                    #### get the final sdf_selector from queries_sdfs_selector ####
+                    # sdf_selector = queries_sdfs_selector[:, -1]
+                    # neg_neg = ((queries_sdfs_selector[:, 0] == 0).float() + (queries_sdfs_selector[:, -1] == 1).float()) > 1.5 #### both inside of the object
+                    sdf_selector = 1 - queries_sdfs_selector[:, 0]
+                    # neg_neg
+                    # sdf_selector = queries_sdfs_selector
+                    # delta_sdf, sdf_selector = self.query_pts_sdf_fn_for_selector(pts_exp)
+                    bended_pts = batched_index_select(values=bended_pts, indices=sdf_selector.unsqueeze(1), dim=1).squeeze(1) # nn_pts x 3 #
+                    # print(f"bended_pts: {bended_pts.size()}, pts_exp: {pts_exp.size()}")
+                    # pts_exp = bended_pts.squeeze(1)
+                    pts_exp = bended_pts
+                # for cur_pts_ts in range(pts_ts, -1, -1):
+                #     if isinstance(self.bending_network, list):
+                #         for i_obj, cur_bending_network in enumerate(self.bending_network):
+                #             pts_exp = cur_bending_network(pts_exp, input_pts_ts=cur_pts_ts) 
+                #     else:
+                        
+                #         pts_exp = self.bending_network(pts_exp, input_pts_ts=cur_pts_ts)
+            else:
+                if isinstance(self.bending_network, list): # prev sdf network #
+                    # pts_offsets = []
+                    bended_pts = []
+                    queries_sdfs_selector = []
+                    for i_obj, cur_bending_network in enumerate(self.bending_network):
+                        bended_pts_exp = cur_bending_network(pts_exp, input_pts_ts=pts_ts)
+                        # pts_offsets.append(bended_pts_exp - pts_exp)
+                        _, cur_bended_pts_selecotr = self.query_pts_sdf_fn_for_selector(bended_pts_exp)
+                        bended_pts.append(bended_pts_exp)
+                        queries_sdfs_selector.append(cur_bended_pts_selecotr)
+                    bended_pts = torch.stack(bended_pts, dim=1) # nn_pts x 2 x 3 for bended pts #
+                    queries_sdfs_selector = torch.stack(queries_sdfs_selector, dim=1) # nn_pts x 2
+                    # queries_sdfs_selector = (queries_sdfs_selector.sum(dim=1) > 0.5).float().long()
+                    sdf_selector = queries_sdfs_selector[:, -1]
+                    # sdf_selector = queries_sdfs_selector
+                    
+                    
+                    # delta_sdf, sdf_selector = self.query_pts_sdf_fn_for_selector(pts_exp)
+                    bended_pts = batched_index_select(values=bended_pts, indices=sdf_selector.unsqueeze(1), dim=1).squeeze(1) # nn_pts x 3 #
+                    # print(f"bended_pts: {bended_pts.size()}, pts_exp: {pts_exp.size()}")
+                    pts_exp = bended_pts.squeeze(1)
+                    
+                    # pts_offsets = torch.stack(pts_offsets, dim=0)
+                    # pts_offsets = torch.sum(pts_offsets, dim=0)
+                    # pts_exp = pts_exp + pts_offsets
+                else:
+                    pts_exp = self.bending_network(pts_exp, input_pts_ts=pts_ts)
+            if len(pts.size()) == 3:
+                pts = pts_exp.contiguous().view(nnb, nns, -1).contiguous()
+            else:
+                pts = pts_exp
+            return pts, sdf_selector
+        else:
+            # pts: nn_batch x nn_samples x 3
+            if len(pts.size()) == 3:
+                nnb, nns = pts.size(0), pts.size(1)
+                pts_exp = pts.contiguous().view(nnb * nns, -1).contiguous()
+            else:
+                pts_exp = pts
+            # print(f"prior to deforming: {pts.size()}")
+            
+            dist_pts_to_bkg_pts = torch.sum(
+                (pts_exp.unsqueeze(1) - self.bkg_pts.unsqueeze(0)) ** 2, dim=-1 ## nn_pts_exp x nn_bkg_pts
+            )
+            dist_mask = dist_pts_to_bkg_pts <= self.dist_interp_thres # 
+            dist_mask_float = dist_mask.float()
+            
+            # dist_mask_float #
+            cur_fr_bkg_def_exp = self.cur_fr_bkg_pts_defs.unsqueeze(0).repeat(pts_exp.size(0), 1, 1).contiguous()
+            cur_fr_pts_def = torch.sum(
+                cur_fr_bkg_def_exp * dist_mask_float.unsqueeze(-1), dim=1
+            )
+            dist_mask_float_summ = torch.sum(
+                dist_mask_float, dim=1
+            )
+            dist_mask_float_summ = torch.clamp(dist_mask_float_summ, min=1)
+            cur_fr_pts_def = cur_fr_pts_def / dist_mask_float_summ.unsqueeze(-1) # bkg pts deformation #
+            pts_exp = pts_exp - cur_fr_pts_def
+            if len(pts.size()) == 3:
+                pts = pts_exp.contiguous().view(nnb, nns, -1).contiguous()
+            else:
+                pts = pts_exp
+            return pts # 
+    
+    
+    def deform_pts_passive(self, pts, pts_ts=0):
+        
+        if self.use_bending_network:
+            if len(pts.size()) == 3:
+                nnb, nns = pts.size(0), pts.size(1)
+                pts_exp = pts.contiguous().view(nnb * nns, -1).contiguous()
+            else:
+                pts_exp = pts
+            # pts_ts #
+            if self.use_delta_bending:
+                if pts_ts > 0:
+                    for cur_pts_ts in range(pts_ts, 0, -1):
+                        # if isinstance(self.bending_network, list):
+                        #     for i_obj, cur_bending_network in enumerate(self.bending_network):
+                        #         pts_exp = cur_bending_network(pts_exp, input_pts_ts=cur_pts_ts) 
+                        # else:
+                        if isinstance(self.bending_network[-1], fields.BendingNetwork):
+                            pts_exp = self.bending_network[-1](pts_exp, input_pts_ts=cur_pts_ts)
+                        elif isinstance(self.bending_network[-1], fields.BendingNetworkForceForward) or isinstance(self.bending_network[-1], fields.BendingNetworkRigidTransForward):
+                            pts_exp = self.bending_network[-1](pts_exp, input_pts_ts=cur_pts_ts, timestep_to_passive_mesh=self.timestep_to_passive_mesh)
+                        elif isinstance(self.bending_network[-1], fields.BendingNetworkForceFieldForward):
+                            # input_pts, input_pts_ts, timestep_to_passive_mesh, passive_sdf_net, details=None, special_loss_return=False
+                            pts_exp = self.bending_network[-1](pts_exp, input_pts_ts=cur_pts_ts, timestep_to_passive_mesh=self.timestep_to_passive_mesh, passive_sdf_net=self.sdf_network[1])
+                        elif isinstance(self.bending_network[-1], fields.BendingNetworkActiveForceFieldForward):
+                            pts_exp = self.bending_network[-1](pts_exp, input_pts_ts=cur_pts_ts, timestep_to_passive_mesh=self.timestep_to_passive_mesh, passive_sdf_net=self.sdf_network[1], active_bending_net=self.bending_network[0], active_sdf_net=self.sdf_network[0])
+                        elif isinstance(self.bending_network[-1], fields.BendingNetworkActiveForceFieldForwardV2):
+                            pts_exp = self.bending_network[-1](pts_exp, input_pts_ts=cur_pts_ts, timestep_to_passive_mesh=self.timestep_to_passive_mesh, passive_sdf_net=self.sdf_network[1], active_bending_net=self.bending_network[0], active_sdf_net=self.sdf_network[0])
+                        elif isinstance(self.bending_network[-1], fields.BendingNetworkActiveForceFieldForwardV3):
+                            pts_exp = self.bending_network[-1](pts_exp, input_pts_ts=cur_pts_ts, timestep_to_passive_mesh=self.timestep_to_passive_mesh, passive_sdf_net=self.sdf_network[1], active_bending_net=self.bending_network[0], active_sdf_net=self.sdf_network[0])
+                        elif isinstance(self.bending_network[-1], fields.BendingNetworkActiveForceFieldForwardV4):
+                            pts_exp = self.bending_network[-1](pts_exp, input_pts_ts=cur_pts_ts, timestep_to_active_mesh=self.timestep_to_active_mesh, timestep_to_passive_mesh=self.timestep_to_passive_mesh, passive_sdf_net=self.sdf_network[1], active_bending_net=self.bending_network[0], active_sdf_net=self.sdf_network[0])
+                        elif isinstance(self.bending_network[-1], fields.BendingNetworkActiveForceFieldForwardV5):
+                            pts_exp = self.bending_network[-1](pts_exp, input_pts_ts=cur_pts_ts, timestep_to_active_mesh=self.timestep_to_active_mesh, timestep_to_passive_mesh=self.timestep_to_passive_mesh, passive_sdf_net=self.sdf_network[1], active_bending_net=self.bending_network[0], active_sdf_net=self.sdf_network[0])
+                        elif isinstance(self.bending_network[-1], fields.BendingNetworkActiveForceFieldForwardV6):
+                            pts_exp = self.bending_network[-1](pts_exp, input_pts_ts=cur_pts_ts, timestep_to_active_mesh=self.timestep_to_active_mesh, timestep_to_passive_mesh=self.timestep_to_passive_mesh, passive_sdf_net=self.sdf_network[1], active_bending_net=self.bending_network[0], active_sdf_net=self.sdf_network[0])
+                        elif isinstance(self.bending_network[-1], fields.BendingNetworkActiveForceFieldForwardV7):
+                            pts_exp = self.bending_network[-1](pts_exp, input_pts_ts=cur_pts_ts, timestep_to_active_mesh=self.timestep_to_active_mesh, timestep_to_passive_mesh=self.timestep_to_passive_mesh, passive_sdf_net=self.sdf_network[1], active_bending_net=self.bending_network[0], active_sdf_net=self.sdf_network[0])
+                        elif isinstance(self.bending_network[-1], fields.BendingNetworkActiveForceFieldForwardV8):
+                            pts_exp = self.bending_network[-1](pts_exp, input_pts_ts=cur_pts_ts, timestep_to_active_mesh=self.timestep_to_active_mesh, timestep_to_passive_mesh=self.timestep_to_passive_mesh, passive_sdf_net=self.sdf_network[1], active_bending_net=self.bending_network[0], active_sdf_net=self.sdf_network[0])
+                        elif isinstance(self.bending_network[-1], fields.BendingNetworkActiveForceFieldForwardV9):
+                            pts_exp = self.bending_network[-1](pts_exp, input_pts_ts=cur_pts_ts, timestep_to_active_mesh=self.timestep_to_active_mesh, timestep_to_passive_mesh=self.timestep_to_passive_mesh, passive_sdf_net=self.sdf_network[1], active_bending_net=self.bending_network[0], active_sdf_net=self.sdf_network[0])
+                        else:
+                            raise ValueError(f"Unrecognized bending network type: {type(self.bending_network[-1])}")
+                        # pts_exp = self.bending_network[-1](pts_exp, input_pts_ts=cur_pts_ts)
+            else:
+                # if isinstance(self.bending_network, list):
+                #     pts_offsets = []
+                #     for i_obj, cur_bending_network in enumerate(self.bending_network):
+                #         bended_pts_exp = cur_bending_network(pts_exp, input_pts_ts=pts_ts)
+                #         pts_offsets.append(bended_pts_exp - pts_exp)
+                #     pts_offsets = torch.stack(pts_offsets, dim=0)
+                #     pts_offsets = torch.sum(pts_offsets, dim=0)
+                #     pts_exp = pts_exp + pts_offsets
+                # else:
+                pts_exp = self.bending_network[-1](pts_exp, input_pts_ts=pts_ts)
+            if len(pts.size()) == 3:
+                pts = pts_exp.contiguous().view(nnb, nns, -1).contiguous()
+            else:
+                pts = pts_exp
+            return pts
+            
+        # pts: nn_batch x nn_samples x 3
+        if len(pts.size()) == 3:
+            nnb, nns = pts.size(0), pts.size(1)
+            pts_exp = pts.contiguous().view(nnb * nns, -1).contiguous()
+        else:
+            pts_exp = pts
+        # print(f"prior to deforming: {pts.size()}")
+        
+        dist_pts_to_bkg_pts = torch.sum(
+            (pts_exp.unsqueeze(1) - self.bkg_pts.unsqueeze(0)) ** 2, dim=-1 ## nn_pts_exp x nn_bkg_pts
+        )
+        dist_mask = dist_pts_to_bkg_pts <= self.dist_interp_thres # 
+        dist_mask_float = dist_mask.float()
+        
+        # dist_mask_float #
+        cur_fr_bkg_def_exp = self.cur_fr_bkg_pts_defs.unsqueeze(0).repeat(pts_exp.size(0), 1, 1).contiguous()
+        cur_fr_pts_def = torch.sum(
+            cur_fr_bkg_def_exp * dist_mask_float.unsqueeze(-1), dim=1
+        )
+        dist_mask_float_summ = torch.sum(
+            dist_mask_float, dim=1
+        )
+        dist_mask_float_summ = torch.clamp(dist_mask_float_summ, min=1)
+        cur_fr_pts_def = cur_fr_pts_def / dist_mask_float_summ.unsqueeze(-1) # bkg pts deformation #
+        pts_exp = pts_exp - cur_fr_pts_def
+        if len(pts.size()) == 3:
+            pts = pts_exp.contiguous().view(nnb, nns, -1).contiguous()
+        else:
+            pts = pts_exp
+        return pts # 
+    
+    # delta mesh as passive mesh #
+    def query_pts_sdf_fn_for_selector(self, pts):
+        # for negative 
+        # 1) inside the current mesh but outside the previous mesh ---> negative sdf for this field but positive for another field
+        # 2) negative in thie field and also negative in the previous field ---> 
+        # 2) for positive values of this current field ---> 
+        cur_sdf = self.sdf_network.sdf(pts)
+        prev_sdf = self.prev_sdf_network.sdf(pts)
+        neg_neg = ((cur_sdf < 0.).float() + (prev_sdf < 0.).float()) > 1.5
+        neg_pos = ((cur_sdf < 0.).float() + (prev_sdf >= 0.).float()) > 1.5
+        
+        neg_weq_pos = ((cur_sdf <= 0.).float() + (prev_sdf > 0.).float()) > 1.5
+        
+        pos_neg = ((cur_sdf >= 0.).float() + (prev_sdf < 0.).float()) > 1.5
+        pos_pos = ((cur_sdf >= 0.).float() + (prev_sdf >= 0.).float()) > 1.5
+        res_sdf = torch.zeros_like(cur_sdf)
+        res_sdf[neg_neg] = 1. # 
+        res_sdf[neg_pos] = cur_sdf[neg_pos]
+        res_sdf[pos_neg] = cur_sdf[pos_neg]
+        
+        # inside the residual mesh -> must be neg and pos
+        res_sdf_selector = torch.zeros_like(cur_sdf).long() # 
+        # res_sdf_selector[neg_pos] = 1 # is the residual mesh
+        res_sdf_selector[neg_weq_pos] = 1 
+        # res_sdf_selector[]
+        
+        cat_cur_prev_sdf = torch.stack(
+            [cur_sdf, prev_sdf], dim=-1
+        )
+        minn_cur_prev_sdf, _ = torch.min(cat_cur_prev_sdf, dim=-1)
+        res_sdf[pos_pos] = minn_cur_prev_sdf[pos_pos]
+        
+        return res_sdf, res_sdf_selector
+    
+    def query_pts_sdf_fn_for_selector_ndelta(self, pts, i_net):
+        # for negative 
+        # 1) inside the current mesh but outside the previous mesh ---> negative sdf for this field but positive for another field
+        # 2) negative in thie field and also negative in the previous field ---> 
+        # 2) for positive values of this current field ---> 
+        passive_sdf = self.sdf_network[i_net].sdf(pts).squeeze(-1)
+        passive_sdf_selector = torch.zeros_like(passive_sdf).long()
+        passive_sdf_selector[passive_sdf <= 0.] = 1.
+        return passive_sdf, passive_sdf_selector
+        
+        cur_sdf = self.sdf_network.sdf(pts)
+        prev_sdf = self.prev_sdf_network.sdf(pts)
+        neg_neg = ((cur_sdf < 0.).float() + (prev_sdf < 0.).float()) > 1.5
+        neg_pos = ((cur_sdf < 0.).float() + (prev_sdf >= 0.).float()) > 1.5
+        
+        neg_weq_pos = ((cur_sdf <= 0.).float() + (prev_sdf > 0.).float()) > 1.5
+        
+        pos_neg = ((cur_sdf >= 0.).float() + (prev_sdf < 0.).float()) > 1.5
+        pos_pos = ((cur_sdf >= 0.).float() + (prev_sdf >= 0.).float()) > 1.5
+        res_sdf = torch.zeros_like(cur_sdf)
+        res_sdf[neg_neg] = 1. # 
+        res_sdf[neg_pos] = cur_sdf[neg_pos]
+        res_sdf[pos_neg] = cur_sdf[pos_neg]
+        
+        # inside the residual mesh -> must be neg and pos
+        res_sdf_selector = torch.zeros_like(cur_sdf).long() # 
+        # res_sdf_selector[neg_pos] = 1 # is the residual mesh
+        res_sdf_selector[neg_weq_pos] = 1 
+        # res_sdf_selector[]
+        
+        cat_cur_prev_sdf = torch.stack(
+            [cur_sdf, prev_sdf], dim=-1
+        )
+        minn_cur_prev_sdf, _ = torch.min(cat_cur_prev_sdf, dim=-1)
+        res_sdf[pos_pos] = minn_cur_prev_sdf[pos_pos]
+        
+        return res_sdf, res_sdf_selector
+
+
+    def query_func_sdf(self, pts):
+        if isinstance(self.sdf_network, list):
+            tot_sdf_values = []
+            for i_obj, cur_sdf_network in enumerate(self.sdf_network):
+                cur_sdf_values = cur_sdf_network.sdf(pts)
+                tot_sdf_values.append(cur_sdf_values)
+            tot_sdf_values = torch.stack(tot_sdf_values, dim=-1)
+            tot_sdf_values, _ = torch.min(tot_sdf_values, dim=-1) # totsdf values #
+            sdf = tot_sdf_values
+        else:
+            sdf = self.sdf_network.sdf(pts)
+        return sdf
+    
+    def query_func_sdf_passive(self, pts):
+        # if isinstance(self.sdf_network, list):
+        #     tot_sdf_values = []
+        #     for i_obj, cur_sdf_network in enumerate(self.sdf_network):
+        #         cur_sdf_values = cur_sdf_network.sdf(pts)
+        #         tot_sdf_values.append(cur_sdf_values)
+        #     tot_sdf_values = torch.stack(tot_sdf_values, dim=-1)
+        #     tot_sdf_values, _ = torch.min(tot_sdf_values, dim=-1) # totsdf values #
+        #     sdf = tot_sdf_values
+        # else:
+        sdf = self.sdf_network[-1].sdf(pts)
+            
+        return sdf
+        
+
+    def render_core_outside(self, rays_o, rays_d, z_vals, sample_dist, nerf, background_rgb=None, pts_ts=0):
+        """
+        Render background
+        """
+        batch_size, n_samples = z_vals.shape
+
+        # Section length
+        dists = z_vals[..., 1:] - z_vals[..., :-1]
+        dists = torch.cat([dists, torch.Tensor([sample_dist]).expand(dists[..., :1].shape)], -1)
+        mid_z_vals = z_vals + dists * 0.5
+
+        # Section midpoints #
+        pts = rays_o[:, None, :] + rays_d[:, None, :] * mid_z_vals[..., :, None]  # batch_size, n_samples, 3 #
+        
+        # pts = pts.flip((-1,)) * 2 - 1
+        pts = pts * 2 - 1
+        
+        if self.use_selector:
+            pts, sdf_selector = self.deform_pts_with_selector(pts=pts, pts_ts=pts_ts)
+        else:
+            pts = self.deform_pts(pts=pts, pts_ts=pts_ts)
+
+        dis_to_center = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=True).clip(1.0, 1e10)
+        pts = torch.cat([pts / dis_to_center, 1.0 / dis_to_center], dim=-1)       # batch_size, n_samples, 4 #
+
+        dirs = rays_d[:, None, :].expand(batch_size, n_samples, 3)
+
+        pts = pts.reshape(-1, 3 + int(self.n_outside > 0)) ### deformed_pts ###
+        dirs = dirs.reshape(-1, 3)
+
+        if self.use_selector:
+            tot_density, tot_sampled_color = [], []
+            for i_nerf, cur_nerf in enumerate(nerf):
+                cur_density, cur_sampled_color = cur_nerf(pts, dirs)
+                tot_density.append(cur_density)
+                tot_sampled_color.append(cur_sampled_color)
+            tot_density = torch.stack(tot_density, dim=1)
+            tot_sampled_color = torch.stack(tot_sampled_color, dim=1) ### sampled colors
+            # print(f"tot_density: {tot_density.size()}, tot_sampled_color: {tot_sampled_color.size()}, sdf_selector: {sdf_selector.size()}")
+            density = batched_index_select(values=tot_density, indices=sdf_selector.unsqueeze(-1), dim=1).squeeze(1)
+            sampled_color = batched_index_select(values=tot_sampled_color, indices=sdf_selector.unsqueeze(-1), dim=1).squeeze(1)
+        else:
+            density, sampled_color = nerf(pts, dirs)
+        sampled_color = torch.sigmoid(sampled_color)
+        alpha = 1.0 - torch.exp(-F.softplus(density.reshape(batch_size, n_samples)) * dists)
+        alpha = alpha.reshape(batch_size, n_samples)
+        weights = alpha * torch.cumprod(torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[:, :-1]
+        sampled_color = sampled_color.reshape(batch_size, n_samples, 3)
+        color = (weights[:, :, None] * sampled_color).sum(dim=1)
+        if background_rgb is not None:
+            color = color + background_rgb * (1.0 - weights.sum(dim=-1, keepdim=True))
+
+        return {
+            'color': color,
+            'sampled_color': sampled_color,
+            'alpha': alpha,
+            'weights': weights,
+        }
+
+    def up_sample(self, rays_o, rays_d, z_vals, sdf, n_importance, inv_s, pts_ts=0):
+        """
+        Up sampling give a fixed inv_s
+        """
+        batch_size, n_samples = z_vals.shape
+        pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]  # n_rays, n_samples, 3
+        
+        # pts = pts.flip((-1,)) * 2 - 1
+        pts = pts * 2 - 1
+        
+        if self.use_selector:
+            pts, sdf_selector = self.deform_pts_with_selector(pts=pts, pts_ts=pts_ts)
+        else:
+            pts = self.deform_pts(pts=pts, pts_ts=pts_ts)
+        
+        radius = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=False)
+        inside_sphere = (radius[:, :-1] < 1.0) | (radius[:, 1:] < 1.0)
+        sdf = sdf.reshape(batch_size, n_samples)
+        prev_sdf, next_sdf = sdf[:, :-1], sdf[:, 1:]
+        prev_z_vals, next_z_vals = z_vals[:, :-1], z_vals[:, 1:]
+        mid_sdf = (prev_sdf + next_sdf) * 0.5
+        cos_val = (next_sdf - prev_sdf) / (next_z_vals - prev_z_vals + 1e-5)
+
+        # ----------------------------------------------------------------------------------------------------------
+        # Use min value of [ cos, prev_cos ]
+        # Though it makes the sampling (not rendering) a little bit biased, this strategy can make the sampling more
+        # robust when meeting situations like below:
+        #
+        # SDF
+        # ^
+        # |\          -----x----...
+        # | \        /
+        # |  x      x
+        # |---\----/-------------> 0 level
+        # |    \  /
+        # |     \/
+        # |
+        # ----------------------------------------------------------------------------------------------------------
+        prev_cos_val = torch.cat([torch.zeros([batch_size, 1]), cos_val[:, :-1]], dim=-1)
+        cos_val = torch.stack([prev_cos_val, cos_val], dim=-1)
+        cos_val, _ = torch.min(cos_val, dim=-1, keepdim=False)
+        cos_val = cos_val.clip(-1e3, 0.0) * inside_sphere
+
+        dist = (next_z_vals - prev_z_vals)
+        prev_esti_sdf = mid_sdf - cos_val * dist * 0.5
+        next_esti_sdf = mid_sdf + cos_val * dist * 0.5
+        prev_cdf = torch.sigmoid(prev_esti_sdf * inv_s)
+        next_cdf = torch.sigmoid(next_esti_sdf * inv_s)
+        alpha = (prev_cdf - next_cdf + 1e-5) / (prev_cdf + 1e-5)
+        weights = alpha * torch.cumprod(
+            torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[:, :-1]
+
+        z_samples = sample_pdf(z_vals, weights, n_importance, det=True).detach()
+        return z_samples
+
+    def cat_z_vals(self, rays_o, rays_d, z_vals, new_z_vals, sdf, last=False, pts_ts=0):
+        batch_size, n_samples = z_vals.shape
+        _, n_importance = new_z_vals.shape
+        pts = rays_o[:, None, :] + rays_d[:, None, :] * new_z_vals[..., :, None]
+        
+        # pts = pts.flip((-1,)) * 2 - 1
+        pts = pts * 2 - 1
+        
+        if self.use_selector:
+            pts, sdf_selector = self.deform_pts_with_selector(pts=pts, pts_ts=pts_ts)
+        else:
+            pts = self.deform_pts(pts=pts, pts_ts=pts_ts)
+        
+        z_vals = torch.cat([z_vals, new_z_vals], dim=-1)
+        z_vals, index = torch.sort(z_vals, dim=-1)
+
+        if not last:
+            if isinstance(self.sdf_network, list):
+                tot_new_sdf = []
+                for i_obj, cur_sdf_network in enumerate(self.sdf_network):
+                    cur_new_sdf = cur_sdf_network.sdf(pts.reshape(-1, 3)).reshape(batch_size, n_importance)
+                    tot_new_sdf.append(cur_new_sdf)
+                tot_new_sdf = torch.stack(tot_new_sdf, dim=-1)
+                new_sdf, _ = torch.min(tot_new_sdf, dim=-1) # 
+            else:
+                new_sdf = self.sdf_network.sdf(pts.reshape(-1, 3)).reshape(batch_size, n_importance)
+            sdf = torch.cat([sdf, new_sdf], dim=-1)
+            xx = torch.arange(batch_size)[:, None].expand(batch_size, n_samples + n_importance).reshape(-1)
+            index = index.reshape(-1)
+            sdf = sdf[(xx, index)].reshape(batch_size, n_samples + n_importance)
+
+        return z_vals, sdf
+    
+    
+
+    def render_core(self,
+                    rays_o,
+                    rays_d,
+                    z_vals,
+                    sample_dist,
+                    sdf_network,
+                    deviation_network,
+                    color_network,
+                    background_alpha=None,
+                    background_sampled_color=None,
+                    background_rgb=None,
+                    cos_anneal_ratio=0.0,
+                    pts_ts=0):
+        batch_size, n_samples = z_vals.shape
+
+        # Section length
+        dists = z_vals[..., 1:] - z_vals[..., :-1]
+        dists = torch.cat([dists, torch.Tensor([sample_dist]).expand(dists[..., :1].shape)], -1)
+        mid_z_vals = z_vals + dists * 0.5 # z_vals and dists * 0.5 #
+
+        # Section midpoints
+        pts = rays_o[:, None, :] + rays_d[:, None, :] * mid_z_vals[..., :, None]  # n_rays, n_samples, 3
+        dirs = rays_d[:, None, :].expand(pts.shape)
+
+        pts = pts.reshape(-1, 3) # pts, nn_ou
+        dirs = dirs.reshape(-1, 3)
+        
+        pts = (pts - self.minn_pts) / (self.maxx_pts - self.minn_pts)
+        
+        # pts = pts.flip((-1,)) * 2 - 1
+        pts = pts * 2 - 1
+        
+        if self.use_selector:
+            pts, sdf_selector = self.deform_pts_with_selector(pts=pts, pts_ts=pts_ts)
+        else:
+            pts = self.deform_pts(pts=pts, pts_ts=pts_ts)
+        
+        if isinstance(sdf_network, list):
+            tot_sdf = []
+            tot_feature_vector = []
+            tot_obj_sel = []
+            tot_gradients = []
+            for i_obj, cur_sdf_network in enumerate(sdf_network):
+                cur_sdf_nn_output = cur_sdf_network(pts)
+                cur_sdf, cur_feature_vector = cur_sdf_nn_output[:, :1], cur_sdf_nn_output[:, 1:]
+                tot_sdf.append(cur_sdf)
+                tot_feature_vector.append(cur_feature_vector)
+                
+                gradients = cur_sdf_network.gradient(pts).squeeze()
+                tot_gradients.append(gradients)
+            tot_sdf = torch.stack(tot_sdf, dim=-1)
+            
+            # 
+            if self.use_selector:
+                sdf = batched_index_select(tot_sdf, sdf_selector.unsqueeze(1).unsqueeze(1), dim=2).squeeze(-1)
+                obj_sel = sdf_selector.unsqueeze(1)
+            else:
+                sdf, obj_sel = torch.min(tot_sdf, dim=-1)
+            feature_vector = torch.stack(tot_feature_vector, dim=1)
+            
+            # batched_index_select
+            # print(f"before sel: {feature_vector.size()}, obj_sel: {obj_sel.size()}")
+            feature_vector = batched_index_select(values=feature_vector, indices=obj_sel, dim=1).squeeze(1)
+            
+            # feature_vector = feature_vector[obj_sel.unsqueeze(-1), :].squeeze(1)
+            # print(f"after sel: {feature_vector.size()}")
+            tot_gradients = torch.stack(tot_gradients, dim=1)
+            # gradients = tot_gradients[obj_sel.unsqueeze(-1)].squeeze(1)
+            gradients = batched_index_select(values=tot_gradients, indices=obj_sel, dim=1).squeeze(1)
+            # print(f"gradients: {gradients.size()}, tot_gradients: {tot_gradients.size()}")
+            
+        else:
+            sdf_nn_output = sdf_network(pts)
+            sdf = sdf_nn_output[:, :1]
+            feature_vector = sdf_nn_output[:, 1:]
+            gradients = sdf_network.gradient(pts).squeeze()
+        
+        if self.use_selector:
+            tot_sampled_color = []
+            for i_color_net, cur_color_network in enumerate(color_network):
+                cur_sampled_color = cur_color_network(pts, gradients, dirs, feature_vector) # .reshape(batch_size, n_samples, 3)
+                tot_sampled_color.append(cur_sampled_color)
+            # print(f"tot_density: {tot_density.size()}, tot_sampled_color: {tot_sampled_color.size()}, sdf_selector: {sdf_selector.size()}")
+            tot_sampled_color = torch.stack(tot_sampled_color, dim=1)
+            sampled_color = batched_index_select(values=tot_sampled_color, indices=sdf_selector.unsqueeze(-1), dim=1).squeeze(1).reshape(batch_size, n_samples, 3)
+        else:    
+            sampled_color = color_network(pts, gradients, dirs, feature_vector).reshape(batch_size, n_samples, 3)
+
+        if self.use_selector and isinstance(deviation_network, list):
+            tot_inv_s = []
+            for i_dev_net, cur_deviation_network in enumerate(deviation_network):
+                cur_inv_s = cur_deviation_network(torch.zeros([1, 3]))[:, :1].clip(1e-6, 1e6)
+                tot_inv_s.append(cur_inv_s)
+            tot_inv_s = torch.stack(tot_inv_s, dim=1)
+            inv_s = batched_index_select(values=tot_inv_s, indices=sdf_selector.unsqueeze(-1), dim=1).squeeze(1)
+            # inv_s = 
+        else:
+            # deviation network #
+            inv_s = deviation_network(torch.zeros([1, 3]))[:, :1].clip(1e-6, 1e6)           # Single parameter
+        inv_s = inv_s.expand(batch_size * n_samples, 1)
+
+        true_cos = (dirs * gradients).sum(-1, keepdim=True)
+
+        # "cos_anneal_ratio" grows from 0 to 1 in the beginning training iterations. The anneal strategy below makes
+        # the cos value "not dead" at the beginning training iterations, for better convergence.
+        iter_cos = -(F.relu(-true_cos * 0.5 + 0.5) * (1.0 - cos_anneal_ratio) +
+                     F.relu(-true_cos) * cos_anneal_ratio)  # always non-positive
+
+        # Estimate signed distances at section points
+        estimated_next_sdf = sdf + iter_cos * dists.reshape(-1, 1) * 0.5
+        estimated_prev_sdf = sdf - iter_cos * dists.reshape(-1, 1) * 0.5
+
+        prev_cdf = torch.sigmoid(estimated_prev_sdf * inv_s)
+        next_cdf = torch.sigmoid(estimated_next_sdf * inv_s)
+
+        p = prev_cdf - next_cdf
+        c = prev_cdf
+
+        alpha = ((p + 1e-5) / (c + 1e-5)).reshape(batch_size, n_samples).clip(0.0, 1.0)
+
+        pts_norm = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=True).reshape(batch_size, n_samples)
+        inside_sphere = (pts_norm < 1.0).float().detach()
+        relax_inside_sphere = (pts_norm < 1.2).float().detach()
+
+        # Render with background
+        if background_alpha is not None:
+            alpha = alpha * inside_sphere + background_alpha[:, :n_samples] * (1.0 - inside_sphere)
+            alpha = torch.cat([alpha, background_alpha[:, n_samples:]], dim=-1)
+            sampled_color = sampled_color * inside_sphere[:, :, None] +\
+                            background_sampled_color[:, :n_samples] * (1.0 - inside_sphere)[:, :, None]
+            sampled_color = torch.cat([sampled_color, background_sampled_color[:, n_samples:]], dim=1)
+
+        weights = alpha * torch.cumprod(torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[:, :-1]
+        weights_sum = weights.sum(dim=-1, keepdim=True)
+
+        color = (sampled_color * weights[:, :, None]).sum(dim=1)
+        if background_rgb is not None:    # Fixed background, usually black
+            color = color + background_rgb * (1.0 - weights_sum)
+
+        # Eikonal loss
+        gradient_error = (torch.linalg.norm(gradients.reshape(batch_size, n_samples, 3), ord=2,
+                                            dim=-1) - 1.0) ** 2
+        gradient_error = (relax_inside_sphere * gradient_error).sum() / (relax_inside_sphere.sum() + 1e-5)
+
+        return {
+            'color': color,
+            'sdf': sdf,
+            'dists': dists,
+            'gradients': gradients.reshape(batch_size, n_samples, 3),
+            's_val': 1.0 / inv_s,
+            'mid_z_vals': mid_z_vals,
+            'weights': weights,
+            'cdf': c.reshape(batch_size, n_samples),
+            'gradient_error': gradient_error,
+            'inside_sphere': inside_sphere
+        }
+        
+    def per_sdf_query(self, pts):
+        tot_sdfs = []
+        for i_sdf_net, cur_sdf_network in enumerate(self.sdf_network):
+            cur_sdf_value = cur_sdf_network.sdf(pts).squeeze(-1)
+            tot_sdfs.append(cur_sdf_value)
+        tot_sdfs = torch.stack(tot_sdfs, dim=1)
+        return tot_sdfs
+            
+
+    def render(self, rays_o, rays_d, near, far, pts_ts=0, perturb_overwrite=-1, background_rgb=None, cos_anneal_ratio=0.0, use_gt_sdf=False):
+        batch_size = len(rays_o)
+        sample_dist = 2.0 / self.n_samples  # in a unit sphere # # Assuming the region of interest is a unit sphere
+        z_vals = torch.linspace(0.0, 1.0, self.n_samples) # linspace #
+        z_vals = near + (far - near) * z_vals[None, :]
+
+        z_vals_outside = None
+        if self.n_outside > 0:
+            z_vals_outside = torch.linspace(1e-3, 1.0 - 1.0 / (self.n_outside + 1.0), self.n_outside)
+
+        n_samples = self.n_samples
+        perturb = self.perturb
+
+        if perturb_overwrite >= 0:
+            perturb = perturb_overwrite
+        if perturb > 0:
+            t_rand = (torch.rand([batch_size, 1]) - 0.5)
+            z_vals = z_vals + t_rand * 2.0 / self.n_samples
+
+            if self.n_outside > 0: # z values output # n_outside #
+                mids = .5 * (z_vals_outside[..., 1:] + z_vals_outside[..., :-1])
+                upper = torch.cat([mids, z_vals_outside[..., -1:]], -1)
+                lower = torch.cat([z_vals_outside[..., :1], mids], -1)
+                t_rand = torch.rand([batch_size, z_vals_outside.shape[-1]])
+                z_vals_outside = lower[None, :] + (upper - lower)[None, :] * t_rand
+
+        if self.n_outside > 0:
+            z_vals_outside = far / torch.flip(z_vals_outside, dims=[-1]) + 1.0 / self.n_samples
+
+        background_alpha = None
+        background_sampled_color = None
+
+        # Up sample
+        if self.n_importance > 0:
+            with torch.no_grad():
+                pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]
+                
+                pts = (pts - self.minn_pts) / (self.maxx_pts - self.minn_pts)
+                # sdf = self.sdf_network.sdf(pts.reshape(-1, 3)).reshape(batch_size, self.n_samples)
+                # gt_sdf #
+                
+                # 
+                # pts = ((pts - xyz_min) / (xyz_max - xyz_min)).flip((-1,)) * 2 - 1
+                
+                # pts = pts.flip((-1,)) * 2 - 1
+                pts = pts * 2 - 1                             
+                
+                if self.use_selector:
+                    pts, sdf_selector = self.deform_pts_with_selector(pts=pts, pts_ts=pts_ts)
+                else:
+                    pts = self.deform_pts(pts=pts, pts_ts=pts_ts) # give nthe pts
+                
+                pts_exp = pts.reshape(-1, 3)
+                # minn_pts, _ = torch.min(pts_exp, dim=0)
+                # maxx_pts, _ = torch.max(pts_exp, dim=0) # deformation field (not a rigid one) -> the meshes #
+                # print(f"minn_pts: {minn_pts}, maxx_pts: {maxx_pts}")
+
+                # pts_to_near = pts - near.unsqueeze(1)
+                # maxx_pts = 1.5; minn_pts = -1.5
+                # # maxx_pts = 3; minn_pts = -3
+                # # maxx_pts = 1; minn_pts = -1
+                # pts_exp = (pts_exp - minn_pts) / (maxx_pts - minn_pts)
+                
+                ## render and iamges  ####  
+                # if use_gt_sdf:
+                #     ### use the GT sdf field ####
+                #     # print(f"Using gt sdf :")
+                #     sdf = self.gt_sdf(pts_exp.reshape(-1, 3).detach().cpu().numpy())
+                #     sdf = torch.from_numpy(sdf).float().cuda()
+                #     sdf = sdf.reshape(batch_size, self.n_samples)
+                #     ### use the GT sdf field ####
+                # else:
+                #     # pts_exp: (bsz x nn_s) x 3 -> (sdf_network) -> (bsz x nn_s)
+                #     #### use the optimized sdf field ####
+                    
+                #     # sdf = self.sdf_network.sdf(pts_exp).reshape(batch_size, self.n_samples)
+                    
+                if isinstance(self.sdf_network, list):
+                    if self.use_selector:
+                        tot_sdf_values = []
+                        for i_obj, cur_sdf_network in enumerate(self.sdf_network):
+                            cur_sdf_values = cur_sdf_network.sdf(pts_exp).squeeze(-1)
+                            tot_sdf_values.append(cur_sdf_values)
+                        tot_sdf_values = torch.stack(tot_sdf_values, dim=1)
+                        tot_sdf_values = batched_index_select(tot_sdf_values, indices=sdf_selector.unsqueeze(1), dim=1).squeeze(1)
+                        sdf = tot_sdf_values.reshape(batch_size, self.n_samples)
+                    else:
+                            # tot_sdf_values, _ = torch.min(tot_sdf_values, dim=-1) # totsdf values #
+                        tot_sdf_values = []
+                        for i_obj, cur_sdf_network in enumerate(self.sdf_network):
+                            cur_sdf_values = cur_sdf_network.sdf(pts_exp).reshape(batch_size, self.n_samples)
+                            tot_sdf_values.append(cur_sdf_values)
+                        tot_sdf_values = torch.stack(tot_sdf_values, dim=-1)
+                        tot_sdf_values, _ = torch.min(tot_sdf_values, dim=-1) # totsdf values #
+                        sdf = tot_sdf_values
+                else:
+                    sdf = self.sdf_network.sdf(pts_exp).reshape(batch_size, self.n_samples)
+                    
+                    #### use the optimized sdf field ####
+
+                for i in range(self.up_sample_steps):
+                    new_z_vals = self.up_sample(rays_o,
+                                                rays_d,
+                                                z_vals,
+                                                sdf,
+                                                self.n_importance // self.up_sample_steps,
+                                                64 * 2**i,
+                                                pts_ts=pts_ts)
+                    z_vals, sdf = self.cat_z_vals(rays_o,
+                                                  rays_d,
+                                                  z_vals,
+                                                  new_z_vals,
+                                                  sdf,
+                                                  last=(i + 1 == self.up_sample_steps),
+                                                  pts_ts=pts_ts)
+
+            n_samples = self.n_samples + self.n_importance
+
+        # Background model
+        if self.n_outside > 0:
+            z_vals_feed = torch.cat([z_vals, z_vals_outside], dim=-1)
+            z_vals_feed, _ = torch.sort(z_vals_feed, dim=-1)
+            ret_outside = self.render_core_outside(rays_o, rays_d, z_vals_feed, sample_dist, self.nerf, pts_ts=pts_ts)
+
+            background_sampled_color = ret_outside['sampled_color']
+            background_alpha = ret_outside['alpha']
+
+        tot_sdfs = self.per_sdf_query(pts_exp)
+
+        # Render core
+        ret_fine = self.render_core(rays_o, # 
+                                    rays_d,
+                                    z_vals,
+                                    sample_dist,
+                                    self.sdf_network,
+                                    self.deviation_network,
+                                    self.color_network,
+                                    background_rgb=background_rgb,
+                                    background_alpha=background_alpha,
+                                    background_sampled_color=background_sampled_color,
+                                    cos_anneal_ratio=cos_anneal_ratio,
+                                    pts_ts=pts_ts)
+
+        color_fine = ret_fine['color']
+        weights = ret_fine['weights']
+        weights_sum = weights.sum(dim=-1, keepdim=True)
+        gradients = ret_fine['gradients']
+        s_val = ret_fine['s_val'].reshape(batch_size, n_samples).mean(dim=-1, keepdim=True)
+
+        return {
+            'color_fine': color_fine,
+            's_val': s_val,
+            'cdf_fine': ret_fine['cdf'],
+            'weight_sum': weights_sum,
+            'weight_max': torch.max(weights, dim=-1, keepdim=True)[0],
+            'gradients': gradients,
+            'weights': weights,
+            'gradient_error': ret_fine['gradient_error'],
+            'inside_sphere': ret_fine['inside_sphere'],
+            'tot_sdfs': tot_sdfs,
+        }
+    
+    
+    
+    def render_def(self, rays_o, rays_d, near, far, pts_ts=0, perturb_overwrite=-1, background_rgb=None, cos_anneal_ratio=0.0, use_gt_sdf=False, update_tot_def=True):
+        batch_size = len(rays_o)
+        # sample_dist = 2.0 / self.n_samples  # in a unit sphere # # Assuming the region of interest is a unit sphere
+        z_vals = torch.linspace(0.0, 1.0, self.n_samples)
+        z_vals = near + (far - near) * z_vals[None, :]
+
+        z_vals_outside = None
+        if self.n_outside > 0:
+            z_vals_outside = torch.linspace(1e-3, 1.0 - 1.0 / (self.n_outside + 1.0), self.n_outside)
+
+        n_samples = self.n_samples
+        perturb = self.perturb
+
+        if perturb_overwrite >= 0:
+            perturb = perturb_overwrite
+        if perturb > 0:
+            t_rand = (torch.rand([batch_size, 1]) - 0.5)
+            z_vals = z_vals + t_rand * 2.0 / self.n_samples
+
+            if self.n_outside > 0: # z values output # n_outside #
+                mids = .5 * (z_vals_outside[..., 1:] + z_vals_outside[..., :-1])
+                upper = torch.cat([mids, z_vals_outside[..., -1:]], -1)
+                lower = torch.cat([z_vals_outside[..., :1], mids], -1)
+                t_rand = torch.rand([batch_size, z_vals_outside.shape[-1]])
+                z_vals_outside = lower[None, :] + (upper - lower)[None, :] * t_rand
+
+        if self.n_outside > 0:
+            z_vals_outside = far / torch.flip(z_vals_outside, dims=[-1]) + 1.0 / self.n_samples
+
+        background_alpha = None
+        background_sampled_color = None
+
+        pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]
+                
+        pts = (pts - self.minn_pts) / (self.maxx_pts - self.minn_pts)
+        # sdf = self.sdf_network.sdf(pts.reshape(-1, 3)).reshape(batch_size, self.n_samples)
+        # gt_sdf #
+        
+        # 
+        # pts = ((pts - xyz_min) / (xyz_max - xyz_min)).flip((-1,)) * 2 - 1
+        
+        # pts = pts.flip((-1,)) * 2 - 1
+        pts = pts * 2 - 1                             
+        
+        if self.use_selector:
+            pts, sdf_selector = self.deform_pts_with_selector(pts=pts, pts_ts=pts_ts, update_tot_def=update_tot_def)
+        else:
+            pts = self.deform_pts(pts=pts, pts_ts=pts_ts, update_tot_def=update_tot_def) # give nthe pts
+        
+        return {
+            'defed_pts': pts
+        }
+        # 
+    
+    
+    def  extract_fields_from_tets_selector_self(self, bound_min, bound_max, resolution, i_ts, passive=False):
+        # load tet via resolution #
+        # scale them via bounds #
+        # extract the geometry #
+        # /home/xueyi/gen/DeepMetaHandles/data/tets/100_compress.npz # strange #
+        device = bound_min.device
+        # if resolution in [64, 70, 80, 90, 100]:
+        #     tet_fn = f"/home/xueyi/gen/DeepMetaHandles/data/tets/{resolution}_compress.npz"
+        # else:
+        tet_fn = f"/home/xueyi/gen/DeepMetaHandles/data/tets/{100}_compress.npz"
+        if not os.path.exists(tet_fn):
+            tet_fn = f"/data/xueyi/NeuS/data/tets/{100}_compress.npz"
+        tets = np.load(tet_fn)
+        verts = torch.from_numpy(tets['vertices']).float().to(device) # verts positions 
+        indices = torch.from_numpy(tets['tets']).long().to(device) # .to(self.device)
+        # split #
+        # verts; verts; #
+        minn_verts, _ = torch.min(verts, dim=0)
+        maxx_verts, _ = torch.max(verts, dim=0) # (3, ) # exporting the  
+        # scale_verts = maxx_verts - minn_verts
+        scale_bounds = bound_max - bound_min # scale bounds #
+        
+        ### scale the vertices ###
+        scaled_verts = (verts - minn_verts.unsqueeze(0)) / (maxx_verts - minn_verts).unsqueeze(0) ### the maxx and minn verts scales ###
+        
+        # scaled_verts = (verts - minn_verts.unsqueeze(0)) / (maxx_verts - minn_verts).unsqueeze(0) ### the maxx and minn verts scales ###
+
+        scaled_verts = scaled_verts * 2. - 1. # init the sdf filed viathe tet mesh vertices and the sdf values ##
+        # scaled_verts = (scaled_verts * scale_bounds.unsqueeze(0)) + bound_min.unsqueeze(0) ## the scaled verts ###
+        
+        # scaled_verts = scaled_verts - scale_bounds.unsqueeze(0) / 2. # 
+        # scaled_verts = scaled_verts -  bound_min.unsqueeze(0) - scale_bounds.unsqueeze(0) / 2.
+        
+        sdf_values = []
+        N = 64
+        query_bundles = N ** 3 ### N^3
+        query_NNs = scaled_verts.size(0) // query_bundles
+        if query_NNs * query_bundles < scaled_verts.size(0):
+            query_NNs += 1
+        for i_query in range(query_NNs):
+            cur_bundle_st = i_query * query_bundles
+            cur_bundle_ed = (i_query + 1) * query_bundles
+            cur_bundle_ed = min(cur_bundle_ed, scaled_verts.size(0))
+            cur_query_pts = scaled_verts[cur_bundle_st: cur_bundle_ed]
+            # if def_func is not None:
+            cur_query_pts, sdf_selector = self.deform_pts_with_selector(cur_query_pts, pts_ts=i_ts)
+            # cur_query_pts, _
+            # cur_query_vals = query_func(cur_query_pts)
+            
+            
+            
+            if passive:
+                cur_query_vals = self.sdf_network[1].sdf(cur_query_pts) # .squeeze(-1)
+            else:
+                tot_sdf_values = []
+                for i_obj, cur_sdf_network in enumerate(self.sdf_network):
+                    cur_sdf_values = cur_sdf_network.sdf(cur_query_pts).squeeze(-1)
+                    tot_sdf_values.append(cur_sdf_values)
+                tot_sdf_values = torch.stack(tot_sdf_values, dim=1)
+                tot_sdf_values = batched_index_select(tot_sdf_values, indices=sdf_selector.unsqueeze(1), dim=1).squeeze(1)
+                cur_query_vals = tot_sdf_values.unsqueeze(1)
+            # sdf = tot_sdf_values.reshape(batch_size, self.n_samples)
+            # for i_obj, 
+            sdf_values.append(cur_query_vals)
+        sdf_values = torch.cat(sdf_values, dim=0)
+        # print(f"queryed sdf values: {sdf_values.size()}") #
+        
+        gt_sdf_fn = "/home/xueyi/diffsim/DiffHand/assets/hand/100_sdf_values.npy"
+        if not os.path.exists(gt_sdf_fn):
+            gt_sdf_fn = "/data/xueyi/NeuS/data/100_sdf_values.npy"
+        GT_sdf_values = np.load(gt_sdf_fn, allow_pickle=True)
+        
+        GT_sdf_values = torch.from_numpy(GT_sdf_values).float().to(device)
+        
+        # intrinsic, tet values, pts values, sdf network #
+        triangle_table, num_triangles_table, base_tet_edges, v_id = create_mt_variable(device)
+        tet_table, num_tets_table = create_tetmesh_variables(device)
+        
+        sdf_values = sdf_values.squeeze(-1) # how the rendering # 
+        
+        # print(f"GT_sdf_values: {GT_sdf_values.size()}, sdf_values: {sdf_values.size()}, scaled_verts: {scaled_verts.size()}")
+        # print(f"scaled_verts: {scaled_verts.size()}, ")
+        # pos_nx3, sdf_n, tet_fx4, triangle_table, num_triangles_table, base_tet_edges, v_id,
+        # return_tet_mesh=False, ori_v=None, num_tets_table=None, tet_table=None):
+        # marching_tets_tetmesh ##
+        verts, faces, tet_verts, tets = marching_tets_tetmesh(scaled_verts, sdf_values, indices, triangle_table, num_triangles_table, base_tet_edges, v_id, return_tet_mesh=True, ori_v=scaled_verts, num_tets_table=num_tets_table, tet_table=tet_table)
+        ### use the GT sdf values for the marching tets ###
+        GT_verts,  GT_faces,  GT_tet_verts,  GT_tets = marching_tets_tetmesh(scaled_verts, GT_sdf_values, indices, triangle_table, num_triangles_table, base_tet_edges, v_id, return_tet_mesh=True, ori_v=scaled_verts, num_tets_table=num_tets_table, tet_table=tet_table)
+        
+        # print(f"After tet marching with verts: {verts.size()}, faces: {faces.size()}")
+        return verts, faces, sdf_values, GT_verts, GT_faces  # verts, faces #
+
+
+    def extract_geometry(self, bound_min, bound_max, resolution, threshold=0.0):
+        return extract_geometry(bound_min, # extract geometry #
+                                bound_max,
+                                resolution=resolution,
+                                threshold=threshold,
+                                # query_func=lambda pts: -self.sdf_network.sdf(pts),
+                                query_func=lambda pts: -self.query_func_sdf(pts)
+                                )
+        
+    # if self.deform_pts_with_selector:
+    #         pts = self.deform_pts_with_selector(pts=pts, pts_ts=pts_ts)
+    def extract_geometry_tets(self, bound_min, bound_max, resolution, pts_ts=0, threshold=0.0, wdef=False):
+        if wdef:
+            return extract_geometry_tets(bound_min, # extract geometry #
+                                bound_max,
+                                resolution=resolution,
+                                threshold=threshold,
+                                query_func=lambda pts: -self.query_func_sdf(pts), # lambda pts: -self.sdf_network.sdf(pts),
+                                def_func=lambda pts: self.deform_pts(pts, pts_ts=pts_ts) if not self.use_selector else self.deform_pts_with_selector(pts=pts, pts_ts=pts_ts),
+                                selector=True)
+        else:
+            return extract_geometry_tets(bound_min, # extract geometry #
+                                    bound_max,
+                                    resolution=resolution,
+                                    threshold=threshold,
+                                    # query_func=lambda pts: -self.sdf_network.sdf(pts)
+                                     query_func=lambda pts: -self.query_func_sdf(pts), # lambda pts: -self.sdf_network.sdf(pts),
+                                     selector=True
+                                    )
+            
+    def extract_geometry_tets_passive(self, bound_min, bound_max, resolution, pts_ts=0, threshold=0.0, wdef=False):
+        if wdef:
+            return extract_geometry_tets(bound_min, # extract geometry #
+                                bound_max,
+                                resolution=resolution,
+                                threshold=threshold,
+                                query_func=lambda pts: -self.query_func_sdf_passive(pts), # lambda pts: -self.sdf_network.sdf(pts),
+                                def_func=lambda pts: self.deform_pts_passive(pts, pts_ts=pts_ts),
+                                selector=False
+                                )
+            # return extract_geometry_tets(bound_min, # extract geometry #
+            #                     bound_max,
+            #                     resolution=resolution,
+            #                     threshold=threshold,
+            #                     query_func=lambda pts: -self.query_func_sdf_passive(pts), # lambda pts: -self.sdf_network.sdf(pts),
+            #                     def_func=lambda pts: self.deform_pts(pts, pts_ts=pts_ts) if not self.use_selector else self.deform_pts_with_selector(pts=pts, pts_ts=pts_ts),
+            #                     selector=True)
+        else:
+            return extract_geometry_tets(bound_min, # extract geometry #
+                                    bound_max,
+                                    resolution=resolution,
+                                    threshold=threshold,
+                                    # query_func=lambda pts: -self.sdf_network.sdf(pts)
+                                     query_func=lambda pts: -self.query_func_sdf(pts), # lambda pts: -self.sdf_network.sdf(pts),
+                                     selector=False
+                                    )

models/renderer_def_multi_objs_rigidtrans_forward.py

@@ -0,0 +1,1603 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import logging
+import mcubes
+from icecream import ic
+import os
+
+import trimesh
+from pysdf import SDF
+
+import models.fields as fields
+
+from uni_rep.rep_3d.dmtet import marching_tets_tetmesh, create_tetmesh_variables
+
+def batched_index_select(values, indices, dim = 1):
+  value_dims = values.shape[(dim + 1):]
+  values_shape, indices_shape = map(lambda t: list(t.shape), (values, indices))
+  indices = indices[(..., *((None,) * len(value_dims)))]
+  indices = indices.expand(*((-1,) * len(indices_shape)), *value_dims)
+  value_expand_len = len(indices_shape) - (dim + 1)
+  values = values[(*((slice(None),) * dim), *((None,) * value_expand_len), ...)]
+
+  value_expand_shape = [-1] * len(values.shape)
+  expand_slice = slice(dim, (dim + value_expand_len))
+  value_expand_shape[expand_slice] = indices.shape[expand_slice]
+  values = values.expand(*value_expand_shape)
+
+  dim += value_expand_len
+  return values.gather(dim, indices)
+
+
+def create_mt_variable(device):
+    triangle_table = torch.tensor(
+        [
+            [-1, -1, -1, -1, -1, -1],
+            [1, 0, 2, -1, -1, -1],
+            [4, 0, 3, -1, -1, -1],
+            [1, 4, 2, 1, 3, 4],
+            [3, 1, 5, -1, -1, -1],
+            [2, 3, 0, 2, 5, 3],
+            [1, 4, 0, 1, 5, 4],
+            [4, 2, 5, -1, -1, -1],
+            [4, 5, 2, -1, -1, -1],
+            [4, 1, 0, 4, 5, 1],
+            [3, 2, 0, 3, 5, 2],
+            [1, 3, 5, -1, -1, -1],
+            [4, 1, 2, 4, 3, 1],
+            [3, 0, 4, -1, -1, -1],
+            [2, 0, 1, -1, -1, -1],
+            [-1, -1, -1, -1, -1, -1]
+        ], dtype=torch.long, device=device)
+
+    num_triangles_table = torch.tensor([0, 1, 1, 2, 1, 2, 2, 1, 1, 2, 2, 1, 2, 1, 1, 0], dtype=torch.long, device=device)
+    base_tet_edges = torch.tensor([0, 1, 0, 2, 0, 3, 1, 2, 1, 3, 2, 3], dtype=torch.long, device=device)
+    v_id = torch.pow(2, torch.arange(4, dtype=torch.long, device=device))
+    return triangle_table, num_triangles_table, base_tet_edges, v_id
+
+
+
+def  extract_fields_from_tets(bound_min, bound_max, resolution, query_func, def_func=None):
+    # load tet via resolution #
+    # scale them via bounds #
+    # extract the geometry #
+    # /home/xueyi/gen/DeepMetaHandles/data/tets/100_compress.npz # strange #
+    device = bound_min.device
+    # if resolution in [64, 70, 80, 90, 100]:
+    #     tet_fn = f"/home/xueyi/gen/DeepMetaHandles/data/tets/{resolution}_compress.npz"
+    # else:
+    tet_fn = f"/home/xueyi/gen/DeepMetaHandles/data/tets/{100}_compress.npz"
+    tets = np.load(tet_fn)
+    verts = torch.from_numpy(tets['vertices']).float().to(device) # verts positions 
+    indices = torch.from_numpy(tets['tets']).long().to(device) # .to(self.device)
+    # split #
+    # verts; verts; #
+    minn_verts, _ = torch.min(verts, dim=0)
+    maxx_verts, _ = torch.max(verts, dim=0) # (3, ) # exporting the  
+    # scale_verts = maxx_verts - minn_verts
+    scale_bounds = bound_max - bound_min # scale bounds #
+    
+    ### scale the vertices ###
+    scaled_verts = (verts - minn_verts.unsqueeze(0)) / (maxx_verts - minn_verts).unsqueeze(0) ### the maxx and minn verts scales ###
+    
+    # scaled_verts = (verts - minn_verts.unsqueeze(0)) / (maxx_verts - minn_verts).unsqueeze(0) ### the maxx and minn verts scales ###
+
+    scaled_verts = scaled_verts * 2. - 1. # init the sdf filed viathe tet mesh vertices and the sdf values ##
+    # scaled_verts = (scaled_verts * scale_bounds.unsqueeze(0)) + bound_min.unsqueeze(0) ## the scaled verts ###
+    
+    # scaled_verts = scaled_verts - scale_bounds.unsqueeze(0) / 2. # 
+    # scaled_verts = scaled_verts -  bound_min.unsqueeze(0) - scale_bounds.unsqueeze(0) / 2.
+    
+    sdf_values = []
+    N = 64
+    query_bundles = N ** 3 ### N^3
+    query_NNs = scaled_verts.size(0) // query_bundles
+    if query_NNs * query_bundles < scaled_verts.size(0):
+        query_NNs += 1
+    for i_query in range(query_NNs):
+        cur_bundle_st = i_query * query_bundles
+        cur_bundle_ed = (i_query + 1) * query_bundles
+        cur_bundle_ed = min(cur_bundle_ed, scaled_verts.size(0))
+        cur_query_pts = scaled_verts[cur_bundle_st: cur_bundle_ed]
+        if def_func is not None:
+            # cur_query_pts = def_func(cur_query_pts)
+            cur_query_pts, _ = def_func(cur_query_pts)
+        cur_query_vals = query_func(cur_query_pts)
+        sdf_values.append(cur_query_vals)
+    sdf_values = torch.cat(sdf_values, dim=0)
+    # print(f"queryed sdf values: {sdf_values.size()}") #
+    
+    GT_sdf_values = np.load("/home/xueyi/diffsim/DiffHand/assets/hand/100_sdf_values.npy", allow_pickle=True)
+    GT_sdf_values = torch.from_numpy(GT_sdf_values).float().to(device)
+    
+    # intrinsic, tet values, pts values, sdf network #
+    triangle_table, num_triangles_table, base_tet_edges, v_id = create_mt_variable(device)
+    tet_table, num_tets_table = create_tetmesh_variables(device)
+    
+    sdf_values = sdf_values.squeeze(-1) # how the rendering # 
+    
+    # print(f"GT_sdf_values: {GT_sdf_values.size()}, sdf_values: {sdf_values.size()}, scaled_verts: {scaled_verts.size()}")
+    # print(f"scaled_verts: {scaled_verts.size()}, ")
+    # pos_nx3, sdf_n, tet_fx4, triangle_table, num_triangles_table, base_tet_edges, v_id,
+    # return_tet_mesh=False, ori_v=None, num_tets_table=None, tet_table=None):
+    # marching_tets_tetmesh ##
+    verts, faces, tet_verts, tets = marching_tets_tetmesh(scaled_verts, sdf_values, indices, triangle_table, num_triangles_table, base_tet_edges, v_id, return_tet_mesh=True, ori_v=scaled_verts, num_tets_table=num_tets_table, tet_table=tet_table)
+    ### use the GT sdf values for the marching tets ###
+    GT_verts,  GT_faces,  GT_tet_verts,  GT_tets = marching_tets_tetmesh(scaled_verts, GT_sdf_values, indices, triangle_table, num_triangles_table, base_tet_edges, v_id, return_tet_mesh=True, ori_v=scaled_verts, num_tets_table=num_tets_table, tet_table=tet_table)
+    
+    # print(f"After tet marching with verts: {verts.size()}, faces: {faces.size()}")
+    return verts, faces, sdf_values, GT_verts, GT_faces  # verts, faces #
+
+
+def  extract_fields_from_tets_selector(bound_min, bound_max, resolution, query_func, def_func=None):
+    # load tet via resolution #
+    # scale them via bounds #
+    # extract the geometry #
+    # /home/xueyi/gen/DeepMetaHandles/data/tets/100_compress.npz # strange #
+    device = bound_min.device
+    # if resolution in [64, 70, 80, 90, 100]:
+    #     tet_fn = f"/home/xueyi/gen/DeepMetaHandles/data/tets/{resolution}_compress.npz"
+    # else:
+    tet_fn = f"/home/xueyi/gen/DeepMetaHandles/data/tets/{100}_compress.npz"
+    tets = np.load(tet_fn)
+    verts = torch.from_numpy(tets['vertices']).float().to(device) # verts positions 
+    indices = torch.from_numpy(tets['tets']).long().to(device) # .to(self.device)
+    # split #
+    # verts; verts; #
+    minn_verts, _ = torch.min(verts, dim=0)
+    maxx_verts, _ = torch.max(verts, dim=0) # (3, ) # exporting the  
+    # scale_verts = maxx_verts - minn_verts
+    scale_bounds = bound_max - bound_min # scale bounds #
+    
+    ### scale the vertices ###
+    scaled_verts = (verts - minn_verts.unsqueeze(0)) / (maxx_verts - minn_verts).unsqueeze(0) ### the maxx and minn verts scales ###
+    
+    # scaled_verts = (verts - minn_verts.unsqueeze(0)) / (maxx_verts - minn_verts).unsqueeze(0) ### the maxx and minn verts scales ###
+
+    scaled_verts = scaled_verts * 2. - 1. # init the sdf filed viathe tet mesh vertices and the sdf values ##
+    # scaled_verts = (scaled_verts * scale_bounds.unsqueeze(0)) + bound_min.unsqueeze(0) ## the scaled verts ###
+    
+    # scaled_verts = scaled_verts - scale_bounds.unsqueeze(0) / 2. # 
+    # scaled_verts = scaled_verts -  bound_min.unsqueeze(0) - scale_bounds.unsqueeze(0) / 2.
+    
+    sdf_values = []
+    N = 64
+    query_bundles = N ** 3 ### N^3
+    query_NNs = scaled_verts.size(0) // query_bundles
+    if query_NNs * query_bundles < scaled_verts.size(0):
+        query_NNs += 1
+    for i_query in range(query_NNs):
+        cur_bundle_st = i_query * query_bundles
+        cur_bundle_ed = (i_query + 1) * query_bundles
+        cur_bundle_ed = min(cur_bundle_ed, scaled_verts.size(0))
+        cur_query_pts = scaled_verts[cur_bundle_st: cur_bundle_ed]
+        if def_func is not None:
+            # cur_query_pts = def_func(cur_query_pts)
+            cur_query_pts, _ = def_func(cur_query_pts)
+        cur_query_vals = query_func(cur_query_pts)
+        sdf_values.append(cur_query_vals)
+    sdf_values = torch.cat(sdf_values, dim=0)
+    # print(f"queryed sdf values: {sdf_values.size()}") #
+    
+    GT_sdf_values = np.load("/home/xueyi/diffsim/DiffHand/assets/hand/100_sdf_values.npy", allow_pickle=True)
+    GT_sdf_values = torch.from_numpy(GT_sdf_values).float().to(device)
+    
+    # intrinsic, tet values, pts values, sdf network #
+    triangle_table, num_triangles_table, base_tet_edges, v_id = create_mt_variable(device)
+    tet_table, num_tets_table = create_tetmesh_variables(device)
+    
+    sdf_values = sdf_values.squeeze(-1) # how the rendering # 
+    
+    # print(f"GT_sdf_values: {GT_sdf_values.size()}, sdf_values: {sdf_values.size()}, scaled_verts: {scaled_verts.size()}")
+    # print(f"scaled_verts: {scaled_verts.size()}, ")
+    # pos_nx3, sdf_n, tet_fx4, triangle_table, num_triangles_table, base_tet_edges, v_id,
+    # return_tet_mesh=False, ori_v=None, num_tets_table=None, tet_table=None):
+    # marching_tets_tetmesh ##
+    verts, faces, tet_verts, tets = marching_tets_tetmesh(scaled_verts, sdf_values, indices, triangle_table, num_triangles_table, base_tet_edges, v_id, return_tet_mesh=True, ori_v=scaled_verts, num_tets_table=num_tets_table, tet_table=tet_table)
+    ### use the GT sdf values for the marching tets ###
+    GT_verts,  GT_faces,  GT_tet_verts,  GT_tets = marching_tets_tetmesh(scaled_verts, GT_sdf_values, indices, triangle_table, num_triangles_table, base_tet_edges, v_id, return_tet_mesh=True, ori_v=scaled_verts, num_tets_table=num_tets_table, tet_table=tet_table)
+    
+    # print(f"After tet marching with verts: {verts.size()}, faces: {faces.size()}")
+    return verts, faces, sdf_values, GT_verts, GT_faces  # verts, faces #
+
+
+def extract_fields(bound_min, bound_max, resolution, query_func):
+    N = 64
+    X = torch.linspace(bound_min[0], bound_max[0], resolution).split(N)
+    Y = torch.linspace(bound_min[1], bound_max[1], resolution).split(N)
+    Z = torch.linspace(bound_min[2], bound_max[2], resolution).split(N)
+
+    u = np.zeros([resolution, resolution, resolution], dtype=np.float32)
+    with torch.no_grad():
+        for xi, xs in enumerate(X):
+            for yi, ys in enumerate(Y):
+                for zi, zs in enumerate(Z):
+                    xx, yy, zz = torch.meshgrid(xs, ys, zs)
+                    pts = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1)
+                    val = query_func(pts).reshape(len(xs), len(ys), len(zs)).detach().cpu().numpy()
+                    u[xi * N: xi * N + len(xs), yi * N: yi * N + len(ys), zi * N: zi * N + len(zs)] = val
+        # should save u here #
+        # save_u_path = os.path.join("/data2/datasets/diffsim/neus/exp/hand_test/womask_sphere_reverse_value/other_saved", "sdf_values.npy")
+        # np.save(save_u_path, u)
+        # print(f"u saved to {save_u_path}")
+    return u
+
+
+def extract_geometry(bound_min, bound_max, resolution, threshold, query_func):
+    print('threshold: {}'.format(threshold))
+    
+    ## using maching cubes ###
+    u = extract_fields(bound_min, bound_max, resolution, query_func)
+    vertices, triangles = mcubes.marching_cubes(u, threshold) # grid sdf and marching cubes #
+    b_max_np = bound_max.detach().cpu().numpy()
+    b_min_np = bound_min.detach().cpu().numpy()
+
+    vertices = vertices / (resolution - 1.0) * (b_max_np - b_min_np)[None, :] + b_min_np[None, :]
+    ### using maching cubes ###
+    
+    ### using marching tets ###
+    # vertices, triangles = extract_fields_from_tets(bound_min, bound_max, resolution, query_func)
+    # vertices = vertices.detach().cpu().numpy()
+    # triangles = triangles.detach().cpu().numpy()
+    ### using marching tets ###
+    
+    # b_max_np = bound_max.detach().cpu().numpy()
+    # b_min_np = bound_min.detach().cpu().numpy()
+
+    # vertices = vertices / (resolution - 1.0) * (b_max_np - b_min_np)[None, :] + b_min_np[None, :]
+    return vertices, triangles
+
+def extract_geometry_tets(bound_min, bound_max, resolution, threshold, query_func, def_func=None):
+    # print('threshold: {}'.format(threshold))
+    
+    ### using maching cubes ###
+    # u = extract_fields(bound_min, bound_max, resolution, query_func)
+    # vertices, triangles = mcubes.marching_cubes(u, threshold) # grid sdf and marching cubes #
+    # b_max_np = bound_max.detach().cpu().numpy()
+    # b_min_np = bound_min.detach().cpu().numpy()
+
+    # vertices = vertices / (resolution - 1.0) * (b_max_np - b_min_np)[None, :] + b_min_np[None, :]
+    ### using maching cubes ###
+    
+    ## 
+    ### using marching tets ### fiels from tets ##
+    vertices, triangles, tet_sdf_values,  GT_verts, GT_faces  = extract_fields_from_tets(bound_min, bound_max, resolution, query_func, def_func=def_func)
+    # vertices = vertices.detach().cpu().numpy()
+    # triangles = triangles.detach().cpu().numpy()
+    ### using marching tets ###
+    
+    # b_max_np = bound_max.detach().cpu().numpy()
+    # b_min_np = bound_min.detach().cpu().numpy()
+    #
+
+    # vertices = vertices / (resolution - 1.0) * (b_max_np - b_min_np)[None, :] + b_min_np[None, :]
+    return vertices, triangles, tet_sdf_values,  GT_verts, GT_faces 
+
+
+def sample_pdf(bins, weights, n_samples, det=False):
+    # This implementation is from NeRF
+    # Get pdf
+    weights = weights + 1e-5  # prevent nans
+    pdf = weights / torch.sum(weights, -1, keepdim=True)
+    cdf = torch.cumsum(pdf, -1)
+    cdf = torch.cat([torch.zeros_like(cdf[..., :1]), cdf], -1)
+    # Take uniform samples
+    if det:
+        u = torch.linspace(0. + 0.5 / n_samples, 1. - 0.5 / n_samples, steps=n_samples)
+        u = u.expand(list(cdf.shape[:-1]) + [n_samples])
+    else:
+        u = torch.rand(list(cdf.shape[:-1]) + [n_samples])
+
+    # Invert CDF # invert cdf # 
+    u = u.contiguous()
+    inds = torch.searchsorted(cdf, u, right=True)
+    below = torch.max(torch.zeros_like(inds - 1), inds - 1)
+    above = torch.min((cdf.shape[-1] - 1) * torch.ones_like(inds), inds)
+    inds_g = torch.stack([below, above], -1)  # (batch, N_samples, 2)
+
+    matched_shape = [inds_g.shape[0], inds_g.shape[1], cdf.shape[-1]]
+    cdf_g = torch.gather(cdf.unsqueeze(1).expand(matched_shape), 2, inds_g)
+    bins_g = torch.gather(bins.unsqueeze(1).expand(matched_shape), 2, inds_g)
+
+    denom = (cdf_g[..., 1] - cdf_g[..., 0])
+    denom = torch.where(denom < 1e-5, torch.ones_like(denom), denom)
+    t = (u - cdf_g[..., 0]) / denom
+    samples = bins_g[..., 0] + t * (bins_g[..., 1] - bins_g[..., 0])
+
+    return samples
+
+
+def load_GT_vertices(GT_meshes_folder):
+    tot_meshes_fns = os.listdir(GT_meshes_folder)
+    tot_meshes_fns = [fn for fn in tot_meshes_fns if fn.endswith(".obj")]
+    tot_mesh_verts = []
+    tot_mesh_faces = []
+    n_tot_verts = 0
+    for fn in tot_meshes_fns:
+        cur_mesh_fn = os.path.join(GT_meshes_folder, fn)
+        obj_mesh = trimesh.load(cur_mesh_fn, process=False)
+        # obj_mesh.remove_degenerate_faces(height=1e-06)
+
+        verts_obj = np.array(obj_mesh.vertices)
+        faces_obj = np.array(obj_mesh.faces)
+        
+        tot_mesh_verts.append(verts_obj)
+        tot_mesh_faces.append(faces_obj + n_tot_verts)
+        n_tot_verts += verts_obj.shape[0]
+        
+        # tot_mesh_faces.append(faces_obj)
+    tot_mesh_verts = np.concatenate(tot_mesh_verts, axis=0)
+    tot_mesh_faces = np.concatenate(tot_mesh_faces, axis=0)
+    return tot_mesh_verts, tot_mesh_faces
+
+
+class NeuSRenderer:
+    def __init__(self,
+                 nerf,
+                 sdf_network,
+                 deviation_network,
+                 color_network,
+                 n_samples,
+                 n_importance,
+                 n_outside,
+                 up_sample_steps,
+                 perturb):
+        self.nerf = nerf # multiple sdf networks and deviation networks and xxx #
+        self.sdf_network = sdf_network
+        self.deviation_network = deviation_network
+        self.color_network = color_network
+        self.n_samples = n_samples
+        self.n_importance = n_importance
+        self.n_outside = n_outside
+        self.up_sample_steps = up_sample_steps
+        self.perturb = perturb
+        
+        GT_meshes_folder = "/home/xueyi/diffsim/DiffHand/assets/hand"
+        self.mesh_vertices, self.mesh_faces = load_GT_vertices(GT_meshes_folder=GT_meshes_folder)
+        maxx_pts = 25.
+        minn_pts = -15.
+        self.mesh_vertices = (self.mesh_vertices - minn_pts) / (maxx_pts - minn_pts)
+        f = SDF(self.mesh_vertices, self.mesh_faces)
+        self.gt_sdf = f ## a unite sphere or box
+        
+        self.minn_pts = 0
+        self.maxx_pts = 1.
+        
+        # self.minn_pts = -1.5 # gorudn-truth states with the deformation -> update the sdf value fiedl 
+        # self.maxx_pts = 1.5 # 
+        self.bkg_pts = ... # TODO: the bkg pts # bkg_pts; # bkg_pts_defs #
+        self.cur_fr_bkg_pts_defs = ... # TODO: set the cur_bkg_pts_defs for each frame #
+        self.dist_interp_thres = ... # TODO: set the cur_bkg_pts_defs #
+        
+        self.bending_network = ... # TODO: add the bending network # 
+        self.use_bending_network = ... # TODO: set the property #
+        self.use_delta_bending = ... # TODO
+        self.prev_sdf_network = ... # TODO
+        self.use_selector = False
+        # self.bending_network_rigidtrans_forward = ... ## TODO: set the rigidjtrans forward ###
+        # timestep_to_mesh, timestep_to_passive_mesh, bending_net, bending_net_passive, act_sdf_net, details=None, special_loss_return=False
+        self.timestep_to_mesh = ... ## TODO
+        self.timestep_to_passive_mesh = ... ### TODO 
+        self.bending_net = ... ## TODO
+        self.bending_net_passive = ... ### TODO
+        # self.act_sdf_net = ... ### TODO 
+        self.bending_net_kinematic = ... ### TODO
+        self.time_to_act_joints = ... # ## TODO
+        # use bending network # 
+        # two bending netwrok 
+        # two sdf networks # deform pts kinematic #
+        
+    def deform_pts_kinematic(self, pts, pts_ts=0): # deform pts # 
+        
+        if self.use_bending_network:
+            if len(pts.size()) == 3:
+                nnb, nns = pts.size(0), pts.size(1)
+                pts_exp = pts.contiguous().view(nnb * nns, -1).contiguous()
+            else:
+                pts_exp = pts
+            # pts_ts #
+            if self.use_delta_bending:
+                
+                if isinstance(self.bending_net_kinematic, list):
+                    pts_offsets = [] # 
+                    for i_obj, cur_bending_network in enumerate(self.bending_net_kinematic):
+                        if isinstance(cur_bending_network, fields.BendingNetwork):
+                            for cur_pts_ts in range(pts_ts, -1, -1):
+                                cur_pts_exp = cur_bending_network(pts_exp if cur_pts_ts == pts_ts else cur_pts_exp, input_pts_ts=cur_pts_ts) 
+                        # elif isinstance(cur_bending_network, fields.BendingNetworkForward):
+                        #     cur_pts_exp = cur_bending_network(input_pts=pts_exp, input_pts_ts=cur_pts_ts, timestep_to_mesh=self.timestep_to_mesh, timestep_to_passive_mesh=self.timestep_to_passive_mesh, bending_net=self.bending_net, bending_net_passive=self.bending_net_passive, act_sdf_net=self.prev_sdf_network, details=None, special_loss_return=False)
+                        elif isinstance(cur_bending_network, fields.BendingNetworkRigidTrans):
+                            cur_pts_exp = cur_bending_network(pts_exp, input_pts_ts=cur_pts_ts)
+                        else:
+                            raise ValueError('Encountered with unexpected bending network class...')
+                        pts_offsets.append(cur_pts_exp - pts_exp)
+                    pts_offsets = torch.stack(pts_offsets, dim=0)
+                    pts_offsets = torch.sum(pts_offsets, dim=0)
+                    pts_exp = pts_exp + pts_offsets
+                # for cur_pts_ts in range(pts_ts, -1, -1):
+                #     if isinstance(self.bending_network, list): # pts ts # 
+                #         for i_obj, cur_bending_network in enumerate(self.bending_network):
+                #             pts_exp = cur_bending_network(pts_exp, input_pts_ts=cur_pts_ts) 
+                #     else:
+                #         pts_exp = self.bending_network(pts_exp, input_pts_ts=cur_pts_ts)
+            else:
+                if isinstance(self.bending_net_kinematic, list): # prev sdf network #
+                    pts_offsets = []
+                    for i_obj, cur_bending_network in enumerate(self.bending_net_kinematic):
+                        bended_pts_exp = cur_bending_network(pts_exp, input_pts_ts=pts_ts)
+                        pts_offsets.append(bended_pts_exp - pts_exp)
+                    pts_offsets = torch.stack(pts_offsets, dim=0)
+                    pts_offsets = torch.sum(pts_offsets, dim=0)
+                    pts_exp = pts_exp + pts_offsets
+                else:
+                    pts_exp = self.bending_net_kinematic(pts_exp, input_pts_ts=pts_ts)
+            if len(pts.size()) == 3:
+                pts = pts_exp.contiguous().view(nnb, nns, -1).contiguous()
+            else:
+                pts = pts_exp
+            return pts
+            
+        # pts: nn_batch x nn_samples x 3
+        if len(pts.size()) == 3:
+            nnb, nns = pts.size(0), pts.size(1)
+            pts_exp = pts.contiguous().view(nnb * nns, -1).contiguous()
+        else:
+            pts_exp = pts
+        # print(f"prior to deforming: {pts.size()}")
+        
+        dist_pts_to_bkg_pts = torch.sum(
+            (pts_exp.unsqueeze(1) - self.bkg_pts.unsqueeze(0)) ** 2, dim=-1 ## nn_pts_exp x nn_bkg_pts
+        )
+        dist_mask = dist_pts_to_bkg_pts <= self.dist_interp_thres # 
+        dist_mask_float = dist_mask.float()
+        
+        # dist_mask_float #
+        cur_fr_bkg_def_exp = self.cur_fr_bkg_pts_defs.unsqueeze(0).repeat(pts_exp.size(0), 1, 1).contiguous()
+        cur_fr_pts_def = torch.sum(
+            cur_fr_bkg_def_exp * dist_mask_float.unsqueeze(-1), dim=1
+        )
+        dist_mask_float_summ = torch.sum(
+            dist_mask_float, dim=1
+        )
+        dist_mask_float_summ = torch.clamp(dist_mask_float_summ, min=1)
+        cur_fr_pts_def = cur_fr_pts_def / dist_mask_float_summ.unsqueeze(-1) # bkg pts deformation #
+        pts_exp = pts_exp - cur_fr_pts_def
+        if len(pts.size()) == 3:
+            pts = pts_exp.contiguous().view(nnb, nns, -1).contiguous()
+        else:
+            pts = pts_exp
+        return pts # 
+    
+    
+    def deform_pts_kinematic_active(self, pts, pts_ts=0): # deform pts # 
+        
+        if self.use_bending_network:
+            if len(pts.size()) == 3:
+                nnb, nns = pts.size(0), pts.size(1)
+                pts_exp = pts.contiguous().view(nnb * nns, -1).contiguous()
+            else:
+                pts_exp = pts
+            # pts_ts #
+            if self.use_delta_bending:
+                for cur_pts_ts in range(pts_ts, -1, -1):
+                    pts_exp = self.bending_net(pts_exp, input_pts_ts=cur_pts_ts)
+                # if isinstance(self.bending_net_kinematic, list):
+                #     pts_offsets = [] # 
+                #     for i_obj, cur_bending_network in enumerate(self.bending_net_kinematic):
+                #         if isinstance(cur_bending_network, fields.BendingNetwork):
+                #             for cur_pts_ts in range(pts_ts, -1, -1):
+                #                 cur_pts_exp = cur_bending_network(pts_exp if cur_pts_ts == pts_ts else cur_pts_exp, input_pts_ts=cur_pts_ts) 
+                #         # elif isinstance(cur_bending_network, fields.BendingNetworkForward):
+                #         #     cur_pts_exp = cur_bending_network(input_pts=pts_exp, input_pts_ts=cur_pts_ts, timestep_to_mesh=self.timestep_to_mesh, timestep_to_passive_mesh=self.timestep_to_passive_mesh, bending_net=self.bending_net, bending_net_passive=self.bending_net_passive, act_sdf_net=self.prev_sdf_network, details=None, special_loss_return=False)
+                #         elif isinstance(cur_bending_network, fields.BendingNetworkRigidTrans):
+                #             cur_pts_exp = cur_bending_network(pts_exp, input_pts_ts=cur_pts_ts)
+                #         else:
+                #             raise ValueError('Encountered with unexpected bending network class...')
+                #         pts_offsets.append(cur_pts_exp - pts_exp)
+                #     pts_offsets = torch.stack(pts_offsets, dim=0)
+                #     pts_offsets = torch.sum(pts_offsets, dim=0)
+                #     pts_exp = pts_exp + pts_offsets
+                # for cur_pts_ts in range(pts_ts, -1, -1):
+                #     if isinstance(self.bending_network, list): # pts ts # 
+                #         for i_obj, cur_bending_network in enumerate(self.bending_network):
+                #             pts_exp = cur_bending_network(pts_exp, input_pts_ts=cur_pts_ts) 
+                #     else:
+                #         pts_exp = self.bending_network(pts_exp, input_pts_ts=cur_pts_ts)
+            else:
+                pts_exp = self.bending_net(pts_exp, input_pts_ts=pts_ts)
+                # if isinstance(self.bending_net_kinematic, list): # prev sdf network #
+                #     pts_offsets = []
+                #     for i_obj, cur_bending_network in enumerate(self.bending_net_kinematic):
+                #         bended_pts_exp = cur_bending_network(pts_exp, input_pts_ts=pts_ts)
+                #         pts_offsets.append(bended_pts_exp - pts_exp)
+                #     pts_offsets = torch.stack(pts_offsets, dim=0)
+                #     pts_offsets = torch.sum(pts_offsets, dim=0)
+                #     pts_exp = pts_exp + pts_offsets
+                # else:
+                #     pts_exp = self.bending_net_kinematic(pts_exp, input_pts_ts=pts_ts)
+            if len(pts.size()) == 3:
+                pts = pts_exp.contiguous().view(nnb, nns, -1).contiguous()
+            else:
+                pts = pts_exp
+            return pts
+            
+        # pts: nn_batch x nn_samples x 3
+        if len(pts.size()) == 3:
+            nnb, nns = pts.size(0), pts.size(1)
+            pts_exp = pts.contiguous().view(nnb * nns, -1).contiguous()
+        else:
+            pts_exp = pts
+        # print(f"prior to deforming: {pts.size()}")
+        
+        dist_pts_to_bkg_pts = torch.sum(
+            (pts_exp.unsqueeze(1) - self.bkg_pts.unsqueeze(0)) ** 2, dim=-1 ## nn_pts_exp x nn_bkg_pts
+        )
+        dist_mask = dist_pts_to_bkg_pts <= self.dist_interp_thres # 
+        dist_mask_float = dist_mask.float()
+        
+        # dist_mask_float #
+        cur_fr_bkg_def_exp = self.cur_fr_bkg_pts_defs.unsqueeze(0).repeat(pts_exp.size(0), 1, 1).contiguous()
+        cur_fr_pts_def = torch.sum(
+            cur_fr_bkg_def_exp * dist_mask_float.unsqueeze(-1), dim=1
+        )
+        dist_mask_float_summ = torch.sum(
+            dist_mask_float, dim=1
+        )
+        dist_mask_float_summ = torch.clamp(dist_mask_float_summ, min=1)
+        cur_fr_pts_def = cur_fr_pts_def / dist_mask_float_summ.unsqueeze(-1) # bkg pts deformation #
+        pts_exp = pts_exp - cur_fr_pts_def
+        if len(pts.size()) == 3:
+            pts = pts_exp.contiguous().view(nnb, nns, -1).contiguous()
+        else:
+            pts = pts_exp
+        return pts # 
+    
+    # get the pts and render the pts #
+    # pts and the rendering pts #
+    def deform_pts(self, pts, pts_ts=0): # deform pts # 
+        
+        if self.use_bending_network:
+            if len(pts.size()) == 3:
+                nnb, nns = pts.size(0), pts.size(1)
+                pts_exp = pts.contiguous().view(nnb * nns, -1).contiguous()
+            else:
+                pts_exp = pts
+            # pts_ts #
+            if self.use_delta_bending:
+                
+                if isinstance(self.bending_network, list):
+                    pts_offsets = []
+                    for i_obj, cur_bending_network in enumerate(self.bending_network):
+                        if isinstance(cur_bending_network, fields.BendingNetwork):
+                            for cur_pts_ts in range(pts_ts, -1, -1):
+                                cur_pts_exp = cur_bending_network(pts_exp if cur_pts_ts == pts_ts else cur_pts_exp, input_pts_ts=cur_pts_ts) 
+                        elif isinstance(cur_bending_network, fields.BendingNetworkForward):
+                            for cur_pts_ts in range(pts_ts-1, -1, -1):
+                                cur_pts_exp = cur_bending_network(input_pts=pts_exp if cur_pts_ts == pts_ts else cur_pts_exp, input_pts_ts=cur_pts_ts, timestep_to_mesh=self.timestep_to_mesh, timestep_to_passive_mesh=self.timestep_to_passive_mesh, bending_net=self.bending_net, bending_net_passive=self.bending_net_passive, act_sdf_net=self.prev_sdf_network, details=None, special_loss_return=False)
+                        elif isinstance(cur_bending_network, fields.BendingNetworkForwardJointDyn):
+                            # for cur_pts_ts in range(pts_ts-1, -1, -1):
+                            for cur_pts_ts in range(pts_ts, 0, -1):
+                                cur_pts_exp = cur_bending_network(input_pts=pts_exp if cur_pts_ts == pts_ts else cur_pts_exp, input_pts_ts=cur_pts_ts, timestep_to_mesh=self.timestep_to_mesh, timestep_to_passive_mesh=self.timestep_to_passive_mesh, timestep_to_joints_pts=self.time_to_act_joints, bending_net=self.bending_net, bending_net_passive=self.bending_net_passive, act_sdf_net=self.prev_sdf_network, details=None, special_loss_return=False)
+                        # elif isinstance(cur_bending_network, fields.BendingNetworkRigidTrans):
+                        #     cur_pts_exp = cur_bending_network(pts_exp, input_pts_ts=cur_pts_ts)
+                        else:
+                            raise ValueError('Encountered with unexpected bending network class...')
+                        pts_offsets.append(cur_pts_exp - pts_exp)
+                    pts_offsets = torch.stack(pts_offsets, dim=0)
+                    pts_offsets = torch.sum(pts_offsets, dim=0)
+                    pts_exp = pts_exp + pts_offsets
+                # for cur_pts_ts in range(pts_ts, -1, -1):
+                #     if isinstance(self.bending_network, list): # pts ts # 
+                #         for i_obj, cur_bending_network in enumerate(self.bending_network):
+                #             pts_exp = cur_bending_network(pts_exp, input_pts_ts=cur_pts_ts) 
+                #     else:
+                #         pts_exp = self.bending_network(pts_exp, input_pts_ts=cur_pts_ts)
+            else:
+                if isinstance(self.bending_network, list): # prev sdf network #
+                    pts_offsets = []
+                    for i_obj, cur_bending_network in enumerate(self.bending_network):
+                        bended_pts_exp = cur_bending_network(pts_exp, input_pts_ts=pts_ts)
+                        pts_offsets.append(bended_pts_exp - pts_exp)
+                    pts_offsets = torch.stack(pts_offsets, dim=0)
+                    pts_offsets = torch.sum(pts_offsets, dim=0)
+                    pts_exp = pts_exp + pts_offsets
+                else:
+                    pts_exp = self.bending_network(pts_exp, input_pts_ts=pts_ts)
+            if len(pts.size()) == 3:
+                pts = pts_exp.contiguous().view(nnb, nns, -1).contiguous()
+            else:
+                pts = pts_exp
+            return pts
+            
+        # pts: nn_batch x nn_samples x 3
+        if len(pts.size()) == 3:
+            nnb, nns = pts.size(0), pts.size(1)
+            pts_exp = pts.contiguous().view(nnb * nns, -1).contiguous()
+        else:
+            pts_exp = pts
+        # print(f"prior to deforming: {pts.size()}")
+        
+        dist_pts_to_bkg_pts = torch.sum(
+            (pts_exp.unsqueeze(1) - self.bkg_pts.unsqueeze(0)) ** 2, dim=-1 ## nn_pts_exp x nn_bkg_pts
+        )
+        dist_mask = dist_pts_to_bkg_pts <= self.dist_interp_thres # 
+        dist_mask_float = dist_mask.float()
+        
+        # dist_mask_float #
+        cur_fr_bkg_def_exp = self.cur_fr_bkg_pts_defs.unsqueeze(0).repeat(pts_exp.size(0), 1, 1).contiguous()
+        cur_fr_pts_def = torch.sum(
+            cur_fr_bkg_def_exp * dist_mask_float.unsqueeze(-1), dim=1
+        )
+        dist_mask_float_summ = torch.sum(
+            dist_mask_float, dim=1
+        )
+        dist_mask_float_summ = torch.clamp(dist_mask_float_summ, min=1)
+        cur_fr_pts_def = cur_fr_pts_def / dist_mask_float_summ.unsqueeze(-1) # bkg pts deformation #
+        pts_exp = pts_exp - cur_fr_pts_def
+        if len(pts.size()) == 3:
+            pts = pts_exp.contiguous().view(nnb, nns, -1).contiguous()
+        else:
+            pts = pts_exp
+        return pts # 
+    
+    
+    def deform_pts_with_selector(self, pts, pts_ts=0): # deform pts # 
+        
+        if self.use_bending_network:
+            if len(pts.size()) == 3:
+                nnb, nns = pts.size(0), pts.size(1)
+                pts_exp = pts.contiguous().view(nnb * nns, -1).contiguous()
+            else:
+                pts_exp = pts
+            # pts_ts #
+            if self.use_delta_bending:
+                if isinstance(self.bending_network, list):
+                    bended_pts = []
+                    queries_sdfs_selector = []
+                    for i_obj, cur_bending_network in enumerate(self.bending_network):
+                        # if cur_bending_network.use_opt_rigid_translations:
+                        #     bended_pts_exp = cur_bending_network(pts_exp, input_pts_ts=pts_ts)
+                        # else: 
+                        #     # bended_pts_exp = pts_exp.clone()
+                        #     for cur_pts_ts in range(pts_ts, -1, -1):
+                        #         bended_pts_exp = cur_bending_network(pts_exp if cur_pts_ts == pts_ts else bended_pts_exp, input_pts_ts=cur_pts_ts)
+                        
+                        if isinstance(cur_bending_network, fields.BendingNetwork):
+                            for cur_pts_ts in range(pts_ts, -1, -1):
+                                bended_pts_exp = cur_bending_network(pts_exp if cur_pts_ts == pts_ts else bended_pts_exp, input_pts_ts=cur_pts_ts) 
+                        elif isinstance(cur_bending_network, fields.BendingNetworkForward):
+                            for cur_pts_ts in range(pts_ts-1, -1, -1):
+                                bended_pts_exp = cur_bending_network(input_pts=pts_exp if cur_pts_ts == pts_ts else bended_pts_exp, input_pts_ts=cur_pts_ts, timestep_to_mesh=self.timestep_to_mesh, timestep_to_passive_mesh=self.timestep_to_passive_mesh, bending_net=self.bending_net, bending_net_passive=self.bending_net_passive, act_sdf_net=self.prev_sdf_network, details=None, special_loss_return=False)
+                        elif isinstance(cur_bending_network, fields.BendingNetworkForwardJointDyn):
+                            # for cur_pts_ts in range(pts_ts-1, -1, -1):
+                            for cur_pts_ts in range(pts_ts, 0, -1):
+                                bended_pts_exp = cur_bending_network(input_pts=pts_exp if cur_pts_ts == pts_ts else bended_pts_exp, input_pts_ts=cur_pts_ts, timestep_to_mesh=self.timestep_to_mesh, timestep_to_passive_mesh=self.timestep_to_passive_mesh, timestep_to_joints_pts=self.time_to_act_joints, bending_net=self.bending_net, bending_net_passive=self.bending_net_passive, act_sdf_net=self.prev_sdf_network, details=None, special_loss_return=False)        
+                        
+                        if pts_ts == 0:
+                            bended_pts_exp = pts_exp.clone()
+                        _, cur_bended_pts_selecotr = self.query_pts_sdf_fn_for_selector(bended_pts_exp)
+                        bended_pts.append(bended_pts_exp)
+                        queries_sdfs_selector.append(cur_bended_pts_selecotr)
+                    bended_pts = torch.stack(bended_pts, dim=1) # nn_pts x 2 x 3 for bended pts #
+                    queries_sdfs_selector = torch.stack(queries_sdfs_selector, dim=1) # nn_pts x 2
+                    # queries_sdfs_selector = (queries_sdfs_selector.sum(dim=1) > 0.5).float().long()
+                    sdf_selector = queries_sdfs_selector[:, -1]
+                    # sdf_selector = queries_sdfs_selector
+                    # delta_sdf, sdf_selector = self.query_pts_sdf_fn_for_selector(pts_exp)
+                    # print(f"bended_pts: {bended_pts.size()}, sdf_selector: {sdf_selector.size()}, maxx_sdf_selector: {torch.max(sdf_selector)}, minn_sdf_selector: {torch.min(sdf_selector)}")
+                    bended_pts = batched_index_select(values=bended_pts, indices=sdf_selector.unsqueeze(1), dim=1).squeeze(1) # nn_pts x 3 #
+                    # print(f"bended_pts: {bended_pts.size()}, pts_exp: {pts_exp.size()}")
+                    # pts_exp = bended_pts.squeeze(1)
+                    pts_exp = bended_pts
+                            
+                
+                # for cur_pts_ts in range(pts_ts, -1, -1):
+                #     if isinstance(self.bending_network, list):
+                #         for i_obj, cur_bending_network in enumerate(self.bending_network):
+                #             pts_exp = cur_bending_network(pts_exp, input_pts_ts=cur_pts_ts) 
+                #     else:
+                        
+                #         pts_exp = self.bending_network(pts_exp, input_pts_ts=cur_pts_ts)
+            else:
+                if isinstance(self.bending_network, list): # prev sdf network #
+                    # pts_offsets = []
+                    bended_pts = []
+                    queries_sdfs_selector = []
+                    for i_obj, cur_bending_network in enumerate(self.bending_network):
+                        bended_pts_exp = cur_bending_network(pts_exp, input_pts_ts=pts_ts)
+                        # pts_offsets.append(bended_pts_exp - pts_exp)
+                        _, cur_bended_pts_selecotr = self.query_pts_sdf_fn_for_selector(bended_pts_exp)
+                        bended_pts.append(bended_pts_exp)
+                        queries_sdfs_selector.append(cur_bended_pts_selecotr)
+                    bended_pts = torch.stack(bended_pts, dim=1) # nn_pts x 2 x 3 for bended pts #
+                    queries_sdfs_selector = torch.stack(queries_sdfs_selector, dim=1) # nn_pts x 2
+                    # queries_sdfs_selector = (queries_sdfs_selector.sum(dim=1) > 0.5).float().long()
+                    sdf_selector = queries_sdfs_selector[:, -1]
+                    # sdf_selector = queries_sdfs_selector
+                    
+                    
+                    # delta_sdf, sdf_selector = self.query_pts_sdf_fn_for_selector(pts_exp)
+                    bended_pts = batched_index_select(values=bended_pts, indices=sdf_selector.unsqueeze(1), dim=1).squeeze(1) # nn_pts x 3 #
+                    # print(f"bended_pts: {bended_pts.size()}, pts_exp: {pts_exp.size()}")
+                    pts_exp = bended_pts.squeeze(1)
+                    
+                    # pts_offsets = torch.stack(pts_offsets, dim=0)
+                    # pts_offsets = torch.sum(pts_offsets, dim=0)
+                    # pts_exp = pts_exp + pts_offsets
+                else:
+                    pts_exp = self.bending_network(pts_exp, input_pts_ts=pts_ts)
+            if len(pts.size()) == 3:
+                pts = pts_exp.contiguous().view(nnb, nns, -1).contiguous()
+            else:
+                pts = pts_exp
+            return pts, sdf_selector
+            
+        # pts: nn_batch x nn_samples x 3
+        if len(pts.size()) == 3:
+            nnb, nns = pts.size(0), pts.size(1)
+            pts_exp = pts.contiguous().view(nnb * nns, -1).contiguous()
+        else:
+            pts_exp = pts
+        # print(f"prior to deforming: {pts.size()}")
+        
+        dist_pts_to_bkg_pts = torch.sum(
+            (pts_exp.unsqueeze(1) - self.bkg_pts.unsqueeze(0)) ** 2, dim=-1 ## nn_pts_exp x nn_bkg_pts
+        )
+        dist_mask = dist_pts_to_bkg_pts <= self.dist_interp_thres # 
+        dist_mask_float = dist_mask.float()
+        
+        # dist_mask_float #
+        cur_fr_bkg_def_exp = self.cur_fr_bkg_pts_defs.unsqueeze(0).repeat(pts_exp.size(0), 1, 1).contiguous()
+        cur_fr_pts_def = torch.sum(
+            cur_fr_bkg_def_exp * dist_mask_float.unsqueeze(-1), dim=1
+        )
+        dist_mask_float_summ = torch.sum(
+            dist_mask_float, dim=1
+        )
+        dist_mask_float_summ = torch.clamp(dist_mask_float_summ, min=1)
+        cur_fr_pts_def = cur_fr_pts_def / dist_mask_float_summ.unsqueeze(-1) # bkg pts deformation #
+        pts_exp = pts_exp - cur_fr_pts_def
+        if len(pts.size()) == 3:
+            pts = pts_exp.contiguous().view(nnb, nns, -1).contiguous()
+        else:
+            pts = pts_exp
+        return pts # 
+    
+    
+    def deform_pts_passive(self, pts, pts_ts=0):
+        
+        if self.use_bending_network:
+            if len(pts.size()) == 3:
+                nnb, nns = pts.size(0), pts.size(1)
+                pts_exp = pts.contiguous().view(nnb * nns, -1).contiguous()
+            else:
+                pts_exp = pts
+            # pts_ts #
+            
+            
+            
+            if self.use_delta_bending:
+                cur_bending_network = self.bending_network[-1]
+                if isinstance(cur_bending_network, fields.BendingNetwork):
+                    for cur_pts_ts in range(pts_ts, -1, -1):
+                        cur_pts_exp = cur_bending_network(pts_exp if cur_pts_ts == pts_ts else cur_pts_exp, input_pts_ts=cur_pts_ts) 
+                elif isinstance(cur_bending_network, fields.BendingNetworkForward):
+                    for cur_pts_ts in range(pts_ts-1, -1, -1):
+                        cur_pts_exp = cur_bending_network(input_pts=pts_exp if cur_pts_ts == pts_ts else cur_pts_exp, input_pts_ts=cur_pts_ts, timestep_to_mesh=self.timestep_to_mesh, timestep_to_passive_mesh=self.timestep_to_passive_mesh, bending_net=self.bending_net, bending_net_passive=self.bending_net_passive, act_sdf_net=self.prev_sdf_network, details=None, special_loss_return=False)
+                elif isinstance(cur_bending_network, fields.BendingNetworkForwardJointDyn):
+                    # for cur_pts_ts in range(pts_ts-1, -1, -1):
+                    for cur_pts_ts in range(pts_ts, 0, -1):
+                        cur_pts_exp = cur_bending_network(input_pts=pts_exp if cur_pts_ts == pts_ts else cur_pts_exp, input_pts_ts=cur_pts_ts, timestep_to_mesh=self.timestep_to_mesh, timestep_to_passive_mesh=self.timestep_to_passive_mesh, timestep_to_joints_pts=self.time_to_act_joints, bending_net=self.bending_net, bending_net_passive=self.bending_net_passive, act_sdf_net=self.prev_sdf_network, details=None, special_loss_return=False)
+                # elif isinstance(cur_bending_network, fields.BendingNetworkRigidTrans):
+                #     cur_pts_exp = cur_bending_network(pts_exp, input_pts_ts=cur_pts_ts)
+                else:
+                    raise ValueError('Encountered with unexpected bending network class...')
+                # for cur_pts_ts in range(pts_ts, -1, -1):
+                #     if isinstance(self.bending_network, list):
+                #         for i_obj, cur_bending_network in enumerate(self.bending_network):
+                #             pts_exp = cur_bending_network(pts_exp, input_pts_ts=cur_pts_ts) 
+                #     else:
+                #         pts_exp = self.bending_network(pts_exp, input_pts_ts=cur_pts_ts)
+                # time_to_offset = {
+                #     0:[-0.03338804,  0.07566567,  0.0958022 ],
+                #     1:[-0.05909395,  0.05454276,  0.09974975],
+                #     2: [-0.07214502, -0.00118192,  0.09003166],
+                #     3: [-0.10040219, -0.01334709,  0.08493543],
+                #     4: [-0.10047092, -0.01264334,  0.05320398],
+                #     5: [-0.09152254,  0.00722668,  0.0101514 ],
+                # }
+                
+                if pts_ts > 0:
+                    pts_exp = cur_pts_exp
+                else:
+                    pts_exp = pts_exp
+            else:
+                # if isinstance(self.bending_network, list):
+                #     pts_offsets = []
+                #     for i_obj, cur_bending_network in enumerate(self.bending_network):
+                #         bended_pts_exp = cur_bending_network(pts_exp, input_pts_ts=pts_ts)
+                #         pts_offsets.append(bended_pts_exp - pts_exp)
+                #     pts_offsets = torch.stack(pts_offsets, dim=0)
+                #     pts_offsets = torch.sum(pts_offsets, dim=0)
+                #     pts_exp = pts_exp + pts_offsets
+                # else:
+                pts_exp = self.bending_network[-1](pts_exp, input_pts_ts=pts_ts)
+            if len(pts.size()) == 3:
+                pts = pts_exp.contiguous().view(nnb, nns, -1).contiguous()
+            else:
+                pts = pts_exp
+            return pts
+            
+        # pts: nn_batch x nn_samples x 3
+        if len(pts.size()) == 3:
+            nnb, nns = pts.size(0), pts.size(1)
+            pts_exp = pts.contiguous().view(nnb * nns, -1).contiguous()
+        else:
+            pts_exp = pts
+        # print(f"prior to deforming: {pts.size()}")
+        
+        dist_pts_to_bkg_pts = torch.sum(
+            (pts_exp.unsqueeze(1) - self.bkg_pts.unsqueeze(0)) ** 2, dim=-1 ## nn_pts_exp x nn_bkg_pts
+        )
+        dist_mask = dist_pts_to_bkg_pts <= self.dist_interp_thres # 
+        dist_mask_float = dist_mask.float()
+        
+        # dist_mask_float #
+        cur_fr_bkg_def_exp = self.cur_fr_bkg_pts_defs.unsqueeze(0).repeat(pts_exp.size(0), 1, 1).contiguous()
+        cur_fr_pts_def = torch.sum(
+            cur_fr_bkg_def_exp * dist_mask_float.unsqueeze(-1), dim=1
+        )
+        dist_mask_float_summ = torch.sum(
+            dist_mask_float, dim=1
+        )
+        dist_mask_float_summ = torch.clamp(dist_mask_float_summ, min=1)
+        cur_fr_pts_def = cur_fr_pts_def / dist_mask_float_summ.unsqueeze(-1) # bkg pts deformation #
+        pts_exp = pts_exp - cur_fr_pts_def
+        if len(pts.size()) == 3:
+            pts = pts_exp.contiguous().view(nnb, nns, -1).contiguous()
+        else:
+            pts = pts_exp
+        return pts # 
+    
+    
+    def query_pts_sdf_fn_for_selector(self, pts):
+        # for negative 
+        # 1) inside the current mesh but outside the previous mesh ---> negative sdf for this field but positive for another field
+        # 2) negative in thie field and also negative in the previous field ---> 
+        # 2) for positive values of this current field ---> 
+        # maxx_pts, _ = torch.max(pts, dim=0)
+        # minn_pts, _ = torch.min(pts, dim=0)
+        
+        cur_sdf = self.sdf_network.sdf(pts).squeeze(-1)
+        prev_sdf = self.prev_sdf_network.sdf(pts).squeeze(-1)
+        neg_neg = ((cur_sdf < 0.).float() + (prev_sdf < 0.).float()) > 1.5
+        neg_pos = ((cur_sdf < 0.).float() + (prev_sdf >= 0.).float()) > 1.5
+        
+        neg_weq_pos = ((cur_sdf <= 0.).float() + (prev_sdf > 0.).float()) > 1.5
+        
+        pos_neg = ((cur_sdf >= 0.).float() + (prev_sdf < 0.).float()) > 1.5
+        pos_pos = ((cur_sdf >= 0.).float() + (prev_sdf >= 0.).float()) > 1.5
+        res_sdf = torch.zeros_like(cur_sdf)
+        
+        # print(f"res_sdf: {res_sdf.size()}, neg_neg: {neg_neg.size()}, pts: {pts.size()}, maxx_pts: {maxx_pts}, minn_pts: {minn_pts}")
+        # if torch.sum(neg_neg.float()).item() > 0.:
+        res_sdf[neg_neg] = 1. # 
+        # if torch.sum(neg_pos.float()).item() > 0.:
+        res_sdf[neg_pos] = cur_sdf[neg_pos]
+        # if torch.sum(pos_neg.float()).item() > 0.:
+        res_sdf[pos_neg] = cur_sdf[pos_neg]
+        
+        # inside the residual mesh -> must be neg and pos
+        res_sdf_selector = torch.zeros_like(cur_sdf).long() # 
+        # res_sdf_selector[neg_pos] = 1 # is the residual mesh
+        # if torch.sum(neg_weq_pos.float()).item() > 0.:
+        res_sdf_selector[neg_weq_pos] = 1 
+        # res_sdf_selector[]
+        
+        cat_cur_prev_sdf = torch.stack(
+            [cur_sdf, prev_sdf], dim=-1
+        )
+        minn_cur_prev_sdf, _ = torch.min(cat_cur_prev_sdf, dim=-1)
+        
+        if torch.sum(pos_pos.float()).item() > 0.:
+            res_sdf[pos_pos] = minn_cur_prev_sdf[pos_pos]
+        
+        return res_sdf, res_sdf_selector
+    
+    def query_func_sdf(self, pts):
+        # if isinstance(self.sdf_network, list):
+        #     tot_sdf_values = []
+        #     for i_obj, cur_sdf_network in enumerate(self.sdf_network):
+        #         cur_sdf_values = cur_sdf_network.sdf(pts)
+        #         tot_sdf_values.append(cur_sdf_values)
+        #     tot_sdf_values = torch.stack(tot_sdf_values, dim=-1)
+        #     tot_sdf_values, _ = torch.min(tot_sdf_values, dim=-1) # totsdf values #
+        #     sdf = tot_sdf_values
+        # else:
+        #     sdf = self.sdf_network.sdf(pts)
+        
+            
+        cur_sdf = self.sdf_network.sdf(pts)
+        prev_sdf = self.prev_sdf_network.sdf(pts)
+        neg_neg = ((cur_sdf < 0.).float() + (prev_sdf < 0.).float()) > 1.5
+        neg_pos = ((cur_sdf < 0.).float() + (prev_sdf >= 0.).float()) > 1.5
+        
+        neg_weq_pos = ((cur_sdf <= 0.).float() + (prev_sdf > 0.).float()) > 1.5
+        
+        pos_neg = ((cur_sdf >= 0.).float() + (prev_sdf < 0.).float()) > 1.5
+        pos_pos = ((cur_sdf >= 0.).float() + (prev_sdf >= 0.).float()) > 1.5
+        res_sdf = torch.zeros_like(cur_sdf)
+        
+        # print(f"res_sdf: {res_sdf.size()}, neg_neg: {neg_neg.size()}, pts: {pts.size()}, maxx_pts: {maxx_pts}, minn_pts: {minn_pts}")
+        # if torch.sum(neg_neg.float()).item() > 0.:
+        res_sdf[neg_neg] = prev_sdf[neg_neg]
+        # if torch.sum(neg_pos.float()).item() > 0.:
+        res_sdf[neg_pos] = cur_sdf[neg_pos]
+        # if torch.sum(pos_neg.float()).item() > 0.:
+        res_sdf[pos_neg] = cur_sdf[pos_neg]
+        
+        # inside the residual mesh -> must be neg and pos
+        # res_sdf_selector = torch.zeros_like(cur_sdf).long() # 
+        # res_sdf_selector[neg_pos] = 1 # is the residual mesh
+        # if torch.sum(neg_weq_pos.float()).item() > 0.:
+        # res_sdf_selector[neg_weq_pos] = 1 
+        # res_sdf_selector[]
+        
+        cat_cur_prev_sdf = torch.stack(
+            [cur_sdf, prev_sdf], dim=-1
+        )
+        minn_cur_prev_sdf, _ = torch.min(cat_cur_prev_sdf, dim=-1)
+        
+        # if torch.sum(pos_pos.float()).item() > 0.:
+        res_sdf[pos_pos] = minn_cur_prev_sdf[pos_pos]
+        
+        return res_sdf
+    
+    def query_func_active(self, pts):
+        # if isinstance(self.sdf_network, list):
+        #     tot_sdf_values = []
+        #     for i_obj, cur_sdf_network in enumerate(self.sdf_network):
+        #         cur_sdf_values = cur_sdf_network.sdf(pts)
+        #         tot_sdf_values.append(cur_sdf_values)
+        #     tot_sdf_values = torch.stack(tot_sdf_values, dim=-1)
+        #     tot_sdf_values, _ = torch.min(tot_sdf_values, dim=-1) # totsdf values #
+        #     sdf = tot_sdf_values
+        # else:
+        sdf = self.prev_sdf_network.sdf(pts)
+        return sdf
+    
+    def query_func_sdf_passive(self, pts):
+        # if isinstance(self.sdf_network, list):
+        #     tot_sdf_values = []
+        #     for i_obj, cur_sdf_network in enumerate(self.sdf_network):
+        #         cur_sdf_values = cur_sdf_network.sdf(pts)
+        #         tot_sdf_values.append(cur_sdf_values)
+        #     tot_sdf_values = torch.stack(tot_sdf_values, dim=-1)
+        #     tot_sdf_values, _ = torch.min(tot_sdf_values, dim=-1) # totsdf values #
+        #     sdf = tot_sdf_values
+        # else:
+        sdf = self.sdf_network[-1].sdf(pts)
+            
+        return sdf
+        
+
+    def render_core_outside(self, rays_o, rays_d, z_vals, sample_dist, nerf, background_rgb=None, pts_ts=0):
+        """
+        Render background
+        """
+        batch_size, n_samples = z_vals.shape
+
+        # Section length
+        dists = z_vals[..., 1:] - z_vals[..., :-1]
+        dists = torch.cat([dists, torch.Tensor([sample_dist]).expand(dists[..., :1].shape)], -1)
+        mid_z_vals = z_vals + dists * 0.5
+
+        # Section midpoints #
+        pts = rays_o[:, None, :] + rays_d[:, None, :] * mid_z_vals[..., :, None]  # batch_size, n_samples, 3 #
+        
+        # pts = pts.flip((-1,)) * 2 - 1
+        pts = pts * 2 - 1
+        
+        if self.use_selector:
+            pts, sdf_selector = self.deform_pts_with_selector(pts=pts, pts_ts=pts_ts)
+        else:
+            pts = self.deform_pts(pts=pts, pts_ts=pts_ts)
+
+        dis_to_center = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=True).clip(1.0, 1e10)
+        pts = torch.cat([pts / dis_to_center, 1.0 / dis_to_center], dim=-1)       # batch_size, n_samples, 4 #
+
+        dirs = rays_d[:, None, :].expand(batch_size, n_samples, 3)
+
+        pts = pts.reshape(-1, 3 + int(self.n_outside > 0))
+        dirs = dirs.reshape(-1, 3)
+
+        density, sampled_color = nerf(pts, dirs)
+        sampled_color = torch.sigmoid(sampled_color)
+        alpha = 1.0 - torch.exp(-F.softplus(density.reshape(batch_size, n_samples)) * dists)
+        alpha = alpha.reshape(batch_size, n_samples)
+        weights = alpha * torch.cumprod(torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[:, :-1]
+        sampled_color = sampled_color.reshape(batch_size, n_samples, 3)
+        color = (weights[:, :, None] * sampled_color).sum(dim=1)
+        if background_rgb is not None:
+            color = color + background_rgb * (1.0 - weights.sum(dim=-1, keepdim=True))
+
+        return {
+            'color': color,
+            'sampled_color': sampled_color,
+            'alpha': alpha,
+            'weights': weights,
+        }
+
+    def up_sample(self, rays_o, rays_d, z_vals, sdf, n_importance, inv_s, pts_ts=0):
+        """
+        Up sampling give a fixed inv_s
+        """
+        batch_size, n_samples = z_vals.shape
+        pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]  # n_rays, n_samples, 3
+        
+        # pts = pts.flip((-1,)) * 2 - 1
+        pts = pts * 2 - 1
+        
+        if self.use_selector:
+            pts, sdf_selector = self.deform_pts_with_selector(pts=pts, pts_ts=pts_ts)
+        else:
+            pts = self.deform_pts(pts=pts, pts_ts=pts_ts)
+        
+        radius = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=False)
+        inside_sphere = (radius[:, :-1] < 1.0) | (radius[:, 1:] < 1.0)
+        sdf = sdf.reshape(batch_size, n_samples)
+        prev_sdf, next_sdf = sdf[:, :-1], sdf[:, 1:]
+        prev_z_vals, next_z_vals = z_vals[:, :-1], z_vals[:, 1:]
+        mid_sdf = (prev_sdf + next_sdf) * 0.5
+        cos_val = (next_sdf - prev_sdf) / (next_z_vals - prev_z_vals + 1e-5)
+
+        # ----------------------------------------------------------------------------------------------------------
+        # Use min value of [ cos, prev_cos ]
+        # Though it makes the sampling (not rendering) a little bit biased, this strategy can make the sampling more
+        # robust when meeting situations like below:
+        #
+        # SDF
+        # ^
+        # |\          -----x----...
+        # | \        /
+        # |  x      x
+        # |---\----/-------------> 0 level
+        # |    \  /
+        # |     \/
+        # |
+        # ----------------------------------------------------------------------------------------------------------
+        prev_cos_val = torch.cat([torch.zeros([batch_size, 1]), cos_val[:, :-1]], dim=-1)
+        cos_val = torch.stack([prev_cos_val, cos_val], dim=-1)
+        cos_val, _ = torch.min(cos_val, dim=-1, keepdim=False)
+        cos_val = cos_val.clip(-1e3, 0.0) * inside_sphere
+
+        dist = (next_z_vals - prev_z_vals)
+        prev_esti_sdf = mid_sdf - cos_val * dist * 0.5
+        next_esti_sdf = mid_sdf + cos_val * dist * 0.5
+        prev_cdf = torch.sigmoid(prev_esti_sdf * inv_s)
+        next_cdf = torch.sigmoid(next_esti_sdf * inv_s)
+        alpha = (prev_cdf - next_cdf + 1e-5) / (prev_cdf + 1e-5)
+        weights = alpha * torch.cumprod(
+            torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[:, :-1]
+
+        z_samples = sample_pdf(z_vals, weights, n_importance, det=True).detach()
+        return z_samples
+
+    def cat_z_vals(self, rays_o, rays_d, z_vals, new_z_vals, sdf, last=False, pts_ts=0):
+        batch_size, n_samples = z_vals.shape
+        _, n_importance = new_z_vals.shape
+        pts = rays_o[:, None, :] + rays_d[:, None, :] * new_z_vals[..., :, None]
+        
+        # pts = pts.flip((-1,)) * 2 - 1
+        pts = pts * 2 - 1
+        
+        if self.use_selector:
+            pts, sdf_selector = self.deform_pts_with_selector(pts=pts, pts_ts=pts_ts)
+        else:
+            pts = self.deform_pts(pts=pts, pts_ts=pts_ts)
+        
+        z_vals = torch.cat([z_vals, new_z_vals], dim=-1)
+        z_vals, index = torch.sort(z_vals, dim=-1)
+
+        if not last:
+            if isinstance(self.sdf_network, list):
+                tot_new_sdf = []
+                for i_obj, cur_sdf_network in enumerate(self.sdf_network):
+                    cur_new_sdf = cur_sdf_network.sdf(pts.reshape(-1, 3)).reshape(batch_size, n_importance)
+                    tot_new_sdf.append(cur_new_sdf)
+                tot_new_sdf = torch.stack(tot_new_sdf, dim=-1)
+                new_sdf, _ = torch.min(tot_new_sdf, dim=-1) # 
+            else:
+                if self.use_selector:
+                    new_sdf_cur = self.sdf_network.sdf(pts.reshape(-1, 3)) # .reshape(batch_size, n_importance)
+                    new_sdf_prev = self.prev_sdf_network.sdf(pts.reshape(-1, 3)) # .reshape(batch_size, n_importance)
+                    new_sdf = torch.stack([new_sdf_prev, new_sdf_cur], dim=1)
+                    new_sdf = batched_index_select(new_sdf,  sdf_selector.unsqueeze(-1), dim=1).squeeze(1)
+                    new_sdf = new_sdf.reshape(batch_size, n_importance)
+                else:
+                    new_sdf = self.sdf_network.sdf(pts.reshape(-1, 3)).reshape(batch_size, n_importance)
+            sdf = torch.cat([sdf, new_sdf], dim=-1)
+            xx = torch.arange(batch_size)[:, None].expand(batch_size, n_samples + n_importance).reshape(-1)
+            index = index.reshape(-1)
+            sdf = sdf[(xx, index)].reshape(batch_size, n_samples + n_importance)
+
+        return z_vals, sdf
+    
+    
+
+    def render_core(self,
+                    rays_o,
+                    rays_d,
+                    z_vals,
+                    sample_dist,
+                    sdf_network,
+                    deviation_network,
+                    color_network,
+                    background_alpha=None,
+                    background_sampled_color=None,
+                    background_rgb=None,
+                    cos_anneal_ratio=0.0,
+                    pts_ts=0):
+        batch_size, n_samples = z_vals.shape
+
+        # Section length
+        dists = z_vals[..., 1:] - z_vals[..., :-1]
+        dists = torch.cat([dists, torch.Tensor([sample_dist]).expand(dists[..., :1].shape)], -1)
+        mid_z_vals = z_vals + dists * 0.5 # z_vals and dists * 0.5 #
+
+        # Section midpoints
+        pts = rays_o[:, None, :] + rays_d[:, None, :] * mid_z_vals[..., :, None]  # n_rays, n_samples, 3
+        dirs = rays_d[:, None, :].expand(pts.shape)
+
+        pts = pts.reshape(-1, 3) # pts, nn_ou
+        dirs = dirs.reshape(-1, 3)
+        
+        pts = (pts - self.minn_pts) / (self.maxx_pts - self.minn_pts)
+        
+        # pts = pts.flip((-1,)) * 2 - 1
+        pts = pts * 2 - 1
+        
+        if self.use_selector:
+            pts, sdf_selector = self.deform_pts_with_selector(pts=pts, pts_ts=pts_ts)
+        else:
+            pts = self.deform_pts(pts=pts, pts_ts=pts_ts)
+        
+        if isinstance(sdf_network, list):
+            tot_sdf = []
+            tot_feature_vector = []
+            tot_obj_sel = []
+            tot_gradients = []
+            for i_obj, cur_sdf_network in enumerate(sdf_network):
+                cur_sdf_nn_output = cur_sdf_network(pts)
+                cur_sdf, cur_feature_vector = cur_sdf_nn_output[:, :1], cur_sdf_nn_output[:, 1:]
+                tot_sdf.append(cur_sdf)
+                tot_feature_vector.append(cur_feature_vector)
+                
+                gradients = cur_sdf_network.gradient(pts).squeeze()
+                tot_gradients.append(gradients)
+            tot_sdf = torch.stack(tot_sdf, dim=-1)
+            sdf, obj_sel = torch.min(tot_sdf, dim=-1)
+            feature_vector = torch.stack(tot_feature_vector, dim=1)
+            
+            # batched_index_select
+            # print(f"before sel: {feature_vector.size()}, obj_sel: {obj_sel.size()}")
+            feature_vector = batched_index_select(values=feature_vector, indices=obj_sel, dim=1).squeeze(1)
+            
+            
+            # feature_vector = feature_vector[obj_sel.unsqueeze(-1), :].squeeze(1)
+            # print(f"after sel: {feature_vector.size()}")
+            tot_gradients = torch.stack(tot_gradients, dim=1)
+            # gradients = tot_gradients[obj_sel.unsqueeze(-1)].squeeze(1)
+            gradients = batched_index_select(values=tot_gradients, indices=obj_sel, dim=1).squeeze(1)
+            # print(f"gradients: {gradients.size()}, tot_gradients: {tot_gradients.size()}")
+            
+        else:
+            # sdf_nn_output = sdf_network(pts)
+            # sdf = sdf_nn_output[:, :1]
+            # feature_vector = sdf_nn_output[:, 1:]
+            # gradients = sdf_network.gradient(pts).squeeze()
+            
+            if self.use_selector:
+                prev_sdf_nn_output = self.prev_sdf_network(pts)
+                prev_gradients = self.prev_sdf_network.gradient(pts).squeeze()
+                cur_sdf_nn_output = self.sdf_network(pts)
+                cur_gradients = self.sdf_network.gradient(pts).squeeze()
+                
+                sdf_nn_output = torch.stack([prev_sdf_nn_output, cur_sdf_nn_output], dim=1)
+                sdf_nn_output = batched_index_select(sdf_nn_output,  sdf_selector.unsqueeze(-1), dim=1).squeeze(1)
+                
+                sdf = sdf_nn_output[:, :1]
+                feature_vector = sdf_nn_output[:, 1:]
+                
+                gradients = torch.stack([prev_gradients, cur_gradients], dim=1)
+                gradients = batched_index_select(gradients,  sdf_selector.unsqueeze(-1), dim=1).squeeze(1)
+            else:
+                sdf_nn_output = sdf_network(pts)
+                sdf = sdf_nn_output[:, :1]
+                feature_vector = sdf_nn_output[:, 1:]
+                gradients = sdf_network.gradient(pts).squeeze()
+            # new_sdf_cur = self.sdf_network.sdf(pts.reshape(-1, 3)).reshape(batch_size, n_importance)
+            #         new_sdf_prev = self.prev_sdf_network.sdf(pts.reshape(-1, 3)).reshape(batch_size, n_importance)
+            #         new_sdf = torch.stack([new_sdf_prev, new_sdf_cur], dim=1)
+            #         new_sdf = batched_index_select(new_sdf,  sdf_selector.unsqueeze(-1), dim=1).squeeze(1)
+            
+        sampled_color = color_network(pts, gradients, dirs, feature_vector).reshape(batch_size, n_samples, 3)
+
+        # deviation network #
+        inv_s = deviation_network(torch.zeros([1, 3]))[:, :1].clip(1e-6, 1e6)           # Single parameter
+        inv_s = inv_s.expand(batch_size * n_samples, 1)
+
+        true_cos = (dirs * gradients).sum(-1, keepdim=True)
+
+        # "cos_anneal_ratio" grows from 0 to 1 in the beginning training iterations. The anneal strategy below makes
+        # the cos value "not dead" at the beginning training iterations, for better convergence.
+        iter_cos = -(F.relu(-true_cos * 0.5 + 0.5) * (1.0 - cos_anneal_ratio) +
+                     F.relu(-true_cos) * cos_anneal_ratio)  # always non-positive
+
+        # Estimate signed distances at section points
+        estimated_next_sdf = sdf + iter_cos * dists.reshape(-1, 1) * 0.5
+        estimated_prev_sdf = sdf - iter_cos * dists.reshape(-1, 1) * 0.5
+
+        prev_cdf = torch.sigmoid(estimated_prev_sdf * inv_s)
+        next_cdf = torch.sigmoid(estimated_next_sdf * inv_s)
+
+        p = prev_cdf - next_cdf
+        c = prev_cdf
+
+        alpha = ((p + 1e-5) / (c + 1e-5)).reshape(batch_size, n_samples).clip(0.0, 1.0)
+
+        pts_norm = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=True).reshape(batch_size, n_samples)
+        inside_sphere = (pts_norm < 1.0).float().detach()
+        relax_inside_sphere = (pts_norm < 1.2).float().detach()
+
+        # Render with background
+        if background_alpha is not None:
+            alpha = alpha * inside_sphere + background_alpha[:, :n_samples] * (1.0 - inside_sphere)
+            alpha = torch.cat([alpha, background_alpha[:, n_samples:]], dim=-1)
+            sampled_color = sampled_color * inside_sphere[:, :, None] +\
+                            background_sampled_color[:, :n_samples] * (1.0 - inside_sphere)[:, :, None]
+            sampled_color = torch.cat([sampled_color, background_sampled_color[:, n_samples:]], dim=1)
+
+        weights = alpha * torch.cumprod(torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[:, :-1]
+        weights_sum = weights.sum(dim=-1, keepdim=True)
+
+        color = (sampled_color * weights[:, :, None]).sum(dim=1)
+        if background_rgb is not None:    # Fixed background, usually black
+            color = color + background_rgb * (1.0 - weights_sum)
+
+        # Eikonal loss
+        gradient_error = (torch.linalg.norm(gradients.reshape(batch_size, n_samples, 3), ord=2,
+                                            dim=-1) - 1.0) ** 2
+        gradient_error = (relax_inside_sphere * gradient_error).sum() / (relax_inside_sphere.sum() + 1e-5)
+
+        return {
+            'color': color,
+            'sdf': sdf,
+            'dists': dists,
+            'gradients': gradients.reshape(batch_size, n_samples, 3),
+            's_val': 1.0 / inv_s,
+            'mid_z_vals': mid_z_vals,
+            'weights': weights,
+            'cdf': c.reshape(batch_size, n_samples),
+            'gradient_error': gradient_error,
+            'inside_sphere': inside_sphere
+        }
+
+    def render(self, rays_o, rays_d, near, far, pts_ts=0, perturb_overwrite=-1, background_rgb=None, cos_anneal_ratio=0.0, use_gt_sdf=False):
+        batch_size = len(rays_o)
+        sample_dist = 2.0 / self.n_samples  # in a unit sphere # # Assuming the region of interest is a unit sphere
+        z_vals = torch.linspace(0.0, 1.0, self.n_samples) # linspace #
+        z_vals = near + (far - near) * z_vals[None, :]
+
+        z_vals_outside = None
+        if self.n_outside > 0:
+            z_vals_outside = torch.linspace(1e-3, 1.0 - 1.0 / (self.n_outside + 1.0), self.n_outside)
+
+        n_samples = self.n_samples
+        perturb = self.perturb
+
+        if perturb_overwrite >= 0:
+            perturb = perturb_overwrite
+        if perturb > 0:
+            t_rand = (torch.rand([batch_size, 1]) - 0.5)
+            z_vals = z_vals + t_rand * 2.0 / self.n_samples
+
+            if self.n_outside > 0: # z values output # n_outside #
+                mids = .5 * (z_vals_outside[..., 1:] + z_vals_outside[..., :-1])
+                upper = torch.cat([mids, z_vals_outside[..., -1:]], -1)
+                lower = torch.cat([z_vals_outside[..., :1], mids], -1)
+                t_rand = torch.rand([batch_size, z_vals_outside.shape[-1]])
+                z_vals_outside = lower[None, :] + (upper - lower)[None, :] * t_rand
+
+        if self.n_outside > 0:
+            z_vals_outside = far / torch.flip(z_vals_outside, dims=[-1]) + 1.0 / self.n_samples
+
+        background_alpha = None
+        background_sampled_color = None
+
+        # Up sample
+        if self.n_importance > 0:
+            with torch.no_grad():
+                pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]
+                
+                pts = (pts - self.minn_pts) / (self.maxx_pts - self.minn_pts)
+                # sdf = self.sdf_network.sdf(pts.reshape(-1, 3)).reshape(batch_size, self.n_samples)
+                # gt_sdf #
+                
+                # 
+                # pts = ((pts - xyz_min) / (xyz_max - xyz_min)).flip((-1,)) * 2 - 1
+                
+                # pts = pts.flip((-1,)) * 2 - 1
+                pts = pts * 2 - 1                             
+                
+                if self.use_selector:
+                    pts, sdf_selector = self.deform_pts_with_selector(pts=pts, pts_ts=pts_ts)
+                else:
+                    pts = self.deform_pts(pts=pts, pts_ts=pts_ts) # give nthe pts
+                
+                pts_exp = pts.reshape(-1, 3)
+                # minn_pts, _ = torch.min(pts_exp, dim=0)
+                # maxx_pts, _ = torch.max(pts_exp, dim=0) # deformation field (not a rigid one) -> the meshes #
+                # print(f"minn_pts: {minn_pts}, maxx_pts: {maxx_pts}")
+
+                # pts_to_near = pts - near.unsqueeze(1)
+                # maxx_pts = 1.5; minn_pts = -1.5
+                # # maxx_pts = 3; minn_pts = -3
+                # # maxx_pts = 1; minn_pts = -1
+                # pts_exp = (pts_exp - minn_pts) / (maxx_pts - minn_pts)
+                
+                ## render and iamges  ####  
+                if use_gt_sdf:
+                    ### use the GT sdf field ####
+                    # print(f"Using gt sdf :")
+                    sdf = self.gt_sdf(pts_exp.reshape(-1, 3).detach().cpu().numpy())
+                    sdf = torch.from_numpy(sdf).float().cuda()
+                    sdf = sdf.reshape(batch_size, self.n_samples)
+                    ### use the GT sdf field ####
+                else:
+                    # pts_exp: (bsz x nn_s) x 3 -> (sdf_network) -> (bsz x nn_s)
+                    #### use the optimized sdf field ####
+                    
+                    # sdf = self.sdf_network.sdf(pts_exp).reshape(batch_size, self.n_samples)
+                    
+                    if isinstance(self.sdf_network, list):
+                        tot_sdf_values = []
+                        for i_obj, cur_sdf_network in enumerate(self.sdf_network):
+                            cur_sdf_values = cur_sdf_network.sdf(pts_exp).reshape(batch_size, self.n_samples)
+                            tot_sdf_values.append(cur_sdf_values)
+                        tot_sdf_values = torch.stack(tot_sdf_values, dim=-1)
+                        tot_sdf_values, _ = torch.min(tot_sdf_values, dim=-1) # totsdf values #
+                        sdf = tot_sdf_values
+                    else:
+                        # sdf = self.sdf_network.sdf(pts_exp).reshape(batch_size, self.n_samples)
+                        
+                        if self.use_selector:
+                            prev_sdf = self.prev_sdf_network.sdf(pts_exp) # .reshape(batch_size, self.n_samples)
+                            cur_sdf = self.sdf_network.sdf(pts_exp) # .reshape(batch_size, self.n_samples)
+                            sdf = torch.stack([prev_sdf, cur_sdf], dim=1)
+                            sdf = batched_index_select(sdf, indices=sdf_selector.unsqueeze(-1), dim=1).squeeze(1)
+                            sdf = sdf.reshape(batch_size, self.n_samples)
+                        else:
+                            sdf = self.sdf_network.sdf(pts_exp).reshape(batch_size, self.n_samples)
+                    
+                    #### use the optimized sdf field ####
+
+                for i in range(self.up_sample_steps):
+                    new_z_vals = self.up_sample(rays_o,
+                                                rays_d,
+                                                z_vals,
+                                                sdf,
+                                                self.n_importance // self.up_sample_steps,
+                                                64 * 2**i,
+                                                pts_ts=pts_ts)
+                    z_vals, sdf = self.cat_z_vals(rays_o,
+                                                  rays_d,
+                                                  z_vals,
+                                                  new_z_vals,
+                                                  sdf,
+                                                  last=(i + 1 == self.up_sample_steps),
+                                                  pts_ts=pts_ts)
+
+            n_samples = self.n_samples + self.n_importance
+
+        # Background model
+        if self.n_outside > 0:
+            z_vals_feed = torch.cat([z_vals, z_vals_outside], dim=-1)
+            z_vals_feed, _ = torch.sort(z_vals_feed, dim=-1)
+            ret_outside = self.render_core_outside(rays_o, rays_d, z_vals_feed, sample_dist, self.nerf, pts_ts=pts_ts)
+
+            background_sampled_color = ret_outside['sampled_color']
+            background_alpha = ret_outside['alpha']
+
+        # Render core
+        ret_fine = self.render_core(rays_o, # 
+                                    rays_d,
+                                    z_vals,
+                                    sample_dist,
+                                    self.sdf_network,
+                                    self.deviation_network,
+                                    self.color_network,
+                                    background_rgb=background_rgb,
+                                    background_alpha=background_alpha,
+                                    background_sampled_color=background_sampled_color,
+                                    cos_anneal_ratio=cos_anneal_ratio,
+                                    pts_ts=pts_ts)
+
+        color_fine = ret_fine['color']
+        weights = ret_fine['weights']
+        weights_sum = weights.sum(dim=-1, keepdim=True)
+        gradients = ret_fine['gradients']
+        s_val = ret_fine['s_val'].reshape(batch_size, n_samples).mean(dim=-1, keepdim=True)
+
+        return {
+            'color_fine': color_fine,
+            's_val': s_val,
+            'cdf_fine': ret_fine['cdf'],
+            'weight_sum': weights_sum,
+            'weight_max': torch.max(weights, dim=-1, keepdim=True)[0],
+            'gradients': gradients,
+            'weights': weights,
+            'gradient_error': ret_fine['gradient_error'],
+            'inside_sphere': ret_fine['inside_sphere']
+        }
+        
+    # def 
+    def  extract_fields_from_tets_with_selector(self, bound_min, bound_max, resolution, pts_ts ):
+        # load tet via resolution #
+        # scale them via bounds #
+        # extract the geometry #
+        # /home/xueyi/gen/DeepMetaHandles/data/tets/100_compress.npz # strange #
+        device = bound_min.device
+        # if resolution in [64, 70, 80, 90, 100]:
+        #     tet_fn = f"/home/xueyi/gen/DeepMetaHandles/data/tets/{resolution}_compress.npz"
+        # else:
+        tet_fn = f"/home/xueyi/gen/DeepMetaHandles/data/tets/{100}_compress.npz"
+        tets = np.load(tet_fn)
+        verts = torch.from_numpy(tets['vertices']).float().to(device) # verts positions 
+        indices = torch.from_numpy(tets['tets']).long().to(device) # .to(self.device)
+        # split #
+        # verts; verts; #
+        minn_verts, _ = torch.min(verts, dim=0)
+        maxx_verts, _ = torch.max(verts, dim=0) # (3, ) # exporting the  
+        # scale_verts = maxx_verts - minn_verts
+        scale_bounds = bound_max - bound_min # scale bounds #
+        
+        ### scale the vertices ###
+        scaled_verts = (verts - minn_verts.unsqueeze(0)) / (maxx_verts - minn_verts).unsqueeze(0) ### the maxx and minn verts scales ###
+        
+        # scaled_verts = (verts - minn_verts.unsqueeze(0)) / (maxx_verts - minn_verts).unsqueeze(0) ### the maxx and minn verts scales ###
+
+        scaled_verts = scaled_verts * 2. - 1. # init the sdf filed viathe tet mesh vertices and the sdf values ##
+        # scaled_verts = (scaled_verts * scale_bounds.unsqueeze(0)) + bound_min.unsqueeze(0) ## the scaled verts ###
+        
+        # scaled_verts = scaled_verts - scale_bounds.unsqueeze(0) / 2. # 
+        # scaled_verts = scaled_verts -  bound_min.unsqueeze(0) - scale_bounds.unsqueeze(0) / 2.
+        
+        sdf_values = []
+        N = 64
+        query_bundles = N ** 3 ### N^3
+        query_NNs = scaled_verts.size(0) // query_bundles
+        if query_NNs * query_bundles < scaled_verts.size(0):
+            query_NNs += 1
+        for i_query in range(query_NNs):
+            cur_bundle_st = i_query * query_bundles
+            cur_bundle_ed = (i_query + 1) * query_bundles
+            cur_bundle_ed = min(cur_bundle_ed, scaled_verts.size(0))
+            cur_query_pts = scaled_verts[cur_bundle_st: cur_bundle_ed]
+            # if def_func is not None:
+            cur_query_pts, sdf_selector = self.deform_pts_with_selector(pts=cur_query_pts, pts_ts=pts_ts)
+            
+            prev_query_vals = -self.prev_sdf_network.sdf(cur_query_pts)
+            cur_query_vals = -self.sdf_network.sdf(cur_query_pts)
+            cur_query_vals = torch.stack([prev_query_vals, cur_query_vals], dim=1)
+            cur_query_vals = batched_index_select(cur_query_vals, sdf_selector.unsqueeze(-1), dim=1).squeeze(1)
+            
+            # cur_query_vals = query_func(cur_query_pts)
+            sdf_values.append(cur_query_vals)
+        sdf_values = torch.cat(sdf_values, dim=0)
+        # print(f"queryed sdf values: {sdf_values.size()}") #
+        
+        GT_sdf_values = np.load("/home/xueyi/diffsim/DiffHand/assets/hand/100_sdf_values.npy", allow_pickle=True)
+        GT_sdf_values = torch.from_numpy(GT_sdf_values).float().to(device)
+        
+        # intrinsic, tet values, pts values, sdf network #
+        triangle_table, num_triangles_table, base_tet_edges, v_id = create_mt_variable(device)
+        tet_table, num_tets_table = create_tetmesh_variables(device)
+        
+        sdf_values = sdf_values.squeeze(-1) # how the rendering # 
+        
+        # print(f"GT_sdf_values: {GT_sdf_values.size()}, sdf_values: {sdf_values.size()}, scaled_verts: {scaled_verts.size()}")
+        # print(f"scaled_verts: {scaled_verts.size()}, ")
+        # pos_nx3, sdf_n, tet_fx4, triangle_table, num_triangles_table, base_tet_edges, v_id,
+        # return_tet_mesh=False, ori_v=None, num_tets_table=None, tet_table=None):
+        # marching_tets_tetmesh ##
+        verts, faces, tet_verts, tets = marching_tets_tetmesh(scaled_verts, sdf_values, indices, triangle_table, num_triangles_table, base_tet_edges, v_id, return_tet_mesh=True, ori_v=scaled_verts, num_tets_table=num_tets_table, tet_table=tet_table)
+        ### use the GT sdf values for the marching tets ###
+        GT_verts,  GT_faces,  GT_tet_verts,  GT_tets = marching_tets_tetmesh(scaled_verts, GT_sdf_values, indices, triangle_table, num_triangles_table, base_tet_edges, v_id, return_tet_mesh=True, ori_v=scaled_verts, num_tets_table=num_tets_table, tet_table=tet_table)
+        
+        # print(f"After tet marching with verts: {verts.size()}, faces: {faces.size()}")
+        return verts, faces, sdf_values, GT_verts, GT_faces  # verts, faces #
+
+
+    def extract_geometry(self, bound_min, bound_max, resolution, threshold=0.0):
+        return extract_geometry(bound_min, # extract geometry #
+                                bound_max,
+                                resolution=resolution,
+                                threshold=threshold,
+                                # query_func=lambda pts: -self.sdf_network.sdf(pts),
+                                query_func=lambda pts: -self.query_func_sdf(pts)
+                                )
+        
+    # if self.deform_pts_with_selector:
+    #         pts = self.deform_pts_with_selector(pts=pts, pts_ts=pts_ts)
+    def extract_geometry_tets(self, bound_min, bound_max, resolution, pts_ts=0, threshold=0.0, wdef=False):
+        if wdef:
+            return extract_geometry_tets(bound_min, # extract geometry #
+                                bound_max,
+                                resolution=resolution,
+                                threshold=threshold,
+                                query_func=lambda pts: -self.query_func_sdf(pts), # lambda pts: -self.sdf_network.sdf(pts),
+                                def_func=lambda pts: self.deform_pts(pts, pts_ts=pts_ts) if not self.use_selector else self.deform_pts_with_selector(pts=pts, pts_ts=pts_ts))
+        else:
+            return extract_geometry_tets(bound_min, # extract geometry #
+                                    bound_max,
+                                    resolution=resolution,
+                                    threshold=threshold,
+                                    # query_func=lambda pts: -self.sdf_network.sdf(pts)
+                                     query_func=lambda pts: -self.query_func_sdf(pts), # lambda pts: -self.sdf_network.sdf(pts),
+                                    )
+    
+    
+    def extract_geometry_tets_kinematic(self, bound_min, bound_max, resolution, pts_ts=0, threshold=0.0, wdef=False):
+        if wdef:
+            return extract_geometry_tets(bound_min, # extract geometry #
+                                bound_max,
+                                resolution=resolution,
+                                threshold=threshold,
+                                query_func=lambda pts: -self.query_func_sdf(pts), # lambda pts: -self.sdf_network.sdf(pts),
+                                def_func=lambda pts: self.deform_pts_kinematic(pts, pts_ts=pts_ts))
+        else:
+            return extract_geometry_tets(bound_min, # extract geometry #
+                                    bound_max,
+                                    resolution=resolution,
+                                    threshold=threshold,
+                                    # query_func=lambda pts: -self.sdf_network.sdf(pts)
+                                     query_func=lambda pts: -self.query_func_sdf(pts), # lambda pts: -self.sdf_network.sdf(pts),
+            )
+            
+            
+    def extract_geometry_tets_active(self, bound_min, bound_max, resolution, pts_ts=0, threshold=0.0, wdef=False):
+        if wdef:
+            return extract_geometry_tets(bound_min, # extract geometry #
+                                bound_max,
+                                resolution=resolution,
+                                threshold=threshold,
+                                query_func=lambda pts: -self.query_func_active(pts), # lambda pts: -self.sdf_network.sdf(pts),
+                                def_func=lambda pts: self.deform_pts_kinematic_active(pts, pts_ts=pts_ts))
+        else:
+            return extract_geometry_tets(bound_min, # extract geometry #
+                                    bound_max,
+                                    resolution=resolution,
+                                    threshold=threshold,
+                                    # query_func=lambda pts: -self.sdf_network.sdf(pts)
+                                     query_func=lambda pts: -self.query_func_sdf(pts), # lambda pts: -self.sdf_network.sdf(pts),
+            )
+            
+    def extract_geometry_tets_passive(self, bound_min, bound_max, resolution, pts_ts=0, threshold=0.0, wdef=False):
+        if wdef:
+            return extract_geometry_tets(bound_min, # extract geometry #
+                                bound_max,
+                                resolution=resolution,
+                                threshold=threshold,
+                                query_func=lambda pts: -self.query_func_sdf_passive(pts), # lambda pts: -self.sdf_network.sdf(pts),
+                                def_func=lambda pts: self.deform_pts_passive(pts, pts_ts=pts_ts))
+        else:
+            return extract_geometry_tets(bound_min, # extract geometry #
+                                    bound_max,
+                                    resolution=resolution,
+                                    threshold=threshold,
+                                    # query_func=lambda pts: -self.sdf_network.sdf(pts)
+                                     query_func=lambda pts: -self.query_func_sdf(pts), # lambda pts: -self.sdf_network.sdf(pts),
+                                    )

pre-requirements.txt

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

requirements.txt

@@ -0,0 +1,27 @@
+-f https://download.pytorch.org/whl/cpu/torch_stable.html
+-f https://data.pyg.org/whl/torch-2.2.0%2Bcpu.html
+# pip==20.2.4
+torch==2.2.0
+# torchvision==0.13.1
+# torchaudio==0.12.1
+scipy
+trimesh
+icecream
+tqdm
+pyhocon
+open3d
+tensorboard
+
+# blobfile==2.0.1
+# manopth @ git+https://github.com/hassony2/manopth.git
+# numpy==1.23.1
+# psutil==5.9.2
+# scikit-learn
+# scipy==1.9.3
+# tensorboard
+# tensorboardx
+# tqdm
+# trimesh
+# clip
+# chumpy
+# opencv-python

scripts_demo/train_grab_pointset_points_dyn_s1.sh

@@ -0,0 +1,28 @@
+
+export PYTHONPATH=.
+
+
+
+
+export data_case=hand_test_routine_2_light_color_wtime_active_passive
+
+
+
+export trainer=exp_runner_stage_1.py
+
+
+export mode="train_point_set"
+
+export conf=dyn_grab_pointset_points_dyn_s1.conf
+
+export conf_root="./confs_new"
+
+
+export data_path="./data/102_grab_all_data.npy"
+# bash scripts_new/train_grab_pointset_points_dyn_s1.sh
+
+export cuda_ids="0"
+
+
+CUDA_VISIBLE_DEVICES=${cuda_ids} python ${trainer} --mode ${mode} --conf ${conf_root}/${conf} --case ${data_case} --data_path=${data_path}
+

scripts_new/train_grab_mano.sh

@@ -0,0 +1,26 @@
+
+export PYTHONPATH=.
+
+export conf=wmask_refine_passive_rigidtrans_forward.conf
+export data_case=hand_test_routine_2_light_color_wtime_active_passive
+
+
+
+export trainer=exp_runner_stage_1.py
+
+export mode="train_dyn_mano_model"
+
+
+export conf=dyn_grab_pointset_mano.conf
+
+export conf_root="./confs_new"
+
+
+
+# bash scripts_new/train_grab_mano.sh
+
+export cuda_ids="0"
+# 
+
+PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python CUDA_VISIBLE_DEVICES=${cuda_ids} python ${trainer} --mode ${mode} --conf ${conf_root}/${conf} --case ${data_case}
+

scripts_new/train_grab_mano_wreact.sh

@@ -0,0 +1,26 @@
+
+export PYTHONPATH=.
+
+
+
+export data_case=hand_test_routine_2_light_color_wtime_active_passive
+
+
+
+export trainer=exp_runner_stage_1.py
+
+export mode="train_dyn_mano_model_wreact" ## wreact ## wreact ##
+
+export conf=dyn_grab_pointset_mano_dyn.conf
+
+
+export conf_root="./confs_new"
+
+
+
+export cuda_ids="0"
+
+
+
+CUDA_VISIBLE_DEVICES=${cuda_ids} python ${trainer} --mode ${mode} --conf ${conf_root}/${conf} --case ${data_case}
+

scripts_new/train_grab_mano_wreact_optacts.sh

@@ -0,0 +1,26 @@
+
+export PYTHONPATH=.
+
+
+
+export data_case=hand_test_routine_2_light_color_wtime_active_passive
+
+
+
+export trainer=exp_runner_stage_1.py
+
+export mode="train_dyn_mano_model_wreact" ## wreact ## wreact ##
+
+export conf=dyn_grab_pointset_mano_dyn.conf
+
+
+export conf_root="./confs_new"
+
+
+
+export cuda_ids="0"
+
+
+
+CUDA_VISIBLE_DEVICES=${cuda_ids} python ${trainer} --mode ${mode} --conf ${conf_root}/${conf} --case ${data_case}
+

scripts_new/train_grab_pointset.sh

@@ -0,0 +1,99 @@
+
+export PYTHONPATH=.
+
+export cuda_ids="4"
+
+# export cuda_ids="5"
+
+
+export trainer=exp_runner_arti_forward.py
+export conf=wmask_refine_passive_rigidtrans_forward.conf
+export data_case=hand_test_routine_2_light_color_wtime_active_passive
+
+
+# export trainer=exp_runner_arti_multi_objs_compositional.py
+export trainer=exp_runner_arti_multi_objs_compositional_ks.py
+# export trainer=exp_runner_arti_multi_objs_dyn.py
+export trainer=exp_runner_arti_multi_objs_pointset.py
+# /home/xueyi/diffsim/NeuS/confs/wmask_refine_passive_compositional.conf
+export conf=wmask_refine_passive_compositional.conf
+export conf=dyn_arctic_ks.conf
+
+export conf=dyn_arctic_ks_robohand.conf
+# /data/xueyi/diffsim/NeuS/confs/dyn_arctic_ks_robohand_from_mano_model_rules.conf
+export conf=dyn_arctic_ks_robohand_from_mano_model_rules.conf
+export conf=dyn_arctic_robohand_from_mano_model_rules.conf
+export conf=dyn_arctic_robohand_from_mano_model_rules_actions.conf
+export conf=dyn_arctic_robohand_from_mano_model_rules_actions_f2.conf
+export conf=dyn_arctic_robohand_from_mano_model_rules_actions_f2_diffhand.conf
+export conf=dyn_arctic_robohand_from_mano_model_rules_actions_f2_diffhand_v2.conf
+export conf=dyn_arctic_robohand_from_mano_model_rules_actions_f2_diffhand_v4.conf
+
+# /home/xueyi/diffsim/NeuS/confs/dyn_arctic_robohand_from_mano_model_train_mano_dyn_model_states.conf
+export conf=dyn_arctic_robohand_from_mano_model_train_mano_dyn_model_states.conf
+# /home/xueyi/diffsim/NeuS/confs/dyn_arctic_robohand_from_mano_model_train_mano_dyn_model_states.conf
+# export conf=dyn_arctic_robohand_from_mano_model_train_mano_dyn_model_states.conf
+
+export conf=dyn_arctic_robohand_from_mano_model_train_mano_dyn_model_states.conf
+# export conf=dyn_arctic_ks_robohand_from_mano_model_rules_arti.conf
+# a very stiff system for # 
+export conf=dyn_grab_pointset_mano.conf
+
+export mode="train_from_model_rules"
+export mode="train_from_model_rules"
+
+export mode="train_sdf_from_model_rules"
+export mode="train_actions_from_model_rules"
+# export mode="train_actions_from_sim_rules"
+
+
+
+export mode="train_def"
+# export mode="train_actions_from_model_rules"
+export mode="train_mano_actions_from_model_rules"
+
+# virtual force # #  ## virtual forces ##
+# /data/xueyi/diffsim/NeuS/confs/dyn_arctic_ks_robohand_from_mano_model_rules.conf ##
+export mode="train_actions_from_model_rules"
+
+export mode="train_actions_from_mano_model_rules"
+
+export mode="train_real_robot_actions_from_mano_model_rules"
+export mode="train_real_robot_actions_from_mano_model_rules_diffhand"
+
+export mode="train_real_robot_actions_from_mano_model_rules_diffhand_fortest"
+export mode="train_real_robot_actions_from_mano_model_rules_manohand_fortest"
+
+
+export mode="train_real_robot_actions_from_mano_model_rules_diffhand_fortest"
+
+
+
+export mode="train_real_robot_actions_from_mano_model_rules_manohand_fortest_states"
+
+###### diffsim ######
+# /home/xueyi/diffsim/NeuS/confs/dyn_arctic_robohand_from_mano_model_rules_actions_f2_diffhand_v3.conf 
+
+export mode="train_real_robot_actions_from_mano_model_rules_v5_manohand_fortest_states_res_world"
+export mode="train_real_robot_actions_from_mano_model_rules_v5_manohand_fortest_states_res_rl"
+
+export mode="train_dyn_mano_model_states"
+
+
+## train dyn mano model states wreact ##
+export mode="train_dyn_mano_model_states_wreact"
+export conf=dyn_grab_pointset_mano_dyn.conf
+## train dyn mano model states wreact ##
+
+
+export conf_root="./confs_new"
+
+
+#
+# bash scripts_new/train_grab_pointset_dyn.sh
+
+export cuda_ids="2"
+# 
+
+PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python CUDA_VISIBLE_DEVICES=${cuda_ids} python ${trainer} --mode ${mode} --conf ${conf_root}/${conf} --case ${data_case}
+

scripts_new/train_grab_pointset_points_dyn.sh

@@ -0,0 +1,27 @@
+
+export PYTHONPATH=.
+
+
+
+
+export data_case=hand_test_routine_2_light_color_wtime_active_passive
+
+
+
+export trainer=exp_runner_stage_1.py
+
+
+export mode="train_point_set"
+
+export conf=dyn_grab_pointset_points_dyn.conf
+
+export conf_root="./confs_new"
+
+
+# bash scripts_new/train_grab_pointset_points_dyn.sh
+
+export cuda_ids="0"
+
+
+PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python CUDA_VISIBLE_DEVICES=${cuda_ids} python ${trainer} --mode ${mode} --conf ${conf_root}/${conf} --case ${data_case}
+

scripts_new/train_grab_pointset_points_dyn_retar.sh

@@ -0,0 +1,36 @@
+
+export PYTHONPATH=.
+
+
+
+
+export data_case=hand_test_routine_2_light_color_wtime_active_passive
+
+
+
+export trainer=exp_runner_stage_1.py
+
+
+
+### stage 1 -> tracking MANO expanded set using Shadow's object mesh points ###
+export mode="train_point_set_retar"
+export conf=dyn_grab_pointset_points_dyn_retar.conf
+
+# ### stage 2 -> tracking MANO expanded set using Shadow's expanded points ###
+# export mode="train_expanded_set_motions_retar_pts"
+# export conf=dyn_grab_pointset_points_dyn_retar_pts.conf
+
+
+
+
+export conf_root="./confs_new"
+
+
+# bash scripts_new/train_grab_pointset_points_dyn_retar.sh
+
+export cuda_ids="0"
+
+
+
+PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python CUDA_VISIBLE_DEVICES=${cuda_ids} python ${trainer} --mode ${mode} --conf ${conf_root}/${conf} --case ${data_case}
+

scripts_new/train_grab_pointset_points_dyn_retar_pts.sh

@@ -0,0 +1,36 @@
+
+export PYTHONPATH=.
+
+
+
+
+export data_case=hand_test_routine_2_light_color_wtime_active_passive
+
+
+
+export trainer=exp_runner_stage_1.py
+
+
+
+# ### stage 1 -> tracking MANO expanded set using Shadow's object mesh points ###
+# export mode="train_expanded_set_motions_retar"
+# export conf=dyn_grab_pointset_points_dyn_retar.conf
+
+### stage 2 -> tracking MANO expanded set using Shadow's expanded points ###
+export mode="train_point_set_retar_pts"
+export conf=dyn_grab_pointset_points_dyn_retar_pts.conf
+
+
+
+
+export conf_root="./confs_new"
+
+
+# bash scripts_new/train_grab_pointset_points_dyn_retar.sh
+
+export cuda_ids="0"
+
+
+
+PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python CUDA_VISIBLE_DEVICES=${cuda_ids} python ${trainer} --mode ${mode} --conf ${conf_root}/${conf} --case ${data_case}
+

scripts_new/train_grab_pointset_points_dyn_s1.sh

@@ -0,0 +1,27 @@
+
+export PYTHONPATH=.
+
+
+
+
+export data_case=hand_test_routine_2_light_color_wtime_active_passive
+
+
+
+export trainer=exp_runner_stage_1.py
+
+
+export mode="train_point_set"
+
+export conf=dyn_grab_pointset_points_dyn_s1.conf
+
+export conf_root="./confs_new"
+
+
+# bash scripts_new/train_grab_pointset_points_dyn_s1.sh
+
+export cuda_ids="0"
+
+
+PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python CUDA_VISIBLE_DEVICES=${cuda_ids} python ${trainer} --mode ${mode} --conf ${conf_root}/${conf} --case ${data_case}
+

scripts_new/train_grab_pointset_points_dyn_s2.sh

@@ -0,0 +1,27 @@
+
+export PYTHONPATH=.
+
+
+
+
+export data_case=hand_test_routine_2_light_color_wtime_active_passive
+
+
+
+export trainer=exp_runner_stage_1.py
+
+
+export mode="train_point_set"
+
+export conf=dyn_grab_pointset_points_dyn_s2.conf
+
+export conf_root="./confs_new"
+
+
+# bash scripts_new/train_grab_pointset_points_dyn_s2.sh
+
+export cuda_ids="0"
+
+
+PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python CUDA_VISIBLE_DEVICES=${cuda_ids} python ${trainer} --mode ${mode} --conf ${conf_root}/${conf} --case ${data_case}
+

scripts_new/train_grab_pointset_points_dyn_s3.sh

@@ -0,0 +1,27 @@
+
+export PYTHONPATH=.
+
+
+
+
+export data_case=hand_test_routine_2_light_color_wtime_active_passive
+
+
+
+export trainer=exp_runner_stage_1.py
+
+
+export mode="train_point_set"
+
+export conf=dyn_grab_pointset_points_dyn_s3.conf
+
+export conf_root="./confs_new"
+
+
+# bash scripts_new/train_grab_pointset_points_dyn_s3.sh
+
+export cuda_ids="0"
+
+
+PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python CUDA_VISIBLE_DEVICES=${cuda_ids} python ${trainer} --mode ${mode} --conf ${conf_root}/${conf} --case ${data_case}
+

scripts_new/train_grab_pointset_points_dyn_s4.sh

@@ -0,0 +1,27 @@
+
+export PYTHONPATH=.
+
+
+
+
+export data_case=hand_test_routine_2_light_color_wtime_active_passive
+
+
+
+export trainer=exp_runner_stage_1.py
+
+
+export mode="train_point_set"
+
+export conf=dyn_grab_pointset_points_dyn_s4.conf
+
+export conf_root="./confs_new"
+
+
+# bash scripts_new/train_grab_pointset_points_dyn_s4.sh
+
+export cuda_ids="0"
+
+
+PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python CUDA_VISIBLE_DEVICES=${cuda_ids} python ${trainer} --mode ${mode} --conf ${conf_root}/${conf} --case ${data_case}
+

scripts_new/train_grab_shadow_multistages.sh

@@ -0,0 +1,102 @@
+
+export PYTHONPATH=.
+
+export cuda_ids="3"
+
+
+
+export trainer=exp_runner_arti_forward.py
+export conf=wmask_refine_passive_rigidtrans_forward.conf
+export data_case=hand_test_routine_2_light_color_wtime_active_passive
+
+
+# export trainer=exp_runner_arti_multi_objs_compositional.py
+export trainer=exp_runner_arti_multi_objs_compositional_ks.py
+export trainer=exp_runner_arti_multi_objs_arti_dyn.py
+export conf=dyn_arctic_robohand_from_mano_model_rules_actions_f2_diffhand.conf
+export conf=dyn_arctic_robohand_from_mano_model_rules_actions_f2_diffhand_v2.conf
+export conf=dyn_arctic_robohand_from_mano_model_rules_actions_f2_diffhand_v4.conf
+
+# /home/xueyi/diffsim/NeuS/confs/dyn_arctic_robohand_from_mano_model_train_mano_dyn_model_states.conf
+export conf=dyn_arctic_robohand_from_mano_model_train_mano_dyn_model_states.conf
+
+# /home/xueyi/diffsim/NeuS/confs/dyn_grab_mano_model_states.conf
+export conf=dyn_grab_mano_model_states.conf
+
+export conf=dyn_grab_shadow_model_states.conf
+
+
+
+
+export mode="train_def"
+# export mode="train_actions_from_model_rules"
+export mode="train_mano_actions_from_model_rules"
+
+# virtual force #
+# /data/xueyi/diffsim/NeuS/confs/dyn_arctic_ks_robohand_from_mano_model_rules.conf ##
+
+export mode="train_real_robot_actions_from_mano_model_rules_diffhand_fortest"
+
+
+
+export mode="train_real_robot_actions_from_mano_model_rules_manohand_fortest_states"
+
+
+### using the diffsim ###
+
+###### diffsim ######
+# /home/xueyi/diffsim/NeuS/confs/dyn_arctic_robohand_from_mano_model_rules_actions_f2_diffhand_v3.conf 
+
+
+export mode="train_real_robot_actions_from_mano_model_rules_v5_manohand_fortest_states_res_world"
+export mode="train_real_robot_actions_from_mano_model_rules_v5_manohand_fortest_states_res_rl"
+
+export mode="train_dyn_mano_model_states"
+
+# 
+export mode="train_real_robot_actions_from_mano_model_rules_v5_shadowhand_fortest_states_grab"
+## dyn grab shadow 
+
+
+### optimize for the manipulatable hand actions ###
+export mode="train_real_robot_actions_from_mano_model_rules_v5_shadowhand_fortest_states_grab_redmax_acts"
+
+
+
+# ## acts ##
+# export conf=dyn_grab_shadow_model_states_224.conf
+# export conf=dyn_grab_shadow_model_states_89.conf
+# export conf=dyn_grab_shadow_model_states_102.conf
+# export conf=dyn_grab_shadow_model_states_7.conf
+# export conf=dyn_grab_shadow_model_states_47.conf
+# export conf=dyn_grab_shadow_model_states_67.conf
+# export conf=dyn_grab_shadow_model_states_76.conf
+# export conf=dyn_grab_shadow_model_states_85.conf
+# # export conf=dyn_grab_shadow_model_states_91.conf
+# # export conf=dyn_grab_shadow_model_states_167.conf
+# export conf=dyn_grab_shadow_model_states_107.conf
+# export conf=dyn_grab_shadow_model_states_306.conf
+# export conf=dyn_grab_shadow_model_states_313.conf
+# export conf=dyn_grab_shadow_model_states_322.conf
+
+
+# /home/xueyi/diffsim/NeuS/confs_new/dyn_grab_arti_shadow_multi_stages.conf
+export conf=dyn_grab_arti_shadow_multi_stages.conf
+# export conf=dyn_grab_shadow_model_states_398.conf
+# export conf=dyn_grab_shadow_model_states_363.conf
+# export conf=dyn_grab_shadow_model_states_358.conf
+
+# bash scripts_new/train_grab_shadow_multistages.sh
+
+
+
+
+
+export conf_root="./confs_new"
+
+
+export cuda_ids="4"
+
+
+PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python CUDA_VISIBLE_DEVICES=${cuda_ids} python ${trainer} --mode ${mode} --conf ${conf_root}/${conf} --case ${data_case}
+

scripts_new/train_grab_shadow_singlestage.sh

@@ -0,0 +1,101 @@
+
+export PYTHONPATH=.
+
+export cuda_ids="3"
+
+
+
+export trainer=exp_runner_arti_forward.py
+export conf=wmask_refine_passive_rigidtrans_forward.conf
+export data_case=hand_test_routine_2_light_color_wtime_active_passive
+
+
+# export trainer=exp_runner_arti_multi_objs_compositional.py
+export trainer=exp_runner_arti_multi_objs_compositional_ks.py
+export trainer=exp_runner_arti_multi_objs_arti_dyn.py
+export conf=dyn_arctic_robohand_from_mano_model_rules_actions_f2_diffhand.conf
+export conf=dyn_arctic_robohand_from_mano_model_rules_actions_f2_diffhand_v2.conf
+export conf=dyn_arctic_robohand_from_mano_model_rules_actions_f2_diffhand_v4.conf
+
+# /home/xueyi/diffsim/NeuS/confs/dyn_arctic_robohand_from_mano_model_train_mano_dyn_model_states.conf
+export conf=dyn_arctic_robohand_from_mano_model_train_mano_dyn_model_states.conf
+
+# /home/xueyi/diffsim/NeuS/confs/dyn_grab_mano_model_states.conf
+export conf=dyn_grab_mano_model_states.conf
+
+export conf=dyn_grab_shadow_model_states.conf
+
+
+
+
+export mode="train_def"
+# export mode="train_actions_from_model_rules"
+export mode="train_mano_actions_from_model_rules"
+
+# virtual force #
+# /data/xueyi/diffsim/NeuS/confs/dyn_arctic_ks_robohand_from_mano_model_rules.conf ##
+
+export mode="train_real_robot_actions_from_mano_model_rules_diffhand_fortest"
+
+
+
+export mode="train_real_robot_actions_from_mano_model_rules_manohand_fortest_states"
+
+
+###### diffsim ######
+# /home/xueyi/diffsim/NeuS/confs/dyn_arctic_robohand_from_mano_model_rules_actions_f2_diffhand_v3.conf 
+
+
+export mode="train_real_robot_actions_from_mano_model_rules_v5_manohand_fortest_states_res_world"
+export mode="train_real_robot_actions_from_mano_model_rules_v5_manohand_fortest_states_res_rl"
+
+export mode="train_dyn_mano_model_states"
+
+# 
+export mode="train_real_robot_actions_from_mano_model_rules_v5_shadowhand_fortest_states_grab"
+## dyn grab shadow 
+
+### optimze for the redmax hand actions from joint states ###
+export mode="train_real_robot_actions_from_mano_model_rules_shadowhand"
+
+### optimize for the manipulatable hand actions ###
+export mode="train_real_robot_actions_from_mano_model_rules_v5_shadowhand_fortest_states_grab_redmax_acts"
+
+
+
+# ## acts ##
+# export conf=dyn_grab_shadow_model_states_224.conf
+# export conf=dyn_grab_shadow_model_states_89.conf
+# export conf=dyn_grab_shadow_model_states_102.conf
+# export conf=dyn_grab_shadow_model_states_7.conf
+# export conf=dyn_grab_shadow_model_states_47.conf
+# export conf=dyn_grab_shadow_model_states_67.conf
+# export conf=dyn_grab_shadow_model_states_76.conf
+# export conf=dyn_grab_shadow_model_states_85.conf
+# # export conf=dyn_grab_shadow_model_states_91.conf
+# # export conf=dyn_grab_shadow_model_states_167.conf
+# export conf=dyn_grab_shadow_model_states_107.conf
+# export conf=dyn_grab_shadow_model_states_306.conf
+# export conf=dyn_grab_shadow_model_states_313.conf
+# export conf=dyn_grab_shadow_model_states_322.conf
+
+# /home/xueyi/diffsim/NeuS/confs_new/dyn_grab_arti_shadow_multi_stages.conf
+# /home/xueyi/diffsim/NeuS/confs_new/dyn_grab_arti_shadow_single_stage.conf
+export conf=dyn_grab_arti_shadow_single_stage.conf
+# export conf=dyn_grab_shadow_model_states_398.conf
+# export conf=dyn_grab_shadow_model_states_363.conf
+# export conf=dyn_grab_shadow_model_states_358.conf
+
+# bash scripts_new/train_grab_shadow_singlestage.sh
+
+
+
+
+export conf_root="./confs_new"
+
+
+export cuda_ids="3"
+
+
+PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python CUDA_VISIBLE_DEVICES=${cuda_ids} python ${trainer} --mode ${mode} --conf ${conf_root}/${conf} --case ${data_case}
+