From 429b7a2d2f7b366b4cd83828cf64fd9acca0b8ec Mon Sep 17 00:00:00 2001 From: LJones-alt Date: Wed, 25 Feb 2026 13:55:31 +0000 Subject: [PATCH 1/2] minor changs --- docs/find_best_checkpoints.py | 2 +- .../vial_insert/20260216134807/config.json | 236 ++++++++++++++++++ docs/run_experiments.sh | 58 ++--- .../robomimic/backup_controller_handler.py | 5 +- .../robomimic/backup_controller_insert.py | 163 ++++++++---- .../robomimic/play_controller.py | 2 +- .../robomimic/play_ensemble_v05.py | 47 +++- .../robomimic/play_with_backup.py | 2 +- .../robomimic/switching_logic_algorithm.tex | 62 +++++ .../robomimic/train_ensemble.py | 6 +- .../isaaclab/envs/mdp/observations.py | 1 + .../isaaclab_mimic/datagen/generation.py | 6 +- .../agents/robomimic/bc_rnn_low_dim.json | 4 +- .../config/franka/dev_ik_rel_vial_insert.py | 2 +- .../cube_lift/mdp/observations.py | 10 +- .../cube_lift/mdp/terminations.py | 9 +- 16 files changed, 507 insertions(+), 108 deletions(-) create mode 100644 docs/insert/Dev-IK-Rel-Insert-v0_old/model5/vial_insert/20260216134807/config.json create mode 100644 scripts/imitation_learning/robomimic/switching_logic_algorithm.tex diff --git a/docs/find_best_checkpoints.py b/docs/find_best_checkpoints.py index 13b927f7e306..4c35e1afe447 100644 --- a/docs/find_best_checkpoints.py +++ b/docs/find_best_checkpoints.py @@ -72,7 +72,7 @@ def main(): """Main function to process all model directories.""" # Define paths - handle both host and Docker container paths script_dir = Path(__file__).parent.resolve() - base_dir = script_dir / "insert/Dev-IK-Rel-Insert-v0" + base_dir = script_dir / "/workspace/isaaclab/robomimic/Dev-IK-Rel-Insert-v0" output_dir = base_dir / "best_models" # Create output directory (including parent directories if needed) diff --git a/docs/insert/Dev-IK-Rel-Insert-v0_old/model5/vial_insert/20260216134807/config.json b/docs/insert/Dev-IK-Rel-Insert-v0_old/model5/vial_insert/20260216134807/config.json new file mode 100644 index 000000000000..c1f0e6cea50b --- /dev/null +++ b/docs/insert/Dev-IK-Rel-Insert-v0_old/model5/vial_insert/20260216134807/config.json @@ -0,0 +1,236 @@ +{ + "algo_name": "bc", + "experiment": { + "name": "vial_insert", + "validate": true, + "logging": { + "terminal_output_to_txt": true, + "log_tb": true, + "log_wandb": true, + "wandb_proj_name": "vial_insert" + }, + "save": { + "enabled": true, + "every_n_seconds": null, + "every_n_epochs": null, + "epochs": [], + "on_best_validation": true, + "on_best_rollout_return": true, + "on_best_rollout_success_rate": true + }, + "epoch_every_n_steps": 100, + "validation_epoch_every_n_steps": 10, + "env": "Dev-IK-Rel-Insert-v0", + "additional_envs": null, + "render": false, + "render_video": true, + "keep_all_videos": false, + "video_skip": 5, + "rollout": { + "enabled": false, + "n": 50, + "horizon": 400, + "rate": 50, + "warmstart": 0, + "terminate_on_success": true + }, + "env_meta_update_dict": {}, + "ckpt_path": null + }, + "train": { + "data": [ + { + "path": "docs/insert/glassware_lift_demos_generated_split.hdf5" + } + ], + "output_dir": "/workspace/isaaclab/robomimic/Dev-IK-Rel-Insert-v0/model14/", + "normalize_weights_by_ds_size": false, + "num_data_workers": 4, + "hdf5_cache_mode": "low_dim", + "hdf5_use_swmr": true, + "hdf5_load_next_obs": false, + "hdf5_normalize_obs": false, + "hdf5_filter_key": "train", + "hdf5_validation_filter_key": "valid", + "seq_length": 16, + "pad_seq_length": true, + "frame_stack": 1, + "pad_frame_stack": true, + "dataset_keys": [ + "actions", + "rewards", + "dones" + ], + "action_keys": [ + "actions" + ], + "action_config": { + "actions": { + "normalization": null + } + }, + "goal_mode": null, + "cuda": true, + "batch_size": 100, + "num_epochs": 1000, + "seed": 101, + "max_grad_norm": null + }, + "algo": { + "optim_params": { + "policy": { + "optimizer_type": "adam", + "learning_rate": { + "initial": 0.0001, + "decay_factor": 0.01, + "epoch_schedule": [], + "scheduler_type": "multistep" + }, + "regularization": { + "L2": 0.0 + } + } + }, + "loss": { + "l2_weight": 1.0, + "l1_weight": 0.0, + "cos_weight": 0.0 + }, + "actor_layer_dims": [ + 1024, + 1024 + ], + "gaussian": { + "enabled": false, + "fixed_std": false, + "init_std": 0.1, + "min_std": 0.01, + "std_activation": "softplus", + "low_noise_eval": true + }, + "gmm": { + "enabled": true, + "num_modes": 10, + "min_std": 0.0001, + "std_activation": "softplus", + "low_noise_eval": true + }, + "vae": { + "enabled": false, + "latent_dim": 14, + "latent_clip": null, + "kl_weight": 1.0, + "decoder": { + "is_conditioned": true, + "reconstruction_sum_across_elements": false + }, + "prior": { + "learn": false, + "is_conditioned": false, + "use_gmm": false, + "gmm_num_modes": 10, + "gmm_learn_weights": false, + "use_categorical": false, + "categorical_dim": 10, + "categorical_gumbel_softmax_hard": false, + "categorical_init_temp": 1.0, + "categorical_temp_anneal_step": 0.001, + "categorical_min_temp": 0.3 + }, + "encoder_layer_dims": [ + 300, + 400 + ], + "decoder_layer_dims": [ + 300, + 400 + ], + "prior_layer_dims": [ + 300, + 400 + ] + }, + "rnn": { + "enabled": true, + "horizon": 100, + "hidden_dim": 400, + "rnn_type": "LSTM", + "num_layers": 2, + "open_loop": false, + "kwargs": { + "bidirectional": false + } + }, + "transformer": { + "enabled": false, + "context_length": 10, + "embed_dim": 512, + "num_layers": 6, + "num_heads": 8, + "emb_dropout": 0.1, + "attn_dropout": 0.1, + "block_output_dropout": 0.1, + "sinusoidal_embedding": false, + "activation": "gelu", + "supervise_all_steps": false, + "nn_parameter_for_timesteps": true, + "pred_future_acs": false + } + }, + "observation": { + "modalities": { + "obs": { + "low_dim": [ + "eef_pos", + "eef_quat", + "gripper_pos", + "object_position", + "target_object_position", + "joint_pos", + "joint_vel", + "actions" + ], + "rgb": [], + "depth": [], + "scan": [] + }, + "goal": { + "low_dim": [], + "rgb": [], + "depth": [], + "scan": [] + } + }, + "encoder": { + "low_dim": { + "core_class": null, + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "rgb": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "depth": { + "core_class": "VisualCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + }, + "scan": { + "core_class": "ScanCore", + "core_kwargs": {}, + "obs_randomizer_class": null, + "obs_randomizer_kwargs": {} + } + } + }, + "meta": { + "hp_base_config_file": null, + "hp_keys": [], + "hp_values": [] + } +} \ No newline at end of file diff --git a/docs/run_experiments.sh b/docs/run_experiments.sh index 9a526d7a878f..03501729f73b 100644 --- a/docs/run_experiments.sh +++ b/docs/run_experiments.sh @@ -2,29 +2,15 @@ ## script to run experiments in the background -task="Dev-IK-Rel-v1" -horizon=500 +task="Dev-IK-Rel-Insert-v0" +horizon=1000 num_rollouts=100 run_file="scripts/imitation_learning/robomimic/play_ensemble_v05.py" -exp_type="lift_none" +exp_type="insert_low_calib" ##### this now runs the experiments -# ## ensemble 15 -# ensemble_size=15 -# seed=101 -# exp_name="${task}_ensemble_${ensemble_size}_seed_${seed}_${exp_type}" -# ./isaaclab.sh -p $run_file --task $task --horizon $horizon --num_rollouts $num_rollouts --ensemble_size $ensemble_size --seed $seed --exp_name $exp_name --headless - -# seed=107 -# exp_name="${task}_ensemble_${ensemble_size}_seed_${seed}_${exp_type}" -# ./isaaclab.sh -p $run_file --task $task --horizon $horizon --num_rollouts $num_rollouts --ensemble_size $ensemble_size --seed $seed --exp_name $exp_name --headless - -# seed=115 -# exp_name="${task}_ensemble_${ensemble_size}_seed_${seed}_${exp_type}" -# ./isaaclab.sh -p $run_file --task $task --horizon $horizon --num_rollouts $num_rollouts --ensemble_size $ensemble_size --seed $seed --exp_name $exp_name --headless - -## ensemble 10 +# # ## ensemble 15 ensemble_size=10 seed=101 exp_name="${task}_ensemble_${ensemble_size}_seed_${seed}_${exp_type}" @@ -38,16 +24,30 @@ seed=115 exp_name="${task}_ensemble_${ensemble_size}_seed_${seed}_${exp_type}" ./isaaclab.sh -p $run_file --task $task --horizon $horizon --num_rollouts $num_rollouts --ensemble_size $ensemble_size --seed $seed --exp_name $exp_name --headless -## ensemble 5 -ensemble_size=5 -seed=101 -exp_name="${task}_ensemble_${ensemble_size}_seed_${seed}_${exp_type}" -./isaaclab.sh -p $run_file --task $task --horizon $horizon --num_rollouts $num_rollouts --ensemble_size $ensemble_size --seed $seed --exp_name $exp_name --headless +# # ## ensemble 10 +# ensemble_size=10 +# seed=101 +# exp_name="${task}_ensemble_${ensemble_size}_seed_${seed}_${exp_type}" +# ./isaaclab.sh -p $run_file --task $task --horizon $horizon --num_rollouts $num_rollouts --ensemble_size $ensemble_size --seed $seed --exp_name $exp_name --headless -seed=107 -exp_name="${task}_ensemble_${ensemble_size}_seed_${seed}_${exp_type}" -./isaaclab.sh -p $run_file --task $task --horizon $horizon --num_rollouts $num_rollouts --ensemble_size $ensemble_size --seed $seed --exp_name $exp_name --headless +# seed=107 +# exp_name="${task}_ensemble_${ensemble_size}_seed_${seed}_${exp_type}" +# ./isaaclab.sh -p $run_file --task $task --horizon $horizon --num_rollouts $num_rollouts --ensemble_size $ensemble_size --seed $seed --exp_name $exp_name --headless -seed=115 -exp_name="${task}_ensemble_${ensemble_size}_seed_${seed}_${exp_type}" -./isaaclab.sh -p $run_file --task $task --horizon $horizon --num_rollouts $num_rollouts --ensemble_size $ensemble_size --seed $seed --exp_name $exp_name --headless +# seed=115 +# exp_name="${task}_ensemble_${ensemble_size}_seed_${seed}_${exp_type}" +# ./isaaclab.sh -p $run_file --task $task --horizon $horizon --num_rollouts $num_rollouts --ensemble_size $ensemble_size --seed $seed --exp_name $exp_name --headless + +# ## ensemble 5 +# ensemble_size=5 +# seed=101 +# exp_name="${task}_ensemble_${ensemble_size}_seed_${seed}_${exp_type}" +# ./isaaclab.sh -p $run_file --task $task --horizon $horizon --num_rollouts $num_rollouts --ensemble_size $ensemble_size --seed $seed --exp_name $exp_name --headless + +# seed=107 +# exp_name="${task}_ensemble_${ensemble_size}_seed_${seed}_${exp_type}" +# ./isaaclab.sh -p $run_file --task $task --horizon $horizon --num_rollouts $num_rollouts --ensemble_size $ensemble_size --seed $seed --exp_name $exp_name --headless + +# seed=115 +# exp_name="${task}_ensemble_${ensemble_size}_seed_${seed}_${exp_type}" +# ./isaaclab.sh -p $run_file --task $task --horizon $horizon --num_rollouts $num_rollouts --ensemble_size $ensemble_size --seed $seed --exp_name $exp_name --headless diff --git a/scripts/imitation_learning/robomimic/backup_controller_handler.py b/scripts/imitation_learning/robomimic/backup_controller_handler.py index f6e63d51df63..ae3d35908a0d 100644 --- a/scripts/imitation_learning/robomimic/backup_controller_handler.py +++ b/scripts/imitation_learning/robomimic/backup_controller_handler.py @@ -41,7 +41,7 @@ def _get_controller(self): print(f"[DEBUG] Using Place Backup Controller") return BackupControllerPlaceSM(dt=0.01*2, num_envs=self.num_envs, position_threshold=0.02, device='cuda', offset = torch.tensor([[-0.1, 0, 0.1]], device='cuda:0'), goal_quat = self.object_goal_rot, goal_pos= self.object_goal_pos) case "insert": - return BackupControllerInsertSM(0.01*2, self.num_envs,'cuda' , 0.02) + return BackupControllerInsertSM(dt=0.01*2, num_envs=self.num_envs, position_threshold=0.02, device='cuda', offset = torch.tensor([[-0.1, 0, 0.1]], device='cuda:0'), goal_quat = self.object_goal_rot, goal_pos= self.object_goal_pos) case _: raise ValueError(f"Unknown task type {self.tasktype} for backup controller") @@ -71,7 +71,8 @@ def _get_goal_rot(self): return self.env.unwrapped.command_manager.get_command("object_pose") case "insert": - return self.env.unwrapped.scene["vialrack"].data.root_pose_w + goal_rot = self.env.unwrapped.command_manager.get_command("object_pose") + return self.env.unwrapped.command_manager.get_command("object_pose") def reset(self): self.backup_controller.reset_idx() diff --git a/scripts/imitation_learning/robomimic/backup_controller_insert.py b/scripts/imitation_learning/robomimic/backup_controller_insert.py index 5c12bdd60aea..51dfa4f2ca4c 100644 --- a/scripts/imitation_learning/robomimic/backup_controller_insert.py +++ b/scripts/imitation_learning/robomimic/backup_controller_insert.py @@ -25,7 +25,8 @@ class BackupSMWaitTime: LIFT_OBJECT = wp.constant(0.3) MIDPOINT = wp.constant(0.1) APPROACH_ABOVE_GOAL = wp.constant(0.1) - APPROACH_GOAL = wp.constant(6) + APPROACH_GOAL = wp.constant(0.1) + UNGRASP_OBJECT = wp.constant(1.0) class BackupSM: """States for the pick state machine.""" @@ -38,12 +39,27 @@ class BackupSM: MIDPOINT = wp.constant(5) APPROACH_ABOVE_GOAL = wp.constant(6) APPROACH_GOAL = wp.constant(7) - #UNGRASP_OBJECT = wp.constant(7) + UNGRASP_OBJECT = wp.constant(8) @wp.func def distance_below_threshold(current_pos: wp.vec3, desired_pos: wp.vec3, threshold: float) -> bool: return wp.length(current_pos - desired_pos) < threshold +@wp.func +def rotation_within_threshold(current_rot: wp.quat, desired_rot: wp.quat, threshold_rad: float) -> bool: + # Quaternion angular distance: angle = 2 * acos(|q1 · q2|) + # |dot| clamped to [0,1] to guard against floating point > 1 + dot = wp.abs( + current_rot[0] * desired_rot[0] + + current_rot[1] * desired_rot[1] + + current_rot[2] * desired_rot[2] + + current_rot[3] * desired_rot[3] + ) + dot = wp.min(dot, float(1.0)) + angle = float(2.0) * wp.acos(dot) + # print(f"angle : {angle}, threshold : {threshold_rad}") + return angle < threshold_rad + @wp.kernel def infer_state_machine( @@ -55,6 +71,7 @@ def infer_state_machine( current_object_pose: wp.array(dtype=wp.transform), # current object pose final_object_pose: wp.array(dtype=wp.transform), rest_ee_pose : wp.array(dtype=wp.transform), # rest position + goal_quat : wp.array(dtype=wp.transform), # goal rotation des_ee_pose: wp.array(dtype=wp.transform), # what we return gripper_state: wp.array(dtype=float), offset: wp.array(dtype=wp.transform), @@ -69,22 +86,26 @@ def infer_state_machine( state = sm_state[tid] # decide next state if state == BackupSM.REST: + # print("[SM State REST]") # First lift up to a safe height, then go to rest position current_ee_pos = wp.transform_get_translation(ee_pose[tid]) rest_pos = wp.transform_get_translation(rest_ee_pose[tid]) - rest_rot = wp.transform_get_rotation(rest_ee_pose[tid]) + rest_rot = wp.transform_get_rotation(goal_quat[tid]) # If we're below safe height, first move up - safe_height = 0.35 # Safe z-height to avoid collisions + safe_height = 0.15 # Safe z-height to avoid collisions if current_ee_pos.z < safe_height: # Lift straight up first - lift_pos = wp.vec3(current_ee_pos.x, current_ee_pos.y, safe_height) + lift_pos = wp.vec3(current_ee_pos.x, current_ee_pos.y, 0.2) des_ee_pose[tid] = wp.transform(lift_pos, rest_rot) + #print("[SM REST doing safe ee lift]") else: # Already at safe height, move to rest position - des_ee_pose[tid] = rest_ee_pose[tid] + pose_pos = wp.transform_get_translation(rest_ee_pose[tid]) + pose_rot = wp.transform_get_rotation(goal_quat[tid]) + des_ee_pose[tid] = wp.transform(pose_pos, pose_rot) - gripper_state[tid] = GripperState.CLOSE + gripper_state[tid] = GripperState.OPEN # wait for a while if distance_below_threshold( wp.transform_get_translation(ee_pose[tid]), @@ -96,12 +117,14 @@ def infer_state_machine( # move to next state and reset wait time sm_state[tid] = BackupSM.APPROACH_ABOVE_OBJECT sm_wait_time[tid] = 0.0 + else: + print("[SM REST] wait timer ") elif state == BackupSM.APPROACH_ABOVE_OBJECT: # print("[SM] APPR above obj") pose_pos = wp.transform_get_translation(current_object_pose[tid]) - pose_rot = wp.transform_get_rotation(rest_ee_pose[tid]) + pose_rot = wp.transform_get_rotation(goal_quat[tid]) # Go to a safe height ABOVE the object (z + 0.15m) with x offset - safe_above_pos = wp.vec3(pose_pos.x - 0.1, pose_pos.y, pose_pos.z + 0.15) + safe_above_pos = wp.vec3(pose_pos.x, pose_pos.y+0.1, pose_pos.z + 0.20) des_ee_pose[tid] = wp.transform(safe_above_pos, pose_rot) gripper_state[tid] = GripperState.OPEN # Open gripper in preparation if distance_below_threshold( @@ -119,16 +142,22 @@ def infer_state_machine( # print("[SM] approach object") pose_pos = wp.transform_get_translation(current_object_pose[tid]) - pose_rot = wp.transform_get_rotation(rest_ee_pose[tid]) + pose_rot = wp.transform_get_rotation(goal_quat[tid]) # Apply offset in x-direction (5 cm = 0.05 m) - offset_pos = wp.vec3(pose_pos.x -0.1, pose_pos.y+0.02, pose_pos.z+0.02) #sample vial + 0.04 + offset_pos = wp.vec3(pose_pos.x, pose_pos.y+0.1, pose_pos.z) #sample vial + 0.04 des_ee_pose[tid] = wp.transform(offset_pos, pose_rot) gripper_state[tid] = GripperState.OPEN - if distance_below_threshold( + pos_ok = distance_below_threshold( wp.transform_get_translation(ee_pose[tid]), wp.transform_get_translation(des_ee_pose[tid]), - position_threshold, - ): + 0.02, + ) + rot_ok = rotation_within_threshold( + wp.transform_get_rotation(ee_pose[tid]), + wp.transform_get_rotation(goal_quat[tid]), + float(0.05), # ~8.6 degrees tolerance + ) + if pos_ok and rot_ok: if sm_wait_time[tid] >= BackupSMWaitTime.APPROACH_OBJECT: # move to next state and reset wait time print("[SM_INFO] : Moving from APPR_OBJ to GRASP_OBJECT") @@ -137,9 +166,9 @@ def infer_state_machine( elif state == BackupSM.GRASP_OBJECT: # print("[SM] GRASP") pose_pos = wp.transform_get_translation(current_object_pose[tid]) - pose_rot = wp.transform_get_rotation(rest_ee_pose[tid]) + pose_rot = wp.transform_get_rotation(goal_quat[tid]) # Stay at the approach position while closing gripper - offset_pos = wp.vec3(pose_pos.x - 0.1, pose_pos.y + 0.02, pose_pos.z + 0.02) + offset_pos = wp.vec3(pose_pos.x, pose_pos.y+0.1, pose_pos.z) des_ee_pose[tid] = wp.transform(offset_pos, pose_rot) gripper_state[tid] = GripperState.CLOSE # wait for a while @@ -150,11 +179,12 @@ def infer_state_machine( sm_wait_time[tid] = 0.0 elif state == BackupSM.LIFT_OBJECT: - #print("[SM] LIFT") # uses starting position - pose_pos = wp.transform_get_translation(object_pose[tid]) - pose_rot = wp.transform_get_rotation(rest_ee_pose[tid]) + #print("[SM] LIFT") # uses current object position (not start pose) + # pose_pos = wp.transform_get_translation(object_pose[tid]) + pose_pos = wp.transform_get_translation(current_object_pose[tid]) + pose_rot = wp.transform_get_rotation(goal_quat[tid]) # Apply offset in x-direction (5 cm = 0.05 m) - offset_pos = wp.vec3(pose_pos.x -0.1, pose_pos.y+0.02, pose_pos.z+0.2) #sample vial + 0.04 + offset_pos = wp.vec3(pose_pos.x, pose_pos.y+0.1, 0.25) #sample vial + 0.04 des_ee_pose[tid] = wp.transform(offset_pos, pose_rot) gripper_state[tid] = GripperState.CLOSE # wait for a while @@ -172,10 +202,10 @@ def infer_state_machine( elif state == BackupSM.MIDPOINT: #print("[SM] MIDPOINT") pose_pos = wp.transform_get_translation(rest_ee_pose[tid]) - pose_rot = wp.transform_get_rotation(rest_ee_pose[tid]) + pose_rot = wp.transform_get_rotation(goal_quat[tid]) # Apply offset in x-direction (5 cm = 0.05 m) #offset_pos = wp.vec3(pose_pos.x -0.1, pose_pos.y+0.02, pose_pos.z+0.2) #sample vial + 0.04 - des_ee_pose[tid] = rest_ee_pose[tid] + des_ee_pose[tid] = wp.transform(pose_pos, pose_rot) gripper_state[tid] = GripperState.CLOSE # wait for a while if distance_below_threshold( @@ -191,25 +221,28 @@ def infer_state_machine( sm_wait_time[tid] = 0.0 elif state == BackupSM.APPROACH_ABOVE_GOAL: #print("[SM] Approach above goal") - offset_pos = wp.transform_get_translation(offset[tid]) - offset_rot = wp.transform_get_rotation(offset[tid]) - offset_pos = wp.vec3(offset_pos.x, offset_pos.y, offset_pos.z) # raise 25 cm - new_offset = wp.transform(offset_pos, offset_rot) - above_target_pose = wp.transform_multiply(new_offset, final_object_pose[tid]) + offset_pos = wp.transform_get_translation(final_object_pose[tid]) + offset_rot = wp.transform_get_rotation(goal_quat[tid]) + safe_above_pos = wp.vec3(offset_pos.x, offset_pos.y + 0.1, offset_pos.z + 0.25) + #offset_pos = wp.vec3(offset_pos.x, offset_pos.y, offset_pos.z) # raise 25 cm + above_target_pose = wp.transform(safe_above_pos, offset_rot) + #above_target_pose = wp.transform_multiply(new_offset, final_object_pose[tid]) # Blend time for smooth approach - APPROACH_BLEND_TIME = 0.4 # seconds, tune as needed - alpha = wp.clamp(sm_wait_time[tid] / APPROACH_BLEND_TIME, 0.0, 1.0) - # Interpolate position and rotation - current_pos = wp.transform_get_translation(ee_pose[tid]) - current_rot = wp.transform_get_rotation(ee_pose[tid]) + # APPROACH_BLEND_TIME = 0.4 # seconds, tune as needed + # alpha = wp.clamp(sm_wait_time[tid] / APPROACH_BLEND_TIME, 0.0, 1.0) + # # Interpolate position and rotation + # current_pos = wp.transform_get_translation(ee_pose[tid]) + # current_rot = wp.transform_get_rotation(ee_pose[tid]) - target_pos = wp.transform_get_translation(above_target_pose) - target_rot = wp.transform_get_rotation(above_target_pose) + # target_pos = wp.transform_get_translation(above_target_pose) + # target_rot = wp.transform_get_rotation(above_target_pose) - pos_interp = wp.lerp(current_pos, target_pos, alpha) - rot_interp = wp.quat_slerp(current_rot, target_rot, alpha) - # Set interpolated pose - des_ee_pose[tid] = wp.transform(pos_interp, rot_interp) + # pos_interp = wp.lerp(current_pos, target_pos, alpha) + # rot_interp = wp.quat_slerp(current_rot, target_rot, alpha) + # # Set interpolated pose + # des_ee_pose[tid] = wp.transform(pos_interp, rot_interp) + #safe_above_pos = wp.vec3(pose_pos.x - 0.1, pose_pos.y, pose_pos.z + 0.15) + des_ee_pose[tid] = above_target_pose gripper_state[tid] = GripperState.CLOSE # Evaluate readiness to transition if distance_below_threshold( @@ -224,23 +257,36 @@ def infer_state_machine( #print("[SM] approach goal") # print("[SM_INFO] : in state ungrasp") pose_pos = wp.transform_get_translation(final_object_pose[tid]) - pose_rot = wp.transform_get_rotation(final_object_pose[tid]) + pose_rot = wp.transform_get_rotation(goal_quat[tid]) # Apply offset in x-direction (5 cm = 0.05 m) - offset_pos = wp.vec3(pose_pos.x, pose_pos.y+0.1, pose_pos.z+0.02) + offset_pos = wp.vec3(pose_pos.x, pose_pos.y+0.1, pose_pos.z+0.07) des_ee_pose[tid] = wp.transform(offset_pos, pose_rot) gripper_state[tid] = GripperState.CLOSE # wait for a while if distance_below_threshold( wp.transform_get_translation(ee_pose[tid]), wp.transform_get_translation(des_ee_pose[tid]), - 0.01, + 0.02, ): if sm_wait_time[tid] >= BackupSMWaitTime.APPROACH_GOAL: # move to next state and reset wait time - print("[SM_INFO] : Moving from apprach goal to rest") - sm_state[tid] = BackupSM.REST + print("[SM_INFO] : Moving from approach goal to ungrasp") + sm_state[tid] = BackupSM.UNGRASP_OBJECT sm_wait_time[tid] = 0.0 - + elif state == BackupSM.UNGRASP_OBJECT: + # print("[SM] GRASP") + pose_pos = wp.transform_get_translation(final_object_pose[tid]) + pose_rot = wp.transform_get_rotation(goal_quat[tid]) + # Stay at the approach position while closing gripper + offset_pos = wp.vec3(pose_pos.x , pose_pos.y + 0.1, pose_pos.z+0.01) + des_ee_pose[tid] = wp.transform(offset_pos, pose_rot) + gripper_state[tid] = GripperState.OPEN + # wait for a while + if sm_wait_time[tid] >= BackupSMWaitTime.UNGRASP_OBJECT: + # move to next state and reset wait time + print("[SM_INFO] : Moving from UNGRASP_OBJECT to REST") + sm_state[tid] = BackupSM.REST + sm_wait_time[tid] = 0.0 # increment wait time sm_wait_time[tid] = sm_wait_time[tid] + dt[tid] @@ -251,7 +297,7 @@ def infer_state_machine( class BackupControllerInsertSM: - def __init__(self, dt: float, num_envs: int, device: torch.device | str = "cpu", position_threshold=0.02, offset = torch.tensor([[-0.1, 0, 0.1]], device='cuda:0')): + def __init__(self, dt: float, num_envs: int, device: torch.device | str = "cpu", position_threshold=0.02, offset = torch.tensor([[-0.1, 0, 0.1]], device='cuda:0'), goal_quat = torch.tensor([[0, 0, 0, 0,0,0,0]], device='cuda:0'), goal_pos = torch.tensor([[0, 0, 0, 0,0,0,0]], device='cuda:0')): """Initialize the state machine. Args: @@ -261,13 +307,14 @@ def __init__(self, dt: float, num_envs: int, device: torch.device | str = "cpu", """ ## debug - print(f'dt : {dt},num_envs {num_envs}, device {device}, pos_thresh {position_threshold} ') + print(f'dt : {dt},num_envs {num_envs}, \n device {device}, pos_thresh {position_threshold} \n goal_quat {goal_quat} \n goal_pos {goal_pos}') # save parameters self.dt = float(dt) self.num_envs = num_envs self.device = device self.position_threshold = position_threshold - + self.goal_quat = goal_quat + self.goal_pos = goal_pos # initialize state machine self.sm_dt = torch.full((self.num_envs,), self.dt, device=self.device) self.sm_state = torch.full((self.num_envs,), 0, dtype=torch.int32, device=self.device) @@ -278,6 +325,7 @@ def __init__(self, dt: float, num_envs: int, device: torch.device | str = "cpu", self.des_gripper_state = torch.full((self.num_envs,), 0.0, device=self.device) self.final_object_pose = torch.zeros((self.num_envs, 7), device=self.device) self.rest_pose = torch.zeros((self.num_envs, 7), device=self.device) + # approach above object offset self.offset = torch.zeros((self.num_envs, 7), device=self.device) @@ -294,6 +342,7 @@ def __init__(self, dt: float, num_envs: int, device: torch.device | str = "cpu", self.final_object_pose_wp = wp.from_torch(self.final_object_pose, wp.transform) self.rest_pose_wp = wp.from_torch(self.rest_pose, wp.transform) self.offset_wp = wp.from_torch(self.offset, wp.transform) + self.goal_quat_wp = wp.from_torch(self.goal_quat, wp.transform) ## For Debug self.debug_des_pose = torch.zeros((self.num_envs, 7), device=self.device) self.debug_cur_pose = torch.zeros((self.num_envs, 7), device=self.device) @@ -358,22 +407,25 @@ def compute(self, ee_pose: torch.Tensor, start_object_pose: torch.Tensor, curren # print("-" * 80) # print(f"Got given state {sm_state}") current_object_pose = self._ensure_7d(current_object_pose, self.device) - final_object_pose = self._ensure_7d(final_object_pose, self.device) + final_object_pose = self._ensure_7d(self.goal_pos, self.device) ee_pose = self._ensure_7d(ee_pose, self.device) start_object_pose = self._ensure_7d(start_object_pose, self.device) rest_pose = self._ensure_7d(rest_pose, self.device) + goal_quat = self._ensure_7d(self.goal_quat, self.device) # now safe to reorder ee_pose = ee_pose[:, [0, 1, 2, 4, 5, 6, 3]] start_object_pose = start_object_pose[:, [0, 1, 2, 4, 5, 6, 3]] current_object_pose = current_object_pose[:, [0, 1, 2, 4, 5, 6, 3]] final_object_pose = final_object_pose[:, [0, 1, 2, 4, 5, 6, 3]] rest_pose = rest_pose[:, [0, 1, 2, 4, 5, 6, 3]] + goal_quat = goal_quat[:, [0, 1, 2, 4, 5, 6, 3]] # convert to warp ee_pose_wp = wp.from_torch(ee_pose.contiguous(), wp.transform) start_object_pose_wp = wp.from_torch(start_object_pose.contiguous(), wp.transform) current_object_pose_wp = wp.from_torch(current_object_pose.contiguous(), wp.transform) final_object_pose_wp = wp.from_torch(final_object_pose.contiguous(), wp.transform) rest_pose_wp = wp.from_torch(rest_pose.contiguous(), wp.transform) + goal_quat_wp = wp.from_torch(goal_quat.contiguous(), wp.transform) # run state machine # get state from env - convert to warp again self.sm_state[env_num] = sm_state @@ -391,6 +443,7 @@ def compute(self, ee_pose: torch.Tensor, start_object_pose: torch.Tensor, curren current_object_pose_wp, final_object_pose_wp, rest_pose_wp, + goal_quat_wp, self.des_ee_pose_wp, self.des_gripper_state_wp, self.offset_wp, @@ -404,17 +457,21 @@ def compute(self, ee_pose: torch.Tensor, start_object_pose: torch.Tensor, curren des_debug = self.debug_des_pose.detach().cpu().numpy() cur_debug = self.debug_cur_pose.detach().cpu().numpy() debug_state = self.debug_state.detach().cpu().numpy() + goal_quat_debug = self.goal_quat.detach().cpu().numpy() for i in range(self.num_envs): des_pos, des_quat = des_debug[i][:3], des_debug[i][3:] cur_pos, cur_quat = cur_debug[i][:3], cur_debug[i][3:] + goal_quat = goal_quat_debug[i][3:] # compute Euclidean distance with numpy dist = np.linalg.norm(cur_pos - des_pos) # print(f"StateMachine debug state {debug_state}") - # print(f"[Env {i}]") - # print(f" Current -> pos: {cur_pos}, quat: {cur_quat}") - # print(f" Desired -> pos: {des_pos}, quat: {des_quat}") - #print(f" Euclidean distance: {dist:.4f}") - #print(f"returned state : {self.sm_state}") + # # print(f"[Env {i}]") + # print(f" Goal pose : {final_object_pose}") + # print(f" Current -> pos: {cur_pos}, quat: {cur_quat}") + # print(f" Desired -> pos: {des_pos}, quat: {des_quat}") + # print(f" Goal quat -> : {goal_quat}") + # print(f" Euclidean distance: {dist:.4f}") + # #print(f"returned state : {self.sm_state}") # convert transformations back to (w, x, y, z) des_ee_pose = self.des_ee_pose[:, [0, 1, 2, 6, 3, 4, 5]] # convert to torch diff --git a/scripts/imitation_learning/robomimic/play_controller.py b/scripts/imitation_learning/robomimic/play_controller.py index 166b1d518540..88454818a962 100644 --- a/scripts/imitation_learning/robomimic/play_controller.py +++ b/scripts/imitation_learning/robomimic/play_controller.py @@ -118,7 +118,7 @@ def rollout_ensemble(env, success_term, horizon, device): ### SET USE RECOVERY TO FALSE #### # Set up recovery controller - backup_controller = BackupControllerSM(0.01*2, num_envs,'cuda' , 0.01) + backup_controller = BackupControllerInsertSM(0.01*2, num_envs,'cuda' , 0.01) backup_controller.reset_idx() # none so reset all robot = env.unwrapped.scene["robot"] diff --git a/scripts/imitation_learning/robomimic/play_ensemble_v05.py b/scripts/imitation_learning/robomimic/play_ensemble_v05.py index c10cfccf1be6..8bfcfe6fa1d8 100644 --- a/scripts/imitation_learning/robomimic/play_ensemble_v05.py +++ b/scripts/imitation_learning/robomimic/play_ensemble_v05.py @@ -331,7 +331,7 @@ def rollout_ensemble(ensemble, env, success_term, horizon, device, parameters, certain_joint_positions = [] ###### SET UP RECOVERY #### - use_recovery=False + use_recovery=False recovery_activated_during_rollout = 0 print("rollout recovery enabled ? : ", use_recovery) @@ -344,7 +344,7 @@ def rollout_ensemble(ensemble, env, success_term, horizon, device, parameters, max_recovery_steps = 500 # Safety limit for recovery mode ##### CONFIG RECOVERY CONTROLLER #### - backup_controller = BackupController(env, device, tasktype="lift") + backup_controller = BackupController(env, device, tasktype="insert") state_guess = 0 last_state = 0 recovery_mode = False @@ -597,7 +597,7 @@ def main(): # pick_place_ensemble_30 = load_ensemble(device, ensemble_path='scripts/imitation_learning/robomimic/pick_place_ensemble_30_paths.txt') # pick_place_ensemble = load_ensemble(device, ensemble_path='docs/lift/Dev-IK-Rel-v1/best_models/best_model_paths.txt') - pick_place_ensemble = load_ensemble(device, ensemble_path='scripts/imitation_learning/robomimic/insert_paths.txt') + pick_place_ensemble = load_ensemble(device, ensemble_path='scripts/imitation_learning/robomimic/low_insert_paths.txt') # Lets set these to the 0.99 confidence parameters = { @@ -639,6 +639,45 @@ def main(): }, }, + 'vial_insert': { + 0 : { + "confidence_level": 5.3916760000000016e-05, + "window_size": 10, + "max_peaks": 8 + }, + 1 : { + "confidence_level": 5.681572000000001e-05, + "window_size": 10, + "max_peaks": 8 + }, + 2 : { + "confidence_level": 5.7037600000000004e-05, + "window_size": 10, + "max_peaks": 8 + }, + 3 : { + "confidence_level": 0.00012664534000000001, + "window_size": 10, + "max_peaks": 8 + }, + 4 : { + "confidence_level": 5.852008e-05, + "window_size": 10, + "max_peaks": 8 + }, + 5 : { + "confidence_level": 8.595834000000004e-05, + "window_size": 10, + "max_peaks": 8 + }, + 6 : { + "confidence_level": 1.57e-08, + "window_size": 10, + "max_peaks": 8 + }, + + + } } @@ -684,7 +723,7 @@ def main(): for trial in range(args_cli.num_rollouts): print(f"[INFO] Starting trial {trial}") - terminated, traj, recovery_activated_during_rollout, failure = rollout_ensemble(pick_place_ensemble[:args_cli.ensemble_size], env, success_term, args_cli.horizon, device, parameters['beaker_lift'], use_recovery=args_cli.use_recovery, rollout_num=trial) + terminated, traj, recovery_activated_during_rollout, failure = rollout_ensemble(pick_place_ensemble[:args_cli.ensemble_size], env, success_term, args_cli.horizon, device, parameters['vial_insert'], use_recovery=args_cli.use_recovery, rollout_num=trial) # save the uncertainties print("Finished rollout, recovery needed : ", recovery_activated_during_rollout) #print("actions shape : ", traj['actions']) diff --git a/scripts/imitation_learning/robomimic/play_with_backup.py b/scripts/imitation_learning/robomimic/play_with_backup.py index ef7c15bbaf15..39900a4b6f48 100644 --- a/scripts/imitation_learning/robomimic/play_with_backup.py +++ b/scripts/imitation_learning/robomimic/play_with_backup.py @@ -118,7 +118,7 @@ def rollout_ensemble(env, success_term, horizon, device): # Set up recovery controller - backup_controller = BackupController(env, device, tasktype="lift") + backup_controller = BackupController(env, device, tasktype="insert") state_guess = 0 for i in range(horizon): diff --git a/scripts/imitation_learning/robomimic/switching_logic_algorithm.tex b/scripts/imitation_learning/robomimic/switching_logic_algorithm.tex new file mode 100644 index 000000000000..198c0267ee15 --- /dev/null +++ b/scripts/imitation_learning/robomimic/switching_logic_algorithm.tex @@ -0,0 +1,62 @@ +\begin{algorithm}[t] +\caption{Uncertainty-Based Switching Logic for Safe Robot Manipulation} +\label{alg:switching_logic} +\begin{algorithmic}[1] +\Procedure{SwitchingLogic}{$\text{uncertainties}, \text{step}, \text{parameters}$} + \State \textit{// Initialization} + \State $\text{counts} \gets 0$ + \State $\text{triggered} \gets \text{False}$ + + \State \textit{// Iterate over each joint/action dimension} + \For{$j = 0$ \textbf{to} $|\text{uncertainties}| - 1$} + \State $\text{unc} \gets \text{uncertainties}[j]$ + \State $\text{traj}[j][\text{step}] \gets \text{unc}$ + + \State \textit{// Check if sufficient history exists} + \If{$\text{step} > \text{parameters}[j][\text{window\_size}]$} + \State \textit{// Extract sliding window of recent uncertainties} + \State $\text{window\_start} \gets \text{step} - \text{parameters}[j][\text{window\_size}]$ + \State $\text{test\_window} \gets \text{traj}[j][\text{window\_start} : \text{step}]$ + + \State \textit{// Count peaks exceeding confidence threshold} + \State $\text{peaks} \gets 0$ + \For{$x \in \text{test\_window}$} + \If{$x \geq \text{parameters}[j][\text{confidence\_level}]$} + \State $\text{peaks} \gets \text{peaks} + 1$ + \EndIf + \EndFor + + \State \textit{// Check if peaks exceed maximum allowed} + \If{$\text{peaks} \geq \text{parameters}[j][\text{max\_peaks}]$} + \State $\text{counts} \gets \text{counts} + 1$ + \EndIf + \EndIf + \EndFor + + \State \textit{// Trigger safety mechanism if multiple joints show high uncertainty} + \If{$\text{counts} > 2$} + \If{\textbf{not} $\text{triggered}$} + \State $\text{triggered} \gets \text{True}$ + \State \textbf{print} ``Uncertainty triggered safety condition'' + \EndIf + \State \Return \textbf{True} \Comment{Switch to backup controller} + \EndIf + + \State \Return \textbf{False} \Comment{Continue with learned policy} +\EndProcedure + +\State + +\Procedure{HalfwayCheck}{$\text{step}, \text{grasp\_state}, \text{horizon}$} + \State \textit{// Verify subtask completion at midpoint} + \If{$\text{step} = \lfloor \text{horizon} / 2 \rfloor$ \textbf{and not} $\text{grasp\_state}$} + \If{\textbf{not} $\text{triggered}$} + \State \textbf{print} ``Halfway point reached without subtask completion'' + \State $\text{triggered} \gets \text{True}$ + \EndIf + \State \Return \textbf{True} \Comment{Activate recovery} + \EndIf + \State \Return \textbf{False} +\EndProcedure +\end{algorithmic} +\end{algorithm} diff --git a/scripts/imitation_learning/robomimic/train_ensemble.py b/scripts/imitation_learning/robomimic/train_ensemble.py index 9a9f6bb5d954..b4c10121160a 100644 --- a/scripts/imitation_learning/robomimic/train_ensemble.py +++ b/scripts/imitation_learning/robomimic/train_ensemble.py @@ -177,9 +177,9 @@ def train(config: Config, device: str, log_dirs: list[str], ckpt_dirs: list[str] env_meta = FileUtils.get_env_metadata_from_dataset(dataset_path=config.train.data) print(f"config.train.data: {config.train.data}") with config.values_unlocked(): - config.train.data = [{'path': 'docs/place/glassware_insert_demos_generated_split.hdf5'}] + config.train.data = [{'path': 'docs/insert/glassware_insert_demos_generated_split.hdf5'}] ####### HARD CODED MAKE SURE THIS GETS UPDATED ######### - config_list= [{'path': 'docs/place/glassware_insert_demos_generated_split.hdf5'}] + config_list= [{'path': 'docs/insert/glassware_insert_demos_generated_split.hdf5'}] print(f"config.train.data: {config.train.data}") # print(f"config_list: {config_list}") # if isinstance(config.train.data, str): @@ -449,7 +449,7 @@ def keep_data_percentage(trainset, percentage): return trainset -def create_bootstrap_sample(trainset, seed: int, percent: float = 0.25): +def create_bootstrap_sample(trainset, seed: int, percent: float = 1.0): """Create a random sample (without replacement) of a percentage of the training set. Args: diff --git a/source/isaaclab/isaaclab/envs/mdp/observations.py b/source/isaaclab/isaaclab/envs/mdp/observations.py index ac502521aae0..f3daee777aa5 100644 --- a/source/isaaclab/isaaclab/envs/mdp/observations.py +++ b/source/isaaclab/isaaclab/envs/mdp/observations.py @@ -198,6 +198,7 @@ def joint_pos(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg(" """ # extract the used quantities (to enable type-hinting) asset: Articulation = env.scene[asset_cfg.name] + print(f"joint pos {asset.data.joint_pos[:, asset_cfg.joint_ids]}") return asset.data.joint_pos[:, asset_cfg.joint_ids] diff --git a/source/isaaclab_mimic/isaaclab_mimic/datagen/generation.py b/source/isaaclab_mimic/isaaclab_mimic/datagen/generation.py index 1390fdf48eaf..cfdb3ec5e079 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/datagen/generation.py +++ b/source/isaaclab_mimic/isaaclab_mimic/datagen/generation.py @@ -123,9 +123,9 @@ def env_loop( # print("*" * 50, "\033[K") print(summary_text) log_file = "docs/generation_log.txt" - # with open(log_file, 'a') as f: - # f.write(summary_text) - # f.write("\n" + "="*80 + "\n\n") + with open(log_file, 'w') as f: + f.write(summary_text) + f.write("\n" + "="*80 + "\n\n") # termination condition is on enough successes if @guarantee_success or enough attempts otherwise diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/agents/robomimic/bc_rnn_low_dim.json b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/agents/robomimic/bc_rnn_low_dim.json index 6bde3e9661aa..5286c063f046 100755 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/agents/robomimic/bc_rnn_low_dim.json +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/agents/robomimic/bc_rnn_low_dim.json @@ -7,7 +7,7 @@ "terminal_output_to_txt": true, "log_tb": true, "log_wandb": true, - "wandb_proj_name": "vial_insert" + "wandb_proj_name": "vial_insert_1.0_mix" }, "save": { "enabled": true, @@ -55,7 +55,7 @@ "goal_mode": null, "cuda": true, "batch_size": 100, - "num_epochs": 1000, + "num_epochs": 1500, "seed": 101 }, "algo": { diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/dev_ik_rel_vial_insert.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/dev_ik_rel_vial_insert.py index 35707cb590d8..2622a1194628 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/dev_ik_rel_vial_insert.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/dev_ik_rel_vial_insert.py @@ -105,7 +105,7 @@ def __post_init__(self): body_offset=DifferentialInverseKinematicsActionCfg.OffsetCfg(pos=[0.0, 0.0, 0.0]), ) - self.terminations.success= DoneTerm(func=mdp.object_inserted_upright, params={"lower_object_cfg": SceneEntityCfg("vialrack")}) + self.terminations.success= DoneTerm(func=mdp.object_inserted_upright, params={"lower_object_cfg": SceneEntityCfg("vialrack"), "upright_good_deg": 22.5}) #self.terminations.success=DoneTerm(func=mdp.object_stacked_upright, params={"lower_object_cfg": SceneEntityCfg("scale")}) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/mdp/observations.py index a74a3b97cd11..a157efab49ce 100755 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/mdp/observations.py @@ -34,7 +34,7 @@ def object_position_in_robot_root_frame( object_pos_b, _ = subtract_frame_transforms( robot.data.root_state_w[:, :3], robot.data.root_state_w[:, 3:7], object_pos_w ) - + #print(f"object position : {object_pos_b}") return object_pos_b def obstacle_position_in_robot_root_frame( @@ -162,10 +162,10 @@ def is_object_lifted( def ee_frame_pos(env: ManagerBasedRLEnv, ee_frame_cfg: SceneEntityCfg = SceneEntityCfg("ee_frame")) -> torch.Tensor: ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] ee_frame_pos = ee_frame.data.target_pos_w[:, 0, :] - env.scene.env_origins[:, 0:3] - #print(f"EE pos : {ee_frame_pos}") + # print(f"EE pos : {ee_frame_pos}") return ee_frame_pos - +## def ee_frame_quat(env: ManagerBasedRLEnv, ee_frame_cfg: SceneEntityCfg = SceneEntityCfg("ee_frame")) -> torch.Tensor: ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] ee_frame_quat = ee_frame.data.target_quat_w[:, 0, :] @@ -249,7 +249,7 @@ def robot_pose( robot_cfg: SceneEntityCfg = SceneEntityCfg("robot") ): robot: Articulation = env.scene[robot_cfg.name] - # print(robot.data.joint_pos) + print(f"robot data : {robot.data.joint_pos}") return torch.tensor([0.04]) def position_command_error( @@ -301,7 +301,7 @@ def object_tilt (env: ManagerBasedRLEnv, upper_object_cfg: SceneEntityCfg = Scen #print("obs term : tilt : ", tilt_deg) return tilt_deg -def object_knocked(env: ManagerBasedRLEnv, upper_object_cfg: SceneEntityCfg = SceneEntityCfg("object"), max_tilt: float = 45): +def object_knocked(env: ManagerBasedRLEnv, upper_object_cfg: SceneEntityCfg = SceneEntityCfg("object"), max_tilt: float = 90): object: RigidObject = env.scene[upper_object_cfg.name] tilt_deg = upright_tilt_deg(object.data.root_quat_w, object.data.default_root_state[:, 3:7]) # Return a bool tensor with shape (num_envs,) for proper batching diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/mdp/terminations.py index 399ab972ffeb..c9c2d25bd59d 100755 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/mdp/terminations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/mdp/terminations.py @@ -144,6 +144,7 @@ def object_stacked_upright(env: ManagerBasedRLEnv, robot_cfg: SceneEntityCfg = S def upright_tilt_deg(q_cur, q_init) -> torch.Tensor: u_cur, u_init = matrix_from_quat(q_cur)[:, :, 2], matrix_from_quat(q_init)[:, :, 2] + # print(f"upright good {torch.rad2deg(torch.acos((u_cur * u_init).sum(-1).clamp(-1.0, 1.0)))}") return torch.rad2deg(torch.acos((u_cur * u_init).sum(-1).clamp(-1.0, 1.0))) @@ -160,7 +161,7 @@ def object_stacked(env: ManagerBasedRLEnv, robot_cfg: SceneEntityCfg = SceneEnti pos_diff = upper_object.data.root_pos_w - lower_object.data.root_pos_w height_dist = torch.linalg.vector_norm(pos_diff[:, 2:], dim=1) xy_dist = torch.linalg.vector_norm(pos_diff[:, :2], dim=1) - #print(f"For DEBUG : xy_dist : {xy_dist}, height_dist : {height_dist}") + # print(f"For DEBUG : xy_dist : {xy_dist}, height_dist : {height_dist}") stacked = torch.logical_and(xy_dist < xy_threshold, (height_dist - height_diff) < height_threshold) # stacked = torch.logical_and(torch.isclose(robot.data.joint_pos[:, -1], # gripper_open_val.to(env.device), atol=1e-4, rtol=1e-4), stacked) @@ -171,21 +172,23 @@ def object_stacked(env: ManagerBasedRLEnv, robot_cfg: SceneEntityCfg = SceneEnti def object_inserted_upright(env: ManagerBasedRLEnv, robot_cfg: SceneEntityCfg = SceneEntityCfg("robot"), upper_object_cfg: SceneEntityCfg = SceneEntityCfg("object"),lower_object_cfg: SceneEntityCfg = SceneEntityCfg("vialrack"), - xy_threshold: float = 0.01, height_threshold: float = 0.03, height_diff: float = 0.0, + xy_threshold: float = 0.01, height_threshold: float = 0.01, height_diff: float = 0.03, atol=0.0001, rtol=0.0001, - upright_good_deg: float = 30.0, gripper_open_val: torch.Tensor = torch.tensor([0.04]), logging=False) -> torch.Tensor: + upright_good_deg: float = 35, gripper_open_val: torch.Tensor = torch.tensor([0.04]), logging=False) -> torch.Tensor: """Stacked AND sufficiently upright (≤ upright_good_deg).""" robot: Articulation = env.scene[robot_cfg.name] upper: RigidObject = env.scene[upper_object_cfg.name] tilt_deg = upright_tilt_deg(upper.data.root_quat_w, upper.data.default_root_state[:, 3:7]) + print(f"Tilt deg: {tilt_deg}") upright_good = tilt_deg <= upright_good_deg stacked = object_stacked(env, robot_cfg, upper_object_cfg, lower_object_cfg, xy_threshold, height_threshold, height_diff, gripper_open_val) #print(f"For DEBUG : STACKED STATUS : {stacked}") stacked_upright = stacked & upright_good + #print(f"For DEBUG : STACKED STATUS : {stacked_upright}, {stacked}, {upright_good} ") gripper_joint_ids, _ = robot.find_joints(env.cfg.gripper_joint_names) stacked = torch.logical_and( torch.isclose( From eee4486ccea9eaecce7c076aae5d8855a87aa35d Mon Sep 17 00:00:00 2001 From: LJones-alt Date: Wed, 25 Feb 2026 16:30:33 +0000 Subject: [PATCH 2/2] added in insert top down backup controller --- .../robomimic/backup_controller_handler.py | 20 +- .../robomimic/backup_controller_insert_top.py | 478 ++++++++++++++++++ .../robomimic/play_ensemble_v05.py | 25 +- .../cube_lift/config/franka/__init__.py | 14 + .../franka/dev_ik_rel_vial_insert_top_down.py | 127 +++++ .../cube_lift/mdp/observations.py | 4 +- 6 files changed, 651 insertions(+), 17 deletions(-) create mode 100644 scripts/imitation_learning/robomimic/backup_controller_insert_top.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/dev_ik_rel_vial_insert_top_down.py diff --git a/scripts/imitation_learning/robomimic/backup_controller_handler.py b/scripts/imitation_learning/robomimic/backup_controller_handler.py index ae3d35908a0d..a23c301fd063 100644 --- a/scripts/imitation_learning/robomimic/backup_controller_handler.py +++ b/scripts/imitation_learning/robomimic/backup_controller_handler.py @@ -6,6 +6,7 @@ from backup_controller import BackupControllerSM from backup_controller_place import BackupControllerPlaceSM from backup_controller_insert import BackupControllerInsertSM +from backup_controller_insert_top import BackupControllerInsertSM import torch @@ -42,6 +43,9 @@ def _get_controller(self): return BackupControllerPlaceSM(dt=0.01*2, num_envs=self.num_envs, position_threshold=0.02, device='cuda', offset = torch.tensor([[-0.1, 0, 0.1]], device='cuda:0'), goal_quat = self.object_goal_rot, goal_pos= self.object_goal_pos) case "insert": return BackupControllerInsertSM(dt=0.01*2, num_envs=self.num_envs, position_threshold=0.02, device='cuda', offset = torch.tensor([[-0.1, 0, 0.1]], device='cuda:0'), goal_quat = self.object_goal_rot, goal_pos= self.object_goal_pos) + case "insert_top": + ### check this quat + return BackupControllerInsertSM(dt=0.01*2, num_envs=self.num_envs, position_threshold=0.02, device='cuda', offset = torch.tensor([[-0.1, 0, 0.1]], device='cuda:0'), goal_quat = self.object_goal_rot, goal_pos= self.object_goal_pos) case _: raise ValueError(f"Unknown task type {self.tasktype} for backup controller") @@ -55,7 +59,9 @@ def _get_goal_pos(self): goal_pos = self.env.unwrapped.scene["scale"].data.root_pose_w # print(f"[DEBUG] Getting goal pos for place task {goal_pos}") return self.env.unwrapped.scene["scale"].data.root_pose_w - + case "insert_top": + # still vial rack + return self.env.unwrapped.scene["vialrack"].data.root_pose_w case "insert": return self.env.unwrapped.scene["vialrack"].data.root_pose_w @@ -69,7 +75,9 @@ def _get_goal_rot(self): goal_rot = self.env.unwrapped.command_manager.get_command("object_pose") print(f"[DEBUG] Getting goal quat for place task {goal_rot}") return self.env.unwrapped.command_manager.get_command("object_pose") - + case "insert_top": + goal_rot = torch.tensor([[0,1,0,0]], device=self.device)#self.env.unwrapped.command_manager.get_command("object_pose") + return goal_rot #self.env.unwrapped.command_manager.get_command("object_pose") case "insert": goal_rot = self.env.unwrapped.command_manager.get_command("object_pose") return self.env.unwrapped.command_manager.get_command("object_pose") @@ -90,9 +98,15 @@ def _get_rest_pos(self): # lets change this into the 6 element tensor they are expecting roll,pitch,yaw = euler_xyz_from_quat(ee_recovery_rot) rest_pos = torch.cat([rest_pos, ee_recovery_rot], dim =-1) + + #print(f"rpy version :{torch.cat([rest_pos,roll.unsqueeze(0), pitch.unsqueeze(0), yaw.unsqueeze(0)], dim =-1)}") + if self.tasktype == "insert_top": + rest_pos = torch.tensor([[ 0.6226, -0.0621, 0.3555]], device=self.device) + rest_rot = tensor([[ 0, 1, 0, 0]], device=self.device) + roll,pitch,yaw = euler_xyz_from_quat(rest_rot) + rest_pos = torch.cat([rest_pos,rest_rot], dim =-1) print(f"Rest pos {rest_pos}") print(f"rpy : {roll}, {pitch}, {yaw}") - #print(f"rpy version :{torch.cat([rest_pos,roll.unsqueeze(0), pitch.unsqueeze(0), yaw.unsqueeze(0)], dim =-1)}") return rest_pos def _setup_robot(self): diff --git a/scripts/imitation_learning/robomimic/backup_controller_insert_top.py b/scripts/imitation_learning/robomimic/backup_controller_insert_top.py new file mode 100644 index 000000000000..83abe86e3cfd --- /dev/null +++ b/scripts/imitation_learning/robomimic/backup_controller_insert_top.py @@ -0,0 +1,478 @@ +import warp as wp +import gymnasium as gym +import torch +from collections.abc import Sequence +import numpy as np + +# this is a backup state machine controller +# initialize warp +wp.init() + + +class GripperState: + """States for the gripper.""" + + OPEN = wp.constant(1.0) + CLOSE = wp.constant(-1.0) + +class BackupSMWaitTime: + """Additional wait times (in s) for states for before switching.""" + + REST = wp.constant(0.5) + APPROACH_ABOVE_OBJECT = wp.constant(0.3) + APPROACH_OBJECT = wp.constant(0.3) + GRASP_OBJECT = wp.constant(0.1) # 1 second wait + 0.3 for gripper to close + LIFT_OBJECT = wp.constant(0.3) + MIDPOINT = wp.constant(0.1) + APPROACH_ABOVE_GOAL = wp.constant(0.1) + APPROACH_GOAL = wp.constant(0.1) + UNGRASP_OBJECT = wp.constant(1.0) + +class BackupSM: + """States for the pick state machine.""" + + REST = wp.constant(0) + APPROACH_ABOVE_OBJECT = wp.constant(1) + APPROACH_OBJECT = wp.constant(2) + GRASP_OBJECT = wp.constant(3) + LIFT_OBJECT = wp.constant(4) + MIDPOINT = wp.constant(5) + APPROACH_ABOVE_GOAL = wp.constant(6) + APPROACH_GOAL = wp.constant(7) + UNGRASP_OBJECT = wp.constant(8) + +@wp.func +def distance_below_threshold(current_pos: wp.vec3, desired_pos: wp.vec3, threshold: float) -> bool: + return wp.length(current_pos - desired_pos) < threshold + +@wp.func +def rotation_within_threshold(current_rot: wp.quat, desired_rot: wp.quat, threshold_rad: float) -> bool: + # Quaternion angular distance: angle = 2 * acos(|q1 · q2|) + # |dot| clamped to [0,1] to guard against floating point > 1 + dot = wp.abs( + current_rot[0] * desired_rot[0] + + current_rot[1] * desired_rot[1] + + current_rot[2] * desired_rot[2] + + current_rot[3] * desired_rot[3] + ) + dot = wp.min(dot, float(1.0)) + angle = float(2.0) * wp.acos(dot) + # print(f"angle : {angle}, threshold : {threshold_rad}") + return angle < threshold_rad + + +@wp.kernel +def infer_state_machine( + dt: wp.array(dtype=float), + sm_state: wp.array(dtype=int), + sm_wait_time: wp.array(dtype=float), + ee_pose: wp.array(dtype=wp.transform), # current ee pose + object_pose: wp.array(dtype=wp.transform), # start object location + current_object_pose: wp.array(dtype=wp.transform), # current object pose + final_object_pose: wp.array(dtype=wp.transform), + rest_ee_pose : wp.array(dtype=wp.transform), # rest position + goal_quat : wp.array(dtype=wp.transform), # goal rotation + des_ee_pose: wp.array(dtype=wp.transform), # what we return + gripper_state: wp.array(dtype=float), + offset: wp.array(dtype=wp.transform), + position_threshold: float, + debug_des_pose: wp.array(dtype=wp.transform), # NEW + debug_cur_pose: wp.array(dtype=wp.transform), # NEW + debug_state: wp.array(dtype=int), # New +): + # retrieve thread id + tid = wp.tid() + # retrieve state machine state + state = sm_state[tid] + # decide next state + if state == BackupSM.REST: + # print("[SM State REST]") + # First lift up to a safe height, then go to rest position + current_ee_pos = wp.transform_get_translation(ee_pose[tid]) + rest_pos = wp.transform_get_translation(rest_ee_pose[tid]) + rest_rot = wp.transform_get_rotation(goal_quat[tid]) + + # If we're below safe height, first move up + safe_height = 0.15 # Safe z-height to avoid collisions + if current_ee_pos.z < safe_height: + # Lift straight up first + lift_pos = wp.vec3(current_ee_pos.x, current_ee_pos.y, 0.2) + des_ee_pose[tid] = wp.transform(lift_pos, rest_rot) + #print("[SM REST doing safe ee lift]") + else: + # Already at safe height, move to rest position + pose_pos = wp.transform_get_translation(rest_ee_pose[tid]) + pose_rot = wp.transform_get_rotation(goal_quat[tid]) + des_ee_pose[tid] = wp.transform(pose_pos, pose_rot) + + gripper_state[tid] = GripperState.OPEN + # wait for a while + if distance_below_threshold( + wp.transform_get_translation(ee_pose[tid]), + wp.transform_get_translation(rest_ee_pose[tid]), + 0.02, + ): + if sm_wait_time[tid] >= BackupSMWaitTime.REST: + print("[SM_INFO] : Moving from REST to APPROACH_ABOVE_OBJECT") + # move to next state and reset wait time + sm_state[tid] = BackupSM.APPROACH_ABOVE_OBJECT + sm_wait_time[tid] = 0.0 + else: + print("[SM REST] wait timer ") + elif state == BackupSM.APPROACH_ABOVE_OBJECT: + # print("[SM] APPR above obj") + pose_pos = wp.transform_get_translation(current_object_pose[tid]) + pose_rot = wp.transform_get_rotation(goal_quat[tid]) + # Go to a safe height ABOVE the object (z + 0.15m) with x offset + # left this in, but no offset due to new orientation + safe_above_pos = wp.vec3(pose_pos.x, pose_pos.y, pose_pos.z+0.15) + des_ee_pose[tid] = wp.transform(safe_above_pos, pose_rot) + gripper_state[tid] = GripperState.OPEN # Open gripper in preparation + if distance_below_threshold( + wp.transform_get_translation(ee_pose[tid]), + wp.transform_get_translation(des_ee_pose[tid]), + position_threshold, + ): + # wait for a while + if sm_wait_time[tid] >= BackupSMWaitTime.APPROACH_ABOVE_OBJECT: + # move to next state and reset wait time + print("[SM_INFO] : Moving from APPR_ABOVE to APPROACH_OBJECT") + sm_state[tid] = BackupSM.APPROACH_OBJECT + sm_wait_time[tid] = 0.0 + elif state == BackupSM.APPROACH_OBJECT: + # print("[SM] approach object") + + pose_pos = wp.transform_get_translation(current_object_pose[tid]) + pose_rot = wp.transform_get_rotation(goal_quat[tid]) + # Apply offset in x-direction (5 cm = 0.05 m) + offset_pos = wp.vec3(pose_pos.x, pose_pos.y, pose_pos.z) #sample vial + 0.04 + des_ee_pose[tid] = wp.transform(offset_pos, pose_rot) + gripper_state[tid] = GripperState.OPEN + pos_ok = distance_below_threshold( + wp.transform_get_translation(ee_pose[tid]), + wp.transform_get_translation(des_ee_pose[tid]), + 0.01, + ) + + if pos_ok: + if sm_wait_time[tid] >= BackupSMWaitTime.APPROACH_OBJECT: + # move to next state and reset wait time + print("[SM_INFO] : Moving from APPR_OBJ to GRASP_OBJECT") + sm_state[tid] = BackupSM.GRASP_OBJECT + sm_wait_time[tid] = 0.0 + elif state == BackupSM.GRASP_OBJECT: + # print("[SM] GRASP") + pose_pos = wp.transform_get_translation(current_object_pose[tid]) + pose_rot = wp.transform_get_rotation(goal_quat[tid]) + # Stay at the approach position while closing gripper + offset_pos = wp.vec3(pose_pos.x, pose_pos.y, pose_pos.z) + des_ee_pose[tid] = wp.transform(offset_pos, pose_rot) + gripper_state[tid] = GripperState.CLOSE + # wait for a while + if sm_wait_time[tid] >= BackupSMWaitTime.GRASP_OBJECT: + # move to next state and reset wait time + print("[SM_INFO] : Moving from GRSP to LIFT_OBJECT") + sm_state[tid] = BackupSM.LIFT_OBJECT + sm_wait_time[tid] = 0.0 + + elif state == BackupSM.LIFT_OBJECT: + #print("[SM] LIFT") # uses current object position (not start pose) + # pose_pos = wp.transform_get_translation(object_pose[tid]) + pose_pos = wp.transform_get_translation(current_object_pose[tid]) + pose_rot = wp.transform_get_rotation(goal_quat[tid]) + # Apply offset in x-direction (5 cm = 0.05 m) + offset_pos = wp.vec3(pose_pos.x, pose_pos.y, 0.25) #sample vial + 0.04 + des_ee_pose[tid] = wp.transform(offset_pos, pose_rot) + gripper_state[tid] = GripperState.CLOSE + # wait for a while + if distance_below_threshold( + wp.transform_get_translation(ee_pose[tid]), + wp.transform_get_translation(des_ee_pose[tid]), + 0.03, + ): + # wait for a while + if sm_wait_time[tid] >= BackupSMWaitTime.LIFT_OBJECT: + # move to next state and reset wait time + print("[SM_INFO] : Moving from LIFT to MIDPOINT") + sm_state[tid] = BackupSM.MIDPOINT + sm_wait_time[tid] = 0.0 + elif state == BackupSM.MIDPOINT: + #print("[SM] MIDPOINT") + pose_pos = wp.transform_get_translation(rest_ee_pose[tid]) + pose_rot = wp.transform_get_rotation(goal_quat[tid]) + # Apply offset in x-direction (5 cm = 0.05 m) + #offset_pos = wp.vec3(pose_pos.x -0.1, pose_pos.y+0.02, pose_pos.z+0.2) #sample vial + 0.04 + des_ee_pose[tid] = wp.transform(pose_pos, pose_rot) + gripper_state[tid] = GripperState.CLOSE + # wait for a while + if distance_below_threshold( + wp.transform_get_translation(ee_pose[tid]), + wp.transform_get_translation(des_ee_pose[tid]), + 0.03, + ): + # wait for a while + if sm_wait_time[tid] >= BackupSMWaitTime.LIFT_OBJECT: + # move to next state and reset wait time + print("[SM_INFO] : Moving from MID to APPR goal") + sm_state[tid] = BackupSM.APPROACH_ABOVE_GOAL + sm_wait_time[tid] = 0.0 + elif state == BackupSM.APPROACH_ABOVE_GOAL: + #print("[SM] Approach above goal") + offset_pos = wp.transform_get_translation(final_object_pose[tid]) + offset_rot = wp.transform_get_rotation(goal_quat[tid]) + safe_above_pos = wp.vec3(offset_pos.x, offset_pos.y, offset_pos.z+0.15) + #offset_pos = wp.vec3(offset_pos.x, offset_pos.y, offset_pos.z) # raise 25 cm + above_target_pose = wp.transform(safe_above_pos, offset_rot) + #above_target_pose = wp.transform_multiply(new_offset, final_object_pose[tid]) + # Blend time for smooth approach + # APPROACH_BLEND_TIME = 0.4 # seconds, tune as needed + # alpha = wp.clamp(sm_wait_time[tid] / APPROACH_BLEND_TIME, 0.0, 1.0) + # # Interpolate position and rotation + # current_pos = wp.transform_get_translation(ee_pose[tid]) + # current_rot = wp.transform_get_rotation(ee_pose[tid]) + + # target_pos = wp.transform_get_translation(above_target_pose) + # target_rot = wp.transform_get_rotation(above_target_pose) + + # pos_interp = wp.lerp(current_pos, target_pos, alpha) + # rot_interp = wp.quat_slerp(current_rot, target_rot, alpha) + # # Set interpolated pose + # des_ee_pose[tid] = wp.transform(pos_interp, rot_interp) + #safe_above_pos = wp.vec3(pose_pos.x - 0.1, pose_pos.y, pose_pos.z + 0.15) + des_ee_pose[tid] = above_target_pose + gripper_state[tid] = GripperState.CLOSE + # Evaluate readiness to transition + if distance_below_threshold( + wp.transform_get_translation(ee_pose[tid]), + wp.transform_get_translation(des_ee_pose[tid]), + 0.02): + if sm_wait_time[tid] >= BackupSMWaitTime.APPROACH_ABOVE_GOAL: + print("[SM_INFO] : Moving from APPR_ABOVE to APPROACH_GOAL") + sm_state[tid] = BackupSM.APPROACH_GOAL + sm_wait_time[tid] = 0.0 + elif state == BackupSM.APPROACH_GOAL: + #print("[SM] approach goal") + # print("[SM_INFO] : in state ungrasp") + pose_pos = wp.transform_get_translation(final_object_pose[tid]) + pose_rot = wp.transform_get_rotation(goal_quat[tid]) + # Apply offset in x-direction (5 cm = 0.05 m) + # tune the z offset here ? + offset_pos = wp.vec3(pose_pos.x, pose_pos.y, pose_pos.z+0.02) + des_ee_pose[tid] = wp.transform(offset_pos, pose_rot) + gripper_state[tid] = GripperState.CLOSE + # wait for a while + if distance_below_threshold( + wp.transform_get_translation(ee_pose[tid]), + wp.transform_get_translation(des_ee_pose[tid]), + 0.02, + ): + if sm_wait_time[tid] >= BackupSMWaitTime.APPROACH_GOAL: + # move to next state and reset wait time + print("[SM_INFO] : Moving from approach goal to ungrasp") + sm_state[tid] = BackupSM.UNGRASP_OBJECT + sm_wait_time[tid] = 0.0 + elif state == BackupSM.UNGRASP_OBJECT: + # print("[SM] GRASP") + pose_pos = wp.transform_get_translation(final_object_pose[tid]) + pose_rot = wp.transform_get_rotation(goal_quat[tid]) + # Stay at the approach position while closing gripper + offset_pos = wp.vec3(pose_pos.x , pose_pos.y, pose_pos.z+0.02) + des_ee_pose[tid] = wp.transform(offset_pos, pose_rot) + gripper_state[tid] = GripperState.OPEN + # wait for a while + if sm_wait_time[tid] >= BackupSMWaitTime.UNGRASP_OBJECT: + # move to next state and reset wait time + print("[SM_INFO] : Moving from UNGRASP_OBJECT to REST") + sm_state[tid] = BackupSM.REST + sm_wait_time[tid] = 0.0 + + # increment wait time + sm_wait_time[tid] = sm_wait_time[tid] + dt[tid] + debug_des_pose[tid] = des_ee_pose[tid] + debug_cur_pose[tid] = ee_pose[tid] + debug_state[tid] = sm_state[tid] + + + +class BackupControllerInsertSM: + def __init__(self, dt: float, num_envs: int, device: torch.device | str = "cpu", position_threshold=0.02, offset = torch.tensor([[-0.1, 0, 0.1]], device='cuda:0'), goal_quat = torch.tensor([[0, 0, 0, 0,0,0,0]], device='cuda:0'), goal_pos = torch.tensor([[0, 0, 0, 0,0,0,0]], device='cuda:0')): + """Initialize the state machine. + + Args: + dt: The environment time step. + num_envs: The number of environments to simulate. + device: The device to run the state machine on. + """ + ## debug + + print(f'dt : {dt},num_envs {num_envs}, \n device {device}, pos_thresh {position_threshold} \n goal_quat {goal_quat} \n goal_pos {goal_pos}') + # save parameters + self.dt = float(dt) + self.num_envs = num_envs + self.device = device + self.position_threshold = position_threshold + self.goal_quat = goal_quat + self.goal_pos = goal_pos + # initialize state machine + self.sm_dt = torch.full((self.num_envs,), self.dt, device=self.device) + self.sm_state = torch.full((self.num_envs,), 0, dtype=torch.int32, device=self.device) + self.sm_wait_time = torch.zeros((self.num_envs,), device=self.device) + + # desired state + self.des_ee_pose = torch.zeros((self.num_envs, 7), device=self.device) + self.des_gripper_state = torch.full((self.num_envs,), 0.0, device=self.device) + self.final_object_pose = torch.zeros((self.num_envs, 7), device=self.device) + self.rest_pose = torch.zeros((self.num_envs, 7), device=self.device) + + # approach above object offset + + self.offset = torch.zeros((self.num_envs, 7), device=self.device) + self.offset[:, 2] = 0.2 + self.offset[:,0] = -0.1 + self.offset[:, -1] = 1.0 # warp expects quaternion as (x, y, z, w) + + # convert to warp + self.sm_dt_wp = wp.from_torch(self.sm_dt, wp.float32) + self.sm_state_wp = wp.from_torch(self.sm_state, wp.int32) + self.sm_wait_time_wp = wp.from_torch(self.sm_wait_time, wp.float32) + self.des_ee_pose_wp = wp.from_torch(self.des_ee_pose, wp.transform) + self.des_gripper_state_wp = wp.from_torch(self.des_gripper_state, wp.float32) + self.final_object_pose_wp = wp.from_torch(self.final_object_pose, wp.transform) + self.rest_pose_wp = wp.from_torch(self.rest_pose, wp.transform) + self.offset_wp = wp.from_torch(self.offset, wp.transform) + self.goal_quat_wp = wp.from_torch(self.goal_quat, wp.transform) + ## For Debug + self.debug_des_pose = torch.zeros((self.num_envs, 7), device=self.device) + self.debug_cur_pose = torch.zeros((self.num_envs, 7), device=self.device) + self.debug_state = torch.full((self.num_envs,), 0, dtype=torch.int32, device=self.device) + + self.debug_des_pose_wp = wp.from_torch(self.debug_des_pose, wp.transform) + self.debug_cur_pose_wp = wp.from_torch(self.debug_cur_pose, wp.transform) + self.debug_state_wp = wp.from_torch(self.debug_state, wp.int32) + + def reset_idx(self, env_ids: Sequence[int] = None): + """Reset the state machine.""" + print(f'[SM RESET]') + if env_ids is None: + env_ids = slice(None) + self.sm_state[env_ids] = 0 + self.sm_wait_time[env_ids] = 0.0 + + def _ensure_7d(self,tensor: torch.Tensor, device): + # accepts (N,3) -> returns (N,7) with identity quaternion + if tensor.ndim != 2: + raise ValueError(f"expected 2D tensor, got shape {tensor.shape}") + if tensor.shape[1] == 7: + return tensor + if tensor.shape[1] == 3: + quat = torch.tensor([0.0,0.0,0.0,1.0], device=device, dtype=tensor.dtype) + quat = quat.unsqueeze(0).expand(tensor.shape[0], -1) + return torch.cat([tensor, quat], dim=-1) + raise ValueError(f"pose must have 3 or 7 elements, got {tensor.shape[1]}") + + + def compute(self, ee_pose: torch.Tensor, start_object_pose: torch.Tensor, current_object_pose: torch.Tensor, final_object_pose: torch.Tensor, rest_pose: torch.Tensor, sm_state: int =0, env_num = 0) -> torch.Tensor: + """Compute the desired state of the robot's end-effector and the gripper.""" + # convert all transformations from (w, x, y, z) to (x, y, z, w) + # print(f'[SM DEBUG] Doing State Comp') + # print(f'[SM DEBUG] ee_pose {ee_pose}') + # print(f'[SM DEBUG] object pose {object_pose}') + # print(f'[SM DEBUG] desired pose {des_object_pose}') + # print(f'[SM DEBUG] final pose {final_object_pose}') + # print(f'[SM DEBUG] rest pose {rest_pose}') + # before reordering + # print("\n[BackupControllerSM::compute] called with inputs:") + # print(f" sm_state: {sm_state}, env_num: {env_num}") + + # def short_tensor(t, name): + # if t is None: + # print(f" {name}: None") + # elif not isinstance(t, torch.Tensor): + # print(f" {name}: {t} (type={type(t)})") + # else: + # shape = tuple(t.shape) + # device = t.device + # dtype = t.dtype + # # show first few elements flattened + # vals = t.flatten()[:6].detach().cpu().numpy() + # print(f" {name}: shape={shape}, device={device}, dtype={dtype}, values={vals}") + + # short_tensor(ee_pose, "ee_pose") + # short_tensor(object_pose, "object_pose") + # short_tensor(des_object_pose, "des_object_pose") + # short_tensor(final_object_pose, "final_object_pose") + # short_tensor(rest_pose, "rest_pose") + # print("-" * 80) + # print(f"Got given state {sm_state}") + current_object_pose = self._ensure_7d(current_object_pose, self.device) + final_object_pose = self._ensure_7d(self.goal_pos, self.device) + ee_pose = self._ensure_7d(ee_pose, self.device) + start_object_pose = self._ensure_7d(start_object_pose, self.device) + rest_pose = self._ensure_7d(rest_pose, self.device) + goal_quat = self._ensure_7d(self.goal_quat, self.device) + # now safe to reorder + ee_pose = ee_pose[:, [0, 1, 2, 4, 5, 6, 3]] + start_object_pose = start_object_pose[:, [0, 1, 2, 4, 5, 6, 3]] + current_object_pose = current_object_pose[:, [0, 1, 2, 4, 5, 6, 3]] + final_object_pose = final_object_pose[:, [0, 1, 2, 4, 5, 6, 3]] + rest_pose = rest_pose[:, [0, 1, 2, 4, 5, 6, 3]] + goal_quat = goal_quat[:, [0, 1, 2, 4, 5, 6, 3]] + # convert to warp + ee_pose_wp = wp.from_torch(ee_pose.contiguous(), wp.transform) + start_object_pose_wp = wp.from_torch(start_object_pose.contiguous(), wp.transform) + current_object_pose_wp = wp.from_torch(current_object_pose.contiguous(), wp.transform) + final_object_pose_wp = wp.from_torch(final_object_pose.contiguous(), wp.transform) + rest_pose_wp = wp.from_torch(rest_pose.contiguous(), wp.transform) + goal_quat_wp = wp.from_torch(goal_quat.contiguous(), wp.transform) + # run state machine + # get state from env - convert to warp again + self.sm_state[env_num] = sm_state + self.sm_state_wp = wp.from_torch(self.sm_state, wp.int32) + # launch + wp.launch( + kernel=infer_state_machine, + dim=self.num_envs, + inputs=[ + self.sm_dt_wp, + self.sm_state_wp, + self.sm_wait_time_wp, + ee_pose_wp, + start_object_pose_wp, + current_object_pose_wp, + final_object_pose_wp, + rest_pose_wp, + goal_quat_wp, + self.des_ee_pose_wp, + self.des_gripper_state_wp, + self.offset_wp, + self.position_threshold, + self.debug_des_pose_wp, + self.debug_cur_pose_wp, + self.debug_state_wp, + ], + device=self.device, + ) + des_debug = self.debug_des_pose.detach().cpu().numpy() + cur_debug = self.debug_cur_pose.detach().cpu().numpy() + debug_state = self.debug_state.detach().cpu().numpy() + goal_quat_debug = self.goal_quat.detach().cpu().numpy() + for i in range(self.num_envs): + des_pos, des_quat = des_debug[i][:3], des_debug[i][3:] + cur_pos, cur_quat = cur_debug[i][:3], cur_debug[i][3:] + goal_quat = goal_quat_debug[i][3:] + # compute Euclidean distance with numpy + dist = np.linalg.norm(cur_pos - des_pos) + # print(f"StateMachine debug state {debug_state}") + # # print(f"[Env {i}]") + # print(f" Goal pose : {final_object_pose}") + # print(f" Current -> pos: {cur_pos}, quat: {cur_quat}") + # print(f" Desired -> pos: {des_pos}, quat: {des_quat}") + # print(f" Goal quat -> : {goal_quat}") + # print(f" Euclidean distance: {dist:.4f}") + # #print(f"returned state : {self.sm_state}") + # convert transformations back to (w, x, y, z) + des_ee_pose = self.des_ee_pose[:, [0, 1, 2, 6, 3, 4, 5]] + # convert to torch + #print(f'[SM DEBUG] new desired poition : {des_ee_pose}, gripper state {self.des_gripper_state}') + return des_ee_pose , self.des_gripper_state, self.sm_state.clone().item() + #return torch.cat([des_ee_pose, self.des_gripper_state.unsqueeze(-1)], dim=-1) \ No newline at end of file diff --git a/scripts/imitation_learning/robomimic/play_ensemble_v05.py b/scripts/imitation_learning/robomimic/play_ensemble_v05.py index 8bfcfe6fa1d8..df645bdfab14 100644 --- a/scripts/imitation_learning/robomimic/play_ensemble_v05.py +++ b/scripts/imitation_learning/robomimic/play_ensemble_v05.py @@ -331,7 +331,7 @@ def rollout_ensemble(ensemble, env, success_term, horizon, device, parameters, certain_joint_positions = [] ###### SET UP RECOVERY #### - use_recovery=False + use_recovery=True recovery_activated_during_rollout = 0 print("rollout recovery enabled ? : ", use_recovery) @@ -344,7 +344,7 @@ def rollout_ensemble(ensemble, env, success_term, horizon, device, parameters, max_recovery_steps = 500 # Safety limit for recovery mode ##### CONFIG RECOVERY CONTROLLER #### - backup_controller = BackupController(env, device, tasktype="insert") + backup_controller = BackupController(env, device, tasktype="place") state_guess = 0 last_state = 0 recovery_mode = False @@ -593,51 +593,52 @@ def main(): # stack_cube_ensemble = load_ensemble(device, ensemble_path='scripts/imitation_learning/robomimic/stack_cube_ensemble.txt') #pick_place_ensemble = load_ensemble(device, ensemble_path='scripts/imitation_learning/robomimic/ensembles.txt') - #pick_place_ensemble = load_ensemble(device, ensemble_path='docs/place/low/Dev-IK-Rel-Place-v0/best_models/best_model_paths.txt') + pick_place_ensemble = load_ensemble(device, ensemble_path='docs/place/low/Dev-IK-Rel-Place-v0/best_models/best_model_paths.txt') # pick_place_ensemble_30 = load_ensemble(device, ensemble_path='scripts/imitation_learning/robomimic/pick_place_ensemble_30_paths.txt') # pick_place_ensemble = load_ensemble(device, ensemble_path='docs/lift/Dev-IK-Rel-v1/best_models/best_model_paths.txt') - pick_place_ensemble = load_ensemble(device, ensemble_path='scripts/imitation_learning/robomimic/low_insert_paths.txt') + #pick_place_ensemble = load_ensemble(device, ensemble_path='scripts/imitation_learning/robomimic/low_insert_paths.txt') # Lets set these to the 0.99 confidence parameters = { 'beaker_lift' :{ 0 : { - "confidence_level": 0.0002104382060000009, + "confidence_level": 5.373953000000001e-05, "window_size": 10, "max_peaks": 8 }, 1 : { - "confidence_level": 0.00020459246700000051, + "confidence_level": 5.378578e-05, "window_size": 10, "max_peaks": 8 }, 2 : { - "confidence_level": 0.00023551160300000013, + "confidence_level": 6.186318e-05, "window_size": 10, "max_peaks": 8 }, 3 : { - "confidence_level": 0.0003701742110000007, + "confidence_level": 7.819939000000004e-05, "window_size": 10, "max_peaks": 8 }, 4 : { - "confidence_level": 0.0002939644620000011, + "confidence_level": 0.00012073204000000005, "window_size": 10, "max_peaks": 8 }, 5 : { - "confidence_level": 0.0006407762280000002, + "confidence_level": 8.867226000000004e-05, "window_size": 10, "max_peaks": 8 }, 6 : { - "confidence_level": 0.49527883679, + "confidence_level": 1.57e-08, "window_size": 10, "max_peaks": 8 }, + }, 'vial_insert': { 0 : { @@ -723,7 +724,7 @@ def main(): for trial in range(args_cli.num_rollouts): print(f"[INFO] Starting trial {trial}") - terminated, traj, recovery_activated_during_rollout, failure = rollout_ensemble(pick_place_ensemble[:args_cli.ensemble_size], env, success_term, args_cli.horizon, device, parameters['vial_insert'], use_recovery=args_cli.use_recovery, rollout_num=trial) + terminated, traj, recovery_activated_during_rollout, failure = rollout_ensemble(pick_place_ensemble[:args_cli.ensemble_size], env, success_term, args_cli.horizon, device, parameters['beaker_lift'], use_recovery=args_cli.use_recovery, rollout_num=trial) # save the uncertainties print("Finished rollout, recovery needed : ", recovery_activated_during_rollout) #print("actions shape : ", traj['actions']) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/__init__.py index d02f0da3a92a..4e043f645a0c 100755 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/__init__.py @@ -187,6 +187,20 @@ disable_env_checker=True, ) +gym.register( + id="Dev-IK-Rel-Insert-v1", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + kwargs={ + "env_cfg_entry_point": f"{__name__}.dev_ik_rel_vial_insert_top_down:FrankaDevEnvCfg", + "robomimic_bc_cfg_entry_point": os.path.join(agents.__path__[0], "robomimic/bc_rnn_low_dim.json"), + "robomimic_bc_trans_cfg_entry_point": os.path.join(agents.__path__[0], "robomimic/bc_trans.json"), + "robomimic_hbc_cfg_entry_point": os.path.join(agents.__path__[0], "robomimic/hbc.json"), + "robomimic_bcq_cfg_entry_point": os.path.join(agents.__path__[0], "robomimic/bcq.json"), + "robomimic_diffusion_policy_cfg_entry_point": os.path.join(agents.__path__[0], "robomimic/diffusion_policy.json"), + + }, + disable_env_checker=True, +) gym.register( id="Dev-IK-Rel-Insert-vision-v0", diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/dev_ik_rel_vial_insert_top_down.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/dev_ik_rel_vial_insert_top_down.py new file mode 100644 index 000000000000..6b1ba76ef207 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/dev_ik_rel_vial_insert_top_down.py @@ -0,0 +1,127 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab.controllers.differential_ik_cfg import DifferentialIKControllerCfg +from isaaclab.envs.mdp.actions.actions_cfg import DifferentialInverseKinematicsActionCfg +from isaaclab.utils import configclass +from isaaclab.assets import RigidObjectCfg, ArticulationCfg +from isaaclab.managers import ObservationTermCfg as ObsTerm +from isaaclab.managers import EventTermCfg as EventTerm +from isaaclab.assets import AssetBaseCfg +from isaaclab.sim.spawners.from_files import UsdFileCfg +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, NVIDIA_NUCLEUS_DIR +#from source.isaaclab_assets.isaaclab_assets.robots.universal_robots import UR10_CFG +from isaaclab_assets.robots.universal_robots import UR10_CFG +from . import dev_env_cfg +import isaaclab.sim as sim_utils +from isaaclab.sensors import CameraCfg +import math +from isaaclab.managers import SceneEntityCfg +from isaaclab_assets.glassware.glassware import ChemistryGlassware +from isaaclab_tasks.manager_based.manipulation.cube_lift import mdp +from isaaclab_tasks.manager_based.manipulation.cube_lift.mdp import franka_stack_events + +from isaaclab.managers import TerminationTermCfg as DoneTerm +## +# Pre-defined configs +## +from isaaclab_assets.robots.franka import FRANKA_PANDA_HIGH_PD_CFG # isort: skip +from isaaclab_assets.robots.universal_robots import UR10e_ROBOTIQ_GRIPPER_CFG + +## add some cameras in + + + +@configclass +class FrankaDevEnvCfg(dev_env_cfg.FrankaDevEnvCfg): + def __post_init__(self): + # post init of parent + super().__post_init__() + # put the beaker on the stir plate + glassware = ChemistryGlassware() + self.scene.stirplate = glassware.stirplate(pos=[0.5, 0.0, 0.01]) + #self.scene.scale = glassware.scale(pos=[0.3, -0.3, 0.01]) + self.scene.vialrack = glassware.line_vial_rack(pos=[0.3, -0.3, 0.0], scale=1.5) + self.scene.object = glassware.capped_vial(pos=[0.5, 0.0, 0.01], scale =1.0, name="object") + self.observations.policy.target_object_position = ObsTerm(func=mdp.target_position, params={"object_cfg": SceneEntityCfg("vialrack")}) + + ### subtask + self.observations.subtask_terms.stacked = ObsTerm( + func=mdp.object_stacked, + params={ + "robot_cfg": SceneEntityCfg("robot"), + "upper_object_cfg": SceneEntityCfg("object"), + "lower_object_cfg": SceneEntityCfg("vialrack"), + }, + ) + + self.events.randomise_object_scale= EventTerm( + func=mdp.randomize_rigid_body_scale, + mode="prestartup", + params={ + "scale_range": {"x": (1.0, 1.0), "y": (1.0, 1.0), "z": (1.0, 1.0)}, + "asset_cfg": SceneEntityCfg("object"), + }, + ) + self.events.reset_object_position = EventTerm( + func=mdp.reset_place_root_state_uniform, + mode="reset", + params={ + "pose_range": {"x": (0, 0.2), "y": (0, 0.25), "z": (0.02, 0.02)}, + "velocity_range": {}, + "asset_cfg": SceneEntityCfg("object"), + "asset2_cfg" : SceneEntityCfg("stirplate"), + "asset3_cfg" : SceneEntityCfg("vialrack") + }, + ) + self.scene.robot = FRANKA_PANDA_HIGH_PD_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + #self.scene.robot = UR10e_ROBOTIQ_GRIPPER_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + # self.scene.robot = FRANKA_PANDA_HIGH_PD_CFG.replace( + # prim_path="{ENV_REGEX_NS}/Robot", + # init_state=ArticulationCfg.InitialStateCfg( + # joint_pos={ + + # "panda_joint1": 0.3281, + # "panda_joint2": -0.3684, + # "panda_joint3": -0.2787, + # "panda_joint4": -2.6138, + # "panda_joint5": -2.7527, + # "panda_joint6": 2.4991, # +90° → keeps hand level + # "panda_joint7": 0.3331, + # "panda_finger_joint1": 0.04, # open gripper + # "panda_finger_joint2": 0.04, + # } + # ), + # ) + # replace with relative position controller + self.actions.arm_action = DifferentialInverseKinematicsActionCfg( + asset_name="robot", + joint_names=["panda_joint.*"], + body_name="panda_hand", + controller=DifferentialIKControllerCfg(command_type="pose", use_relative_mode=True, ik_method="dls"), + scale=0.5, + body_offset=DifferentialInverseKinematicsActionCfg.OffsetCfg(pos=[0.0, 0.0, 0.0]), + ) + + self.terminations.success= DoneTerm(func=mdp.object_inserted_upright, params={"lower_object_cfg": SceneEntityCfg("vialrack"), "upright_good_deg": 22.5}) + + #self.terminations.success=DoneTerm(func=mdp.object_stacked_upright, params={"lower_object_cfg": SceneEntityCfg("scale")}) + + self.observations.subtask_terms.appr_goal=ObsTerm(func=mdp.is_object_lifted, params={"threshold":0.15} + + ) + + + +@configclass +class FrankaCubeEnvCfg_PLAY(FrankaDevEnvCfg): + def __post_init__(self): + # post init of parent + super().__post_init__() + # make a smaller scene for play + self.scene.num_envs = 50 + self.scene.env_spacing = 2.5 + # disable randomization for play + self.observations.policy.enable_corruption = False diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/mdp/observations.py index a157efab49ce..e9a0793b5cfc 100755 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/mdp/observations.py @@ -162,14 +162,14 @@ def is_object_lifted( def ee_frame_pos(env: ManagerBasedRLEnv, ee_frame_cfg: SceneEntityCfg = SceneEntityCfg("ee_frame")) -> torch.Tensor: ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] ee_frame_pos = ee_frame.data.target_pos_w[:, 0, :] - env.scene.env_origins[:, 0:3] - # print(f"EE pos : {ee_frame_pos}") + print(f"EE pos : {ee_frame_pos}") return ee_frame_pos ## def ee_frame_quat(env: ManagerBasedRLEnv, ee_frame_cfg: SceneEntityCfg = SceneEntityCfg("ee_frame")) -> torch.Tensor: ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] ee_frame_quat = ee_frame.data.target_quat_w[:, 0, :] - #print(f"EE quat : {ee_frame_quat}") + print(f"EE quat : {ee_frame_quat}") return ee_frame_quat