diff --git a/.vscode/browse.vc.db b/.vscode/browse.vc.db new file mode 100644 index 000000000000..5b28fda5fa58 Binary files /dev/null and b/.vscode/browse.vc.db differ diff --git a/.vscode/browse.vc.db-shm b/.vscode/browse.vc.db-shm new file mode 100644 index 000000000000..25a1415fef5f Binary files /dev/null and b/.vscode/browse.vc.db-shm differ diff --git a/.vscode/browse.vc.db-wal b/.vscode/browse.vc.db-wal new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/docker/Dockerfile.base b/docker/Dockerfile.base index 53f528244e25..b87557035f3f 100644 --- a/docker/Dockerfile.base +++ b/docker/Dockerfile.base @@ -77,14 +77,16 @@ RUN mkdir -p ${ISAACSIM_ROOT_PATH}/kit/cache && \ mkdir -p ${DOCKER_USER_HOME}/Documents # for singularity usage, create NVIDIA binary placeholders -RUN touch /bin/nvidia-smi && \ - touch /bin/nvidia-debugdump && \ - touch /bin/nvidia-persistenced && \ - touch /bin/nvidia-cuda-mps-control && \ - touch /bin/nvidia-cuda-mps-server && \ - touch /etc/localtime && \ - mkdir -p /var/run/nvidia-persistenced && \ - touch /var/run/nvidia-persistenced/socket +# NOTE: These are only needed for Singularity, not Docker +# For Docker, comment these out as they override the real nvidia binaries +# RUN touch /bin/nvidia-smi && \ +# touch /bin/nvidia-debugdump && \ +# touch /bin/nvidia-persistenced && \ +# touch /bin/nvidia-cuda-mps-control && \ +# touch /bin/nvidia-cuda-mps-server && \ +# touch /etc/localtime && \ +# mkdir -p /var/run/nvidia-persistenced && \ +# touch /var/run/nvidia-persistenced/socket # installing Isaac Lab dependencies # use pip caching to avoid reinstalling large packages diff --git a/docs/find_best_checkpoints.py b/docs/find_best_checkpoints.py new file mode 100644 index 000000000000..bdda475fba0a --- /dev/null +++ b/docs/find_best_checkpoints.py @@ -0,0 +1,160 @@ +#!/usr/bin/env python3 +""" +Script to find and copy the best validation model checkpoint for each model +in the ensemble training runs. + +The script searches through each model directory in Dev-IK-Rel-Place-v0, +finds the checkpoint with the lowest validation loss, and copies it to a +centralized 'best_models' folder. +""" + +import os +import re +import shutil +from pathlib import Path +from typing import Optional, Tuple + + +def parse_validation_loss(filename: str) -> Optional[float]: + """ + Extract validation loss from checkpoint filename. + + Expected format: model_epoch_XXX_best_validation_YYY.YYY.pth + + Args: + filename: The checkpoint filename + + Returns: + The validation loss as a float, or None if parsing fails + """ + pattern = r'model_epoch_\d+_best_validation_([\d.]+)\.pth' + match = re.search(pattern, filename) + if match: + return float(match.group(1)) + return None + + +def find_best_checkpoint(model_dir: Path) -> Optional[Tuple[Path, float]]: + """ + Find the checkpoint with the lowest validation loss in a model directory. + + Args: + model_dir: Path to the model directory (e.g., model0, model1, etc.) + + Returns: + Tuple of (checkpoint_path, validation_loss) or None if no checkpoints found + """ + # Search for all checkpoint files in the models subdirectory + checkpoint_pattern = "model_epoch_*_best_validation_*.pth" + checkpoints = list(model_dir.rglob(checkpoint_pattern)) + + if not checkpoints: + print(f" āš ļø No checkpoints found in {model_dir.name}") + return None + + best_checkpoint = None + best_loss = float('inf') + + for checkpoint in checkpoints: + loss = parse_validation_loss(checkpoint.name) + if loss is not None and loss < best_loss: + best_loss = loss + best_checkpoint = checkpoint + + if best_checkpoint: + return (best_checkpoint, best_loss) + + print(f" āš ļø Could not parse validation loss from checkpoints in {model_dir.name}") + return None + + +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 / "place/Dev-IK-Rel-Place-v0" + output_dir = base_dir / "best_models" + + # Create output directory (including parent directories if needed) + output_dir.mkdir(parents=True, exist_ok=True) + print(f"šŸ“ Output directory: {output_dir}\n") + + # Find all model directories (model0, model1, ..., modelN) + model_dirs = sorted([d for d in base_dir.iterdir() + if d.is_dir() and d.name.startswith("model") + and d.name != "best_models"]) + + if not model_dirs: + print("āŒ No model directories found!") + return + + print(f"Found {len(model_dirs)} model directories\n") + print("=" * 80) + + # Process each model directory + results = [] + for model_dir in model_dirs: + print(f"\nšŸ” Processing {model_dir.name}...") + + result = find_best_checkpoint(model_dir) + if result: + checkpoint_path, validation_loss = result + print(f" āœ… Best checkpoint: {checkpoint_path.name}") + print(f" šŸ“Š Validation loss: {validation_loss}") + + # Copy to output directory with model name prefix + output_filename = f"{model_dir.name}_{checkpoint_path.name}" + output_path = output_dir / output_filename + + shutil.copy2(checkpoint_path, output_path) + print(f" šŸ“‹ Copied to: {output_filename}") + + results.append({ + 'model': model_dir.name, + 'checkpoint': checkpoint_path.name, + 'loss': validation_loss, + 'output': output_filename + }) + + # Print summary + print("\n" + "=" * 80) + print("\nSUMMARY") + print("=" * 80) + print(f"\nProcessed {len(model_dirs)} models") + print(f"Successfully copied {len(results)} best checkpoints\n") + + if results: + print("Best checkpoints by validation loss:") + print("-" * 80) + sorted_results = sorted(results, key=lambda x: x['loss']) + for i, r in enumerate(sorted_results, 1): + print(f"{i:2d}. {r['model']:8s} | Loss: {r['loss']:12.2f} | {r['checkpoint']}") + + # Save file paths to text file + save_paths_to_file(output_dir, results) + + print(f"\nAll best models saved to: {output_dir}") + else: + print("No checkpoints were copied") + + +def save_paths_to_file(output_dir: Path, results: list) -> None: + """ + Save all copied checkpoint file paths to a text file. + + Args: + output_dir: Directory where the checkpoints were copied + results: List of dictionaries containing model information + """ + txt_file = output_dir / "best_model_paths.txt" + + with open(txt_file, 'w') as f: + for result in sorted(results, key=lambda x: x['model']): + file_path = output_dir / result['output'] + f.write(f"{file_path}\n") + + print(f"šŸ“ File paths saved to: {txt_file}") + + +if __name__ == "__main__": + main() diff --git a/docs/run_experiments.sh b/docs/run_experiments.sh new file mode 100644 index 000000000000..9f7da360e12a --- /dev/null +++ b/docs/run_experiments.sh @@ -0,0 +1,66 @@ + + +## script to run experiments in the background + +task="Dev-IK-Rel-Place-v0" +horizon=5000 +num_rollouts=100 +run_file="scripts/imitation_learning/robomimic/play_ensemble_v05.py" +ensemble_size=15 + +##### this now runs the experiments + +## ensemble 15 +seed=101 +exp_name="${task}_ensemble_${ensemble_size}_seed_${seed}_highdem" +./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}_highdem" +./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}_highdem" +./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}_highdem" +./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}_highdem" +./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}_highdem" +./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}_highdem" +./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}_highdem" +./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}_highdem" +./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 1 +ensemble_size=1 +seed=101 +exp_name="${task}_ensemble_${ensemble_size}_seed_${seed}_highdem" +./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}_highdem" +./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}_highdem" +./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/isaaclab_mimic/annotate_demos.py b/scripts/imitation_learning/isaaclab_mimic/annotate_demos.py index 17322c6e93ce..489a52bbc976 100644 --- a/scripts/imitation_learning/isaaclab_mimic/annotate_demos.py +++ b/scripts/imitation_learning/isaaclab_mimic/annotate_demos.py @@ -210,6 +210,10 @@ def main(): # Disable all termination terms env_cfg.terminations = None + # Note: We keep events enabled as the reset_to() function applies state AFTER events run, + # so the recorded state should override any randomization. If object placement is still + # inaccurate, try using --device cuda:0 instead of --device cpu. + # Set up recorder terms for mimic annotations env_cfg.recorders = MimicRecorderManagerCfg() if not args_cli.auto: diff --git a/scripts/imitation_learning/robomimic/play.py b/scripts/imitation_learning/robomimic/play.py index 4b1476f6bea1..02d293b21eba 100644 --- a/scripts/imitation_learning/robomimic/play.py +++ b/scripts/imitation_learning/robomimic/play.py @@ -178,6 +178,8 @@ def main(): results.append(terminated) print(f"[INFO] Trial {trial}: {terminated}\n") + print(f"\n Environment summary : {env}") + # Print results print(f"\nSuccessful trials: {results.count(True)}, out of {len(results)} trials") print(f"Success rate: {results.count(True) / len(results)}") print(f"Trial Results: {results}\n") diff --git a/scripts/imitation_learning/robomimic/play_diffusion.py b/scripts/imitation_learning/robomimic/play_diffusion.py new file mode 100644 index 000000000000..15d78ca345f7 --- /dev/null +++ b/scripts/imitation_learning/robomimic/play_diffusion.py @@ -0,0 +1,297 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Script to play and evaluate a trained Diffusion Policy from robomimic. + +This script loads a robomimic diffusion policy and plays it in an Isaac Lab environment. +It handles the temporal observation stacking required by diffusion policy. + +Args: + task: Name of the environment. + checkpoint: Path to the robomimic policy checkpoint. + horizon: If provided, override the step horizon of each rollout. + num_rollouts: If provided, override the number of rollouts. + seed: If provided, override the default random seed. + norm_factor_min: If provided, minimum value of the action space normalization factor. + norm_factor_max: If provided, maximum value of the action space normalization factor. +""" + +"""Launch Isaac Sim Simulator first.""" + + +import argparse + +from isaaclab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Evaluate robomimic diffusion policy for Isaac Lab environment.") +parser.add_argument( + "--disable_fabric", action="store_true", default=False, help="Disable fabric and use USD I/O operations." +) +parser.add_argument("--task", type=str, default=None, help="Name of the task.") +parser.add_argument("--checkpoint", type=str, default=None, help="Pytorch model checkpoint to load.") +parser.add_argument("--horizon", type=int, default=800, help="Step horizon of each rollout.") +parser.add_argument("--num_rollouts", type=int, default=1, help="Number of rollouts.") +parser.add_argument("--seed", type=int, default=101, help="Random seed.") +parser.add_argument( + "--norm_factor_min", type=float, default=None, help="Optional: minimum value of the normalization factor." +) +parser.add_argument( + "--norm_factor_max", type=float, default=None, help="Optional: maximum value of the normalization factor." +) +parser.add_argument("--enable_pinocchio", default=False, action="store_true", help="Enable Pinocchio.") + + +# append AppLauncher cli args +AppLauncher.add_app_launcher_args(parser) +# parse the arguments +args_cli = parser.parse_args() + +if args_cli.enable_pinocchio: + # Import pinocchio before AppLauncher to force the use of the version installed by IsaacLab and not the one installed by Isaac Sim + # pinocchio is required by the Pink IK controllers and the GR1T2 retargeter + import pinocchio # noqa: F401 + +# launch omniverse app +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import copy +from collections import deque +import gymnasium as gym +import numpy as np +import random +import torch + +import robomimic.utils.file_utils as FileUtils +import robomimic.utils.torch_utils as TorchUtils + +if args_cli.enable_pinocchio: + import isaaclab_tasks.manager_based.manipulation.pick_place # noqa: F401 + +from isaaclab_tasks.utils import parse_env_cfg + + +def get_observation_horizon(policy): + """Get the observation horizon from diffusion policy config. + + Args: + policy: The loaded robomimic policy wrapper. + + Returns: + int: The observation horizon (default 2 for diffusion policy). + """ + if hasattr(policy.policy, 'algo_config') and hasattr(policy.policy.algo_config, 'horizon'): + return policy.policy.algo_config.horizon.observation_horizon + # Default for diffusion policy if not found + return 2 + + +def prepare_obs_for_diffusion(obs_dict, obs_history): + """Prepare observations for diffusion policy inference. + + Diffusion policy expects observations with shape [B, T, D] where: + - B = batch size (1 for single env inference) + - T = observation_horizon (temporal dimension) + - D = observation feature dimension + + Args: + obs_dict: Current observation dictionary from environment + obs_history: Deque of past observations + + Returns: + Stacked observation dictionary with temporal dimension + """ + # Store current observation (without batch dimension) + current_obs = {} + for k in obs_dict: + if obs_dict[k].dim() > 1: + current_obs[k] = obs_dict[k].squeeze(0) + else: + current_obs[k] = obs_dict[k] + + obs_history.append(current_obs) + + # Stack observations along temporal dimension: [T, D] -> [1, T, D] + stacked_obs = {} + for k in obs_dict: + obs_list = [obs_history[i][k] for i in range(len(obs_history))] + stacked = torch.stack(obs_list, dim=0).unsqueeze(0) # [1, T, D] + stacked_obs[k] = stacked + + # Debug: print shapes on first call + if not hasattr(prepare_obs_for_diffusion, '_printed'): + print("\n[DEBUG] Observation shapes after stacking:") + for k, v in stacked_obs.items(): + print(f" {k}: {v.shape} (ndim={v.ndim})") + prepare_obs_for_diffusion._printed = True + + return stacked_obs + + +def rollout(policy, env, success_term, horizon, device): + """Perform a single rollout of the diffusion policy in the environment. + + Args: + policy: The robomimic diffusion policy to play. + env: The environment to play in. + horizon: The step horizon of each rollout. + device: The device to run the policy on. + + Returns: + terminated: Whether the rollout terminated successfully. + traj: The trajectory of the rollout. + """ + policy.start_episode() + obs_dict, _ = env.reset() + traj = dict(actions=[], obs=[], next_obs=[]) + + # Setup temporal observation history for diffusion policy + observation_horizon = get_observation_horizon(policy) + obs_history = deque(maxlen=observation_horizon) + + # Initialize history with first observation (repeated to fill the queue) + first_obs = {} + for k in obs_dict["policy"]: + if obs_dict["policy"][k].dim() > 1: + first_obs[k] = obs_dict["policy"][k].squeeze(0) + else: + first_obs[k] = obs_dict["policy"][k] + + # Fill history with copies of first observation + for _ in range(observation_horizon): + obs_history.append(copy.deepcopy(first_obs)) + + for i in range(horizon): + # Prepare observations with temporal stacking for diffusion policy + obs = prepare_obs_for_diffusion(obs_dict["policy"], obs_history) + + # Debug: print gripper observation at key steps + # if i == 0 or i == 5: + # print(f"\n[DEBUG] Step {i} - Observations:") + # print(f" gripper_pos: {obs['gripper_pos'][0, -1, :].cpu().numpy()}") # Last timestep + # print(f" eef_pos: {obs['eef_pos'][0, -1, :].cpu().numpy()}") + # print(f" object_position: {obs['object_position'][0, -1, :].cpu().numpy()}") + + # Check if environment has image observations + if hasattr(env.cfg, "image_obs_list"): + # Process image observations for robomimic inference + for image_name in env.cfg.image_obs_list: + if image_name in obs_dict["policy"].keys(): + # Convert from chw uint8 to hwc normalized float + image = torch.squeeze(obs_dict["policy"][image_name]) + image = image.permute(2, 0, 1).clone().float() + image = image / 255.0 + image = image.clip(0.0, 1.0) + # Stack image observations temporally as well + obs[image_name] = image.unsqueeze(0).unsqueeze(0).expand(1, observation_horizon, -1, -1, -1) + + traj["obs"].append(obs) + + # Compute actions from diffusion policy + # Use batched_ob=True because we already have batch dimension [1, T, D] + actions = policy(obs, batched_ob=True) + + # Debug: print action stats periodically + # if i == 0 or i == 10 or i == 50: + # print(f"\n[DEBUG] Step {i} - Raw policy actions:") + # print(f" Shape: {actions.shape}") + # print(f" Min: {actions.min():.4f}, Max: {actions.max():.4f}") + # print(f" Values: {actions}") + # # Check if policy has normalization stats + # if hasattr(policy, 'action_normalization_stats') and policy.action_normalization_stats is not None: + # print(f" Action norm stats available: {policy.action_normalization_stats.keys()}") + # else: + # print(f" WARNING: No action normalization stats in policy!") + + # Unnormalize actions if normalization factors provided + if args_cli.norm_factor_min is not None and args_cli.norm_factor_max is not None: + actions = ( + (actions + 1) * (args_cli.norm_factor_max - args_cli.norm_factor_min) + ) / 2 + args_cli.norm_factor_min + + actions = torch.from_numpy(actions).to(device=device).view(1, env.action_space.shape[1]) + + # Apply actions to environment + obs_dict, _, terminated, truncated, _ = env.step(actions) + + # Record trajectory + traj["actions"].append(actions.tolist()) + traj["next_obs"].append(obs_dict["policy"]) + + # Check if rollout was successful + if bool(success_term.func(env, **success_term.params)[0]): + return True, traj + elif terminated or truncated: + return False, traj + + return False, traj + + +def main(): + """Run a trained diffusion policy from robomimic with Isaac Lab environment.""" + # parse configuration + env_cfg = parse_env_cfg(args_cli.task, device=args_cli.device, num_envs=1, use_fabric=not args_cli.disable_fabric) + + # Set observations to dictionary mode for Robomimic + env_cfg.observations.policy.concatenate_terms = False + + # Set termination conditions + env_cfg.terminations.time_out = None + + # Disable recorder + env_cfg.recorders = None + + # Extract success checking function + success_term = env_cfg.terminations.success + env_cfg.terminations.success = None + + # Create environment + env = gym.make(args_cli.task, cfg=env_cfg).unwrapped + + # Set seed for reproducibility + torch.manual_seed(args_cli.seed) + np.random.seed(args_cli.seed) + random.seed(args_cli.seed) + env.seed(args_cli.seed) + + # Acquire device + device = TorchUtils.get_torch_device(try_to_use_cuda=True) + + # Load policy once (diffusion policy maintains internal action queue) + policy, _ = FileUtils.policy_from_checkpoint(ckpt_path=args_cli.checkpoint, device=device) + print(f"[INFO] Loaded diffusion policy from {args_cli.checkpoint}") + + # Debug: print expected observation shapes from policy + print("\n[DEBUG] Policy expected obs_shapes:") + for k, v in policy.policy.obs_shapes.items(): + print(f" {k}: {v} (len={len(v)})") + + # Get and print observation horizon + obs_horizon = get_observation_horizon(policy) + print(f"[INFO] Diffusion policy observation_horizon: {obs_horizon}") + + # Run policy rollouts + results = [] + for trial in range(args_cli.num_rollouts): + print(f"[INFO] Starting trial {trial}") + terminated, traj = rollout(policy, env, success_term, args_cli.horizon, device) + results.append(terminated) + print(f"[INFO] Trial {trial}: {'Success' if terminated else 'Failed'}\n") + + print(f"\nSuccessful trials: {results.count(True)}, out of {len(results)} trials") + print(f"Success rate: {results.count(True) / len(results):.2%}") + print(f"Trial Results: {results}\n") + + env.close() + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/scripts/imitation_learning/robomimic/play_ensemble_v05.py b/scripts/imitation_learning/robomimic/play_ensemble_v05.py new file mode 100644 index 000000000000..43771bc8d2ac --- /dev/null +++ b/scripts/imitation_learning/robomimic/play_ensemble_v05.py @@ -0,0 +1,781 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# run this script: +# ./isaaclab.sh -p scripts/imitation_learning/robomimic/play_ensemble_v05.py --device cuda --task Dev-IK-Rel-v0 --num_rollouts 100 +# stack cube task: ./isaaclab.sh -p scripts/imitation_learning/robomimic/play.py --device cuda --task Isaac-Stack-Cube-Franka-IK-Rel-v0 --num_rollouts 50 +"""Script to play and evaluate a trained policy from robomimic. + +This script loads a robomimic policy and plays it in an Isaac Lab environment. +Updated for robomimic v0.5 compatibility. + +Args: + task: Name of the environment. + checkpoint: Path to the robomimic policy checkpoint. + horizon: If provided, override the step horizon of each rollout. + num_rollouts: If provided, override the number of rollouts. + seed: If provided, overeride the default random seed. + norm_factor_min: If provided, minimum value of the action space normalization factor. + norm_factor_max: If provided, maximum value of the action space normalization factor. +""" + +"""Launch Isaac Sim Simulator first.""" + + +import argparse + +from isaaclab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Evaluate robomimic policy for Isaac Lab environment.") +parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.") +parser.add_argument("--video_length", type=int, default=800, help="Length of the recorded video (in steps).") +parser.add_argument( + "--disable_fabric", action="store_true", default=False, help="Disable fabric and use USD I/O operations." +) +parser.add_argument("--task", type=str, default=None, help="Name of the task.") +parser.add_argument("--checkpoint", type=str, default=None, help="Pytorch model checkpoint to load.") +parser.add_argument("--horizon", type=int, default=2000, help="Step horizon of each rollout.") +parser.add_argument("--num_rollouts", type=int, default=1, help="Number of rollouts.") +parser.add_argument("--seed", type=int, default=101, help="Random seed.") +parser.add_argument( + "--norm_factor_min", type=float, default=None, help="Optional: minimum value of the normalization factor." +) +parser.add_argument( + "--norm_factor_max", type=float, default=None, help="Optional: maximum value of the normalization factor." +) +parser.add_argument("--enable_pinocchio", default=False, action="store_true", help="Enable Pinocchio.") + +parser.add_argument("--model_name", type=str, default=None) +parser.add_argument("--exp_name", type=str, default=None) + +group = parser.add_mutually_exclusive_group() +group.add_argument("--use_recovery", type=bool ,default= False, help="Enable recovery mechanism. By default recovery is enabled.") + +parser.add_argument("--ensemble_size", type=int, default=10) + +# append AppLauncher cli args +AppLauncher.add_app_launcher_args(parser) +# parse the arguments +args_cli = parser.parse_args() + +if args_cli.enable_pinocchio: + # Import pinocchio before AppLauncher to force the use of the version installed by IsaacLab and not the one installed by Isaac Sim + # pinocchio is required by the Pink IK controllers and the GR1T2 retargeter + import pinocchio # noqa: F401 +import numpy as np +# launch omniverse app +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import time +import copy +import random +import os +import gymnasium as gym +import torch +import math +from collections import deque + +import robomimic.utils.file_utils as FileUtils +import robomimic.utils.torch_utils as TorchUtils +from isaaclab_tasks.manager_based.manipulation.cube_lift.mdp.observations import get_joint_pos, object_grasped +#from chills.tasks.mdp.observations import get_joint_pos, object_grasped +if args_cli.enable_pinocchio: + import isaaclab_tasks.manager_based.manipulation.pick_place # noqa: F401 + +from isaaclab_tasks.utils import parse_env_cfg +#from isaaclab.utils.logging_helper import LoggingHelper, ErrorType, LogType +from isaaclab.utils.math import matrix_from_quat, convert_quat, quat_mul, euler_xyz_from_quat +from evaluation import ensemble_uncertainty +from isaaclab.safety.switchingLogic import SwitchingLogic +from isaaclab.controllers import DifferentialIKController, DifferentialIKControllerCfg +from isaaclab.utils.math import subtract_frame_transforms, euler_xyz_from_quat +from isaaclab.managers import SceneEntityCfg +from isaaclab.envs.mdp.actions.actions_cfg import DifferentialInverseKinematicsActionCfg, BinaryJointPositionActionCfg +from isaaclab.envs.mdp.actions.binary_joint_actions import BinaryJointPositionAction +from backup_controller_handler import BackupController + + +def policy_from_checkpoint_v05(ckpt_path, device, verbose=True): + """ + Load a policy from a checkpoint file - compatible with robomimic v0.5. + + This is a wrapper that handles API differences between robomimic versions. + + Args: + ckpt_path: Path to the checkpoint file. + device: Device to load the policy on. + verbose: Whether to print verbose output. + + Returns: + policy: The loaded policy. + ckpt_dict: The checkpoint dictionary. + """ + import robomimic + from packaging import version + + # Check robomimic version + robomimic_version = getattr(robomimic, '__version__', '0.3.0') + + if verbose: + print(f"[INFO] Loading checkpoint with robomimic version {robomimic_version}") + + # robomimic v0.5+ may have different API + if version.parse(robomimic_version) >= version.parse('0.5.0'): + # v0.5 API - policy_from_checkpoint returns (policy, ckpt_dict) + try: + result = FileUtils.policy_from_checkpoint(ckpt_path=ckpt_path, device=device, verbose=verbose) + if isinstance(result, tuple): + return result + else: + return result, None + except TypeError as e: + if verbose: + print(f"[WARN] policy_from_checkpoint failed with: {e}") + print("[INFO] Trying alternative loading method...") + + # Alternative: manual loading + import torch + import json + from robomimic.algo import algo_factory + from robomimic.config import config_factory + + ckpt_dict = torch.load(ckpt_path, map_location=device) + + # Get config from checkpoint + config = ckpt_dict.get('config', {}) + if isinstance(config, dict): + algo_name = config.get('algo_name', 'bc') + config_obj = config_factory(algo_name) + with config_obj.values_unlocked(): + config_obj.update(config) + config = config_obj + + # Create policy + policy = algo_factory( + algo_name=config.algo_name, + config=config, + obs_key_shapes=ckpt_dict.get('shape_metadata', {}).get('all_shapes', {}), + ac_dim=ckpt_dict.get('shape_metadata', {}).get('ac_dim', 7), + device=device, + ) + + # Load weights + policy.deserialize(ckpt_dict['model']) + policy.set_eval() + + return policy, ckpt_dict + else: + # Older API + return FileUtils.policy_from_checkpoint(ckpt_path=ckpt_path, device=device, verbose=verbose) + + +def rollout(policy, env, success_term, horizon, device): + """Perform a single rollout of the policy in the environment. + + Args: + policy: The robomimicpolicy to play. + env: The environment to play in. + horizon: The step horizon of each rollout. + device: The device to run the policy on. + + Returns: + terminated: Whether the rollout terminated. + traj: The trajectory of the rollout. + """ + policy.start_episode() + obs_dict, _ = env.reset() + + #print(f"obs_dict type: {type(obs_dict)},\nobs_dict contents:\n {obs_dict}") + + traj = dict(actions=[], obs=[], next_obs=[], sub_obs=[], uncertainties=[]) + + # Define the observation keys the trained model expects + TRAINED_OBS_KEYS = ['eef_pos', 'eef_quat', 'gripper_pos', 'joint_pos', 'joint_vel', + 'object_position', 'target_object_position', 'actions'] + + for i in range(horizon): + # Prepare observations + obs = copy.deepcopy(obs_dict["policy"]) + sub_obs = copy.deepcopy(obs_dict["subtask_terms"]) + + for ob in obs: + obs[ob] = torch.squeeze(obs[ob]) + #print(f"found observation : {obs[ob]}") + + for subob in sub_obs: + sub_obs[subob] = torch.squeeze(sub_obs[subob]) + + # Check if environment image observations + if hasattr(env.cfg, "image_obs_list"): + # Process image observations for robomimic inference + for image_name in env.cfg.image_obs_list: + if image_name in obs_dict["policy"].keys(): + # Convert from chw uint8 to hwc normalized float + image = torch.squeeze(obs_dict["policy"][image_name]) + image = image.permute(2, 0, 1).clone().float() + image = image / 255.0 + image = image.clip(0.0, 1.0) + obs[image_name] = image + + traj["obs"].append(obs) + traj["sub_obs"].append(sub_obs) + + # Filter observations to only include keys the model was trained with + policy_obs = {k: obs[k] for k in TRAINED_OBS_KEYS if k in obs} + + # Add dropout layers to the model and calculate uncertainty and remove at the end to not effect final action + + #hooks = inject_dropout_layers(policy=policy, probability=0.1) + # uncertainty_dict = MC_dropout_uncertainty(policy=policy, obs=obs, niters=15) + # traj['uncertainties'].append(uncertainty_dict['variance']) + #remove_dropout_layers(hooks) + + # Compute actions + actions = policy(policy_obs) + + # Unnormalize actions + if args_cli.norm_factor_min is not None and args_cli.norm_factor_max is not None: + actions = ( + (actions + 1) * (args_cli.norm_factor_max - args_cli.norm_factor_min) + ) / 2 + args_cli.norm_factor_min + + actions = torch.from_numpy(actions).to(device=device).view(1, env.action_space.shape[1]) + + # Apply actions + obs_dict, _, terminated, truncated, _ = env.step(actions) + obs = obs_dict["policy"] + sub_obs = obs_dict["subtask_terms"] + + # Record trajectory + traj["actions"].append(actions.tolist()) + traj["next_obs"].append(obs) + + + # Check if rollout was successful + if bool(success_term.func(env, **success_term.params)[0]): + return True, traj + elif terminated or truncated: + return False, traj + + return False, traj + + +# tip over detection +def has_fallen(asset, threshold: float = -0.7, debug: bool = True) -> bool: + """ + Detect if an IsaacLab asset has fallen over by checking its up vector. + + Args: + asset: IsaacLab asset (e.g. robot, object) with .get_world_pose() + threshold: minimum z-component of the up vector (1.0 = perfectly upright). + Typical cutoff ~0.7 (ā‰ˆ45Ā° tilt). + debug: if True, prints internal values for inspection. + + Returns: + True if fallen, False otherwise. + """ + quat = asset.data.body_link_quat_w # (pos, quat), usually torch.Tensors + # print("Beaker quat : ", quat) + # Ensure torch tensor, shape (4,) + if isinstance(quat, (list, tuple)): + quat = torch.tensor(quat, dtype=torch.float32) + quat = quat.view(-1, 4) # shape (1,4) + + # Convert quaternion to wxyz ordering + quat_wxyz = convert_quat(quat, to="wxyz") + + # Rotation matrix (1,3,3) + R = matrix_from_quat(quat_wxyz) + + # World up vector is the third column of rotation matrix + up = R[0, :, 2] # shape (3,) + uprightness = float(up[2]) + + if debug: + print(f"[has_fallen] quat={quat_wxyz.cpu().numpy().round(4)}") + print(f"[has_fallen] up={up.cpu().numpy().round(4)} uprightness={uprightness:.3f}") + + return uprightness > threshold + + +def rollout_ensemble(ensemble, env, success_term, horizon, device, parameters, use_recovery=True, rollout_num=0): + """Perform a single rollout of the policy in the environment. + + Args: + policy: The robomimicpolicy to play. + env: The environment to play in. + horizon: The step horizon of each rollout. + device: The device to run the policy on. + + Returns: + terminated: Whether the rollout terminated. + traj: The trajectory of the rollout. + """ + ## do setup + print("running rollout number : ", rollout_num) + for policy in ensemble: + policy.start_episode() + obs_dict, _ = env.reset() + traj = dict(actions=[], obs=[], next_obs=[], sub_obs=[], + uncertainties=[], min_actions=[], max_actions=[], + time_taken=[], joint_pos=[], recovery=[], failure=[], grasp=[], appr=[]) + + ###### SETUP SWITCHING LOGIC #### + switchingLogic = SwitchingLogic(parameters, horizon=horizon) + certain_joint_positions = [] + + ###### SET UP RECOVERY #### + use_recovery=False + recovery_activated_during_rollout = 0 + print("rollout recovery enabled ? : ", use_recovery) + + ##### RECOVERY CONFIG #### + recovery_cooldown_duration = 300 # timesteps + recovery_cooldown_timer = recovery_cooldown_duration + recovery_cooldown_active = False + recovery = [] + recovery_steps = 0 + max_recovery_steps = 500 # Safety limit for recovery mode + + ##### CONFIG RECOVERY CONTROLLER #### + backup_controller = BackupController(env, device) + state_guess = 0 + last_state = 0 + recovery_mode = False + object_knocked = False + + # Define the observation keys the trained model expects + # These must match the keys used during training (from shape_metadata) + TRAINED_OBS_KEYS = ['eef_pos', 'eef_quat', 'gripper_pos', 'joint_pos', 'joint_vel', + 'object_position', 'target_object_position', 'actions'] + + print(f'Running for horizon {horizon}') + for i in range(horizon): + # Prepare observations + + obs = copy.deepcopy(obs_dict["policy"]) + sub_obs = copy.deepcopy(obs_dict["subtask_terms"]) + + for ob in obs: + obs[ob] = torch.squeeze(obs[ob]) + + for subob in sub_obs: + sub_obs[subob] = torch.squeeze(sub_obs[subob]) + traj["obs"].append(obs) + traj["sub_obs"].append(sub_obs) + + # Check if environment image observations + if hasattr(env.cfg, "image_obs_list"): + # Process image observations for robomimic inference + for image_name in env.cfg.image_obs_list: + if image_name in obs_dict["policy"].keys(): + # Convert from chw uint8 to hwc normalized float + image = torch.squeeze(obs_dict["policy"][image_name]) + image = image.permute(2, 0, 1).clone().float() + image = image / 255.0 + image = image.clip(0.0, 1.0) + obs[image_name] = image + + # Filter observations to only include keys the model was trained with + # This prevents dimension mismatch errors when env has extra observations + policy_obs = {k: obs[k] for k in TRAINED_OBS_KEYS if k in obs} + + # Debug: print observation shapes on first iteration + if i == 0: + print(f"[DEBUG] Available obs keys: {list(obs.keys())}") + print(f"[DEBUG] Filtered policy_obs keys: {list(policy_obs.keys())}") + total_dim = 0 + for k, v in policy_obs.items(): + dim = v.shape[-1] if len(v.shape) > 0 else 1 + print(f"[DEBUG] {k}: shape={v.shape}, dim={dim}") + total_dim += dim + print(f"[DEBUG] Total dimension: {total_dim} (expected 44)") + missing = [k for k in TRAINED_OBS_KEYS if k not in obs] + if missing: + print(f"[DEBUG] MISSING keys: {missing}") + + #### get policy action #### + metrics = ensemble_uncertainty(ensemble, policy_obs) + uncertainty = metrics['variance'] + policy_actions = metrics['mean'] + + # Robust state estimation based on actual observations + is_grasping = bool(sub_obs['grasp'].item()) if sub_obs['grasp'].numel() == 1 else bool(torch.any(sub_obs['grasp'])) + + # Update state guess based on actual gripper/object state + if is_grasping: + # Object is grasped - should be in LIFT or later states + if state_guess < 4: # LIFT_OBJECT = 4 + state_guess = 4 + print(f"[STATE] Updated to LIFT_OBJECT (grasp detected)") + else: + # Not grasping - if we thought we were lifting, we dropped it + if state_guess >= 4 and state_guess < 7: # Was in lift/move states + # Object dropped - need to re-approach + state_guess = 0 # Go back to REST to safely re-approach + print(f"[STATE] Reset to REST (grasp lost)") + + sm_actions, state_guess = backup_controller.get_controller_action(state_guess, 0) + if last_state != state_guess: + print(f"[STATE] Transition: {last_state} -> {state_guess}") + last_state = state_guess + + + # this is the switching logic + uncertainty_triggered = switchingLogic.check(uncertainty, i) + if uncertainty_triggered: + # uncertainty triggered + if use_recovery and not recovery_cooldown_active: + # if we have switched to recovery + if not recovery_mode: + recovery_mode = True + recovery_activated_during_rollout+=1 + recovery_steps = 0 + print(f"[RECOVERY] Step {i} - Starting recovery from state {state_guess}") + + if recovery_mode: + actions = sm_actions + recovery_steps += 1 + + # Exit recovery when task completes or after max steps + max_recovery_steps = 500 + if state_guess == 7 or recovery_steps > max_recovery_steps: # APPROACH_GOAL = 7 + print(f"[RECOVERY] Completed after {recovery_steps} steps, state={state_guess}") + recovery_mode = False + recovery_cooldown_active = True + recovery_cooldown_timer = recovery_cooldown_duration + else: + # Decrement cooldown when not in recovery + if recovery_cooldown_active: + recovery_cooldown_timer -= 1 + if recovery_cooldown_timer <= 0: + recovery_cooldown_active = False + recovery_cooldown_timer = recovery_cooldown_duration + print("Recovery Cooldown Ended") + + actions = policy_actions + + # Unnormalize actions + if args_cli.norm_factor_min is not None and args_cli.norm_factor_max is not None: + actions = ( + (actions + 1) * (args_cli.norm_factor_max - args_cli.norm_factor_min) + ) / 2 + args_cli.norm_factor_min + actions = actions.to(device=device).view(1, env.action_space.shape[1]) + + + # Apply actions + obs_dict, _, terminated, truncated, _ = env.step(actions) + recovery.append(recovery_mode) + # update obs from taking action + obs = obs_dict["policy"] + sub_obs = obs_dict["subtask_terms"] + + # Record trajectory + traj["actions"].append(actions.tolist()) + traj["next_obs"].append(obs) + traj['uncertainties'].append(uncertainty) + traj['max_actions'].append(metrics['max']) + traj['min_actions'].append(metrics['min']) + traj['time_taken'].append(metrics['time_taken']) + traj['joint_pos'].append(obs['abs_joint_pos']) + traj['recovery'].append(recovery) + if torch.any(obs['object_knocked']): + print("[MONITOR] object knocked over") + object_knocked = True + + # Check if rollout was successful + if bool(success_term.func(env, **success_term.params)[0]): + traj['failure'].append(" ") + grasp = torch.any(sub_obs['grasp']) + traj['grasp'].append(grasp.item()) + appr = torch.any(sub_obs['stacked']) + traj['appr'].append(appr.item()) + print(f"grasp : {traj['grasp']}, appr : {traj['appr']}") + # print(f"grasp : {grasp.item()}, appr : {appr.item()}") + return True, traj, recovery_activated_during_rollout, "" + + elif terminated or truncated: + grasp = torch.any(sub_obs['grasp']) + traj['grasp'].append(grasp.item()) + appr = torch.any(sub_obs['stacked']) + traj['appr'].append(appr.item()) + #print(f"grasp : {traj['grasp']}, appr : {traj['appr']}") + if i==horizon: + #if the failure was caused by object collision + traj['failure'].append("timeout") + failure_reason = "timeout" + else: + #if failure due to timeout, record timeout + traj['failure'].append("knocked") + failure_reason = "object_knocked" + return False, traj, recovery_activated_during_rollout, failure_reason + #print("REcovery : ", traj["recovery"]) + traj['failure'].append("timeout") + grasp = torch.any(sub_obs['grasp']) + traj['grasp'].append(grasp.item()) + appr = torch.any(sub_obs['stacked']) + traj['appr'].append(appr.item()) + return False, traj, recovery_activated_during_rollout, "" + + +def load_ensemble(device, ensemble_path): + """Load an ensemble of policies from checkpoint paths. + + Updated for robomimic v0.5 compatibility. + """ + with open(ensemble_path, 'r') as file: + ensemble_paths = [path.strip() for path in file.readlines()] + ensemble = [] + for path in ensemble_paths: + if not path: # Skip empty lines + continue + policy, _ = policy_from_checkpoint_v05(ckpt_path=path, device=device, verbose=True) + ensemble.append(policy) + print(f"[DEBUG] loaded {len(ensemble)} policies") + return ensemble + + +def clear_files(*paths : str): + if len(paths) < 1: + raise ValueError("At least one path is required.") + if os.path.exists(paths[0]): + #overwrite = input(f"You are about to overwrite rollout data at: {paths}\nContinue (y/n):") + overwrite = "y" + match overwrite.lower(): + case 'y' | 'Y': + for path in paths: + with open(path, 'w') as file: + pass + case _: + exit() + + + +def main(): + """Run a trained policy from robomimic with Isaac Lab environment.""" + # parse configuration + env_cfg = parse_env_cfg(args_cli.task, device=args_cli.device, num_envs=1, use_fabric=not args_cli.disable_fabric) + + # Set observations to dictionary mode for Robomimic + env_cfg.observations.policy.concatenate_terms = False + + #create a log handler + # loghelper = LoggingHelper() + + + + # Disable recorder + env_cfg.recorders = None + + # Extract success checking function + success_term = env_cfg.terminations.success + env_cfg.terminations.success = None + + # Get timeout signal + timeout = env_cfg.terminations.time_out + #env_cfg.terminations.time_out = None + + # Create environment + env = gym.make(args_cli.task, cfg=env_cfg).unwrapped + + # Set seed + torch.manual_seed(args_cli.seed) + ##### unset these two lines for deterministic randomness + # np.random.seed(args_cli.seed) + # random.seed(args_cli.seed) + + env.seed(args_cli.seed) + + # Acquire device + device = TorchUtils.get_torch_device(try_to_use_cuda=True) + + # load the stack cube ensemble + # 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/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') + + # Lets set these to the 0.99 confidence + parameters = { + 'beaker_lift' :{ + 0 : { + "confidence_level": 0.018646507043923632, + "window_size": 15, + "max_peaks": 10 + }, + 1 : { + "confidence_level": 0.018898154803458085, + "window_size": 10, + "max_peaks": 9 + }, + 2 : { + "confidence_level": 0.017589774107725654, + "window_size": 15, + "max_peaks": 7 + }, + 3 : { + "confidence_level": 0.02829341592326851, + "window_size": 12, + "max_peaks": 8 + }, + 4 : { + "confidence_level": 0.029170220902127186, + "window_size": 15, + "max_peaks": 11 + }, + 5 : { + "confidence_level": 0.017449474558650185, + "window_size": 12, + "max_peaks": 7 + }, + 6 : { + "confidence_level": 0.6324478323150455, + "window_size": 14, + "max_peaks": 14 + }, + + }, + } + + model_arch = "BC_RNN_GMM" + task = "pick_place" # stack_cube or pick_place + model_name = f"ensemble" + number = 'no_recovery_video' + + results_path = f"docs/training_data/{task}/uncertainty_rollout_{task}/{model_name}/run_{number}" + + uncertainties_path = f"{results_path}/uncertainties{number}.txt" + actions_path = f"{results_path}/actions{number}.txt" + min_actions_path = f"{results_path}/min_actions{number}.txt" + max_actions_path = f"{results_path}/max_actions{number}.txt" + time_taken_path = f"{results_path}/time_taken{number}.txt" + recovery_activated_during_rollout_path = f"{results_path}/recovery_activated_during_rollout.txt" + rollout_uncert_path = f"{uncertainties_path}/rollout_uncerts" + rollout_log_path = f"{uncertainties_path[:len(uncertainties_path)-4]}_rollout_log.txt" + + rollout_log_path = f"{uncertainties_path[:len(uncertainties_path)-4]}_rollout_log.txt" # remove the '.txt' and add rollout_log.txt + + uncertainty_results_path = f"{results_path}/uncertainty_plots" + trajectory_results_path = f"{results_path}/trajectory_plots" + os.makedirs(uncertainty_results_path, exist_ok=True) + os.makedirs(trajectory_results_path, exist_ok=True) + + + # clear files + clear_files( + uncertainties_path, actions_path, + min_actions_path, max_actions_path, + time_taken_path, recovery_activated_during_rollout_path + ) + + # Run policy + results = [] + failed_no_intervention =[] + intervention = [] + intervention_success = [] + knocked=0 + intervention_failure = [] + 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) + # save the uncertainties + print("Finished rollout, recovery needed : ", recovery_activated_during_rollout) + #print("actions shape : ", traj['actions']) + + if not terminated: + print(f"failure : {traj['failure']}") + for j in range(7): + with open(f"docs/rollouts/rollout_logging_no_recovery_joint{j}.txt", "a") as file: + #file.write(f"Logs for joint {j}\n") + for i, var in enumerate(traj['actions']): + # get all the info per joint + act = str(traj['actions'][i][0][j]) + max = str(traj['max_actions'][i].values.tolist()[j]) + min = str(traj['min_actions'][i].values.tolist()[j]) + uncert = str(traj['uncertainties'][i].tolist()[j]) + #print("joint pos : ", str(traj['joint_pos'][i][0].tolist()[j])) + joint_pos = str(traj['joint_pos'][i][0].tolist()[j]) + recovery = "False"#str(traj['recovery'][0][i]) + # print("writing recovery : ", recovery) + line = f"{i} : {act} : {max} : {min} : {uncert} : {joint_pos}: {recovery}: {terminated}\n" + file.write(line) + # this logs the outcome of the rollout - the failure methods + for j in range(7): + joint_uncert = [t[j].item() for t in traj['uncertainties']] + with open(f"docs/rollouts/rollout_uncerts_joint{j}.txt", "a") as f: + line = ",".join(["{:.10f}".format(v) for v in joint_uncert]) + f.write(line + f" : {terminated} \n") + # for i in range(len(traj['failure'])): + with open(f"docs/rollouts/rollout_outcomes.txt", "a") as file: + if terminated: + # print(f"grasp : {traj['grasp']}, appr : {traj['appr']}") + #trialnum : outcome : recovery needed ? : failure mode : grasp subtask : appr subtask + line=f"{trial}: {terminated} : {recovery_activated_during_rollout} : : {True} : {True} \n" + else : + grasp = traj['grasp'][0] + appr = traj['appr'][0] + print(f"Subatask : {grasp} , {appr}") + line=f"{trial} : {terminated} : {recovery_activated_during_rollout} : {traj['failure'][0]} : {grasp} : {appr}\n" + # print(f"grasp : {traj['grasp'][0]}, appr : {traj['appr'][0]}") + file.write(line) + + + results.append(terminated) + if not terminated: + if recovery_activated_during_rollout <=0: + failed_no_intervention.append(1) + + if recovery_activated_during_rollout > 0: + intervention.append(1) + if terminated: + intervention_success.append(1) + else: + intervention_failure.append(1) + else: + intervention.append(0) + if "knocked" in traj['failure'] or (len(traj['failure']) > 0 and traj['failure'][-1] == "knocked"): + knocked+=1 + + print(f'Current success rate: {results.count(True)/len(results):.2%}') + print(f'Current intervention rate: {sum(intervention)/len(intervention):.2%} ({sum(intervention)}/{len(intervention)})') + print(f"[INFO] Trial {trial}: {'SUCCESS' if terminated else 'FAILED'}\n") + #loghelper.stopEpoch(trial) + + # Prepare summary text + exp_name = args_cli.exp_name if args_cli.exp_name else "unnamed" + summary_text = f"\nSummary for experiment {exp_name}\n" + summary_text += f"\nSuccessful trials: {results.count(True)}, out of {len(results)} trials\n" + summary_text += f"Success rate: {results.count(True) / len(results)}\n" + summary_text += f'Failed runs with no intervention : {sum(failed_no_intervention)}\n' + summary_text += f"Intervention rate : {sum(intervention)/len(intervention)}. Of which success : {sum(intervention_success)}, of which failed {sum(intervention_failure)}\n" + summary_text += f"Number of times object knocked over : {knocked}\n" + summary_text += f"Trial Results: {results}\n\n" + summary_text += f"Environment random seed: {args_cli.seed}\n" + + # Print to console + print(summary_text) + + # Write to file + results_file = "docs/ensemble_run_results.txt" + with open(results_file, 'a') as f: + f.write(summary_text) + f.write("\n" + "="*80 + "\n\n") + print(f"Results appended to: {results_file}") + + env.close() + + + +if __name__ == "__main__": + # run the main function + main() + # close sim app + simulation_app.close() diff --git a/scripts/imitation_learning/robomimic/train.py b/scripts/imitation_learning/robomimic/train.py index 945c1f40f980..9203665a3ac9 100644 --- a/scripts/imitation_learning/robomimic/train.py +++ b/scripts/imitation_learning/robomimic/train.py @@ -87,7 +87,7 @@ import isaaclab_tasks.manager_based.manipulation.pick_place # noqa: F401 -def normalize_hdf5_actions(config: Config, log_dir: str) -> str: +def normalize_hdf5_actions(config: Config, log_dir: str) -> list: """Normalizes actions in hdf5 dataset to [-1, 1] range. Args: @@ -95,14 +95,22 @@ def normalize_hdf5_actions(config: Config, log_dir: str) -> str: log_dir: Directory to save normalization parameters. Returns: - Path to the normalized dataset. + List of dataset configs with normalized paths (robomimic v0.5 format). """ - base, ext = os.path.splitext(config.train.data) + # Handle both string path and list format (robomimic v0.5) + if isinstance(config.train.data, str): + original_path = config.train.data + elif isinstance(config.train.data, list): + original_path = config.train.data[0]["path"] + else: + raise ValueError(f"Unexpected config.train.data type: {type(config.train.data)}") + + base, ext = os.path.splitext(original_path) normalized_path = base + "_normalized" + ext # Copy the original dataset print(f"Creating normalized dataset at {normalized_path}") - shutil.copyfile(config.train.data, normalized_path) + shutil.copyfile(original_path, normalized_path) # Open the new dataset and normalize the actions with h5py.File(normalized_path, "r+") as f: @@ -126,11 +134,12 @@ def normalize_hdf5_actions(config: Config, log_dir: str) -> str: f[path] = normalized_data # Save the min and max values to log directory - with open(os.path.join(log_dir, "normalization_params.txt"), "w") as f: - f.write(f"min: {min}\n") - f.write(f"max: {max}\n") + with open(os.path.join(log_dir, "normalization_params.txt"), "w") as norm_f: + norm_f.write(f"min: {min}\n") + norm_f.write(f"max: {max}\n") - return normalized_path + # Return in robomimic v0.5 list format + return [{"path": normalized_path}] def train(config: Config, device: str, log_dir: str, ckpt_dir: str, video_dir: str): @@ -164,17 +173,29 @@ def train(config: Config, device: str, log_dir: str, ckpt_dir: str, video_dir: s # read config to set up metadata for observation modalities (e.g. detecting rgb observations) ObsUtils.initialize_obs_utils_with_config(config) - # make sure the dataset exists - dataset_path = os.path.expanduser(config.train.data) + # robomimic v0.5: config.train.data is a list of dataset configs + # Verify dataset exists (conversion already done in main() before lock) + dataset_path = os.path.expanduser(config.train.data[0]["path"]) if not os.path.exists(dataset_path): raise FileNotFoundError(f"Dataset at provided path {dataset_path} not found!") - - # load basic metadata from training file + + # load basic metadata from training file (use first dataset) print("\n============= Loaded Environment Metadata =============") - env_meta = FileUtils.get_env_metadata_from_dataset(dataset_path=config.train.data) + + env_meta = FileUtils.get_env_metadata_from_dataset(dataset_path=dataset_path) + # robomimic v0.5 changed API: now expects dataset_config dict with "path" key + print(f"config.train.data: {config.train.data}") + dataset_config = config.train.data[0] + + print(f"Dataset Config: {dataset_config}") + shape_meta = FileUtils.get_shape_metadata_from_dataset( - dataset_path=config.train.data, all_obs_keys=config.all_obs_keys, verbose=True + dataset_config=dataset_config, + action_keys=config.train.action_keys, + all_obs_keys=config.all_obs_keys, + verbose=True ) + print(f"shape_meta: {shape_meta}\n") if config.experiment.env is not None: env_meta["env_name"] = config.experiment.env @@ -202,9 +223,29 @@ def train(config: Config, device: str, log_dir: str, ckpt_dir: str, video_dir: s print(envs[env.name]) print("") + print(f"obs keys : {shape_meta['all_obs_keys']}") + + + # load training data BEFORE model creation (robomimic v0.5 requirement) + trainset, validset = TrainUtils.load_data_for_training(config, obs_keys=shape_meta["all_obs_keys"]) + train_sampler = trainset.get_dataset_sampler() + print("\n============= Training Dataset =============") + print(trainset) + print("") + + # robomimic v0.5: set num_train_batches and num_epochs in optim_params before model creation + train_num_steps = config.experiment.epoch_every_n_steps + for k in config.algo.optim_params: + config.algo.optim_params[k]["num_train_batches"] = len(trainset) if train_num_steps is None else train_num_steps + config.algo.optim_params[k]["num_epochs"] = config.train.num_epochs # setup for a new training run - data_logger = DataLogger(log_dir, config=config, log_tb=config.experiment.logging.log_tb) + data_logger = DataLogger( + log_dir, + config=config, + log_tb=config.experiment.logging.log_tb, + log_wandb=config.experiment.logging.log_wandb, + ) model = algo_factory( algo_name=config.algo_name, config=config, @@ -221,13 +262,6 @@ def train(config: Config, device: str, log_dir: str, ckpt_dir: str, video_dir: s print(model) # print model summary print("") - # load training data - trainset, validset = TrainUtils.load_data_for_training(config, obs_keys=shape_meta["all_obs_keys"]) - train_sampler = trainset.get_dataset_sampler() - print("\n============= Training Dataset =============") - print(trainset) - print("") - # maybe retrieve statistics for normalizing observations obs_normalization_stats = None if config.train.hdf5_normalize_obs: @@ -390,11 +424,17 @@ def main(args: argparse.Namespace): # change location of experiment directory config.train.output_dir = os.path.abspath(os.path.join("./logs", args.log_dir, args.task)) - log_dir, ckpt_dir, video_dir = TrainUtils.get_exp_dir(config) + # robomimic v0.5+ returns 4 values: log_dir, output_dir, video_dir, time_dir + log_dir, ckpt_dir, video_dir, _ = TrainUtils.get_exp_dir(config) if args.normalize_training_actions: config.train.data = normalize_hdf5_actions(config, log_dir) + # robomimic v0.5: config.train.data must be a list of dataset configs + # Convert single path string to list format if needed (must be done before lock) + if isinstance(config.train.data, str): + config.train.data = [{"path": os.path.expanduser(config.train.data)}] + # get torch device device = TorchUtils.get_torch_device(try_to_use_cuda=config.train.cuda) @@ -430,7 +470,7 @@ def main(args: argparse.Namespace): parser.add_argument("--task", type=str, default=None, help="Name of the task.") parser.add_argument("--algo", type=str, default=None, help="Name of the algorithm.") - parser.add_argument("--log_dir", type=str, default="robomimic", help="Path to log directory") + parser.add_argument("--log_dir", type=str, default="docs/place/model", help="Path to log directory") parser.add_argument("--normalize_training_actions", action="store_true", default=False, help="Normalize actions") parser.add_argument( "--epochs", diff --git a/scripts/imitation_learning/robomimic/train_ensemble.py b/scripts/imitation_learning/robomimic/train_ensemble.py index 047c62daece7..b2510eed8edb 100644 --- a/scripts/imitation_learning/robomimic/train_ensemble.py +++ b/scripts/imitation_learning/robomimic/train_ensemble.py @@ -175,10 +175,32 @@ def train(config: Config, device: str, log_dirs: list[str], ckpt_dirs: list[str] # load basic metadata from training file print("\n============= Loaded Environment Metadata =============") 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_stack_demos_split_normalized.hdf5'}] + ####### HARD CODED MAKE SURE THIS GETS UPDATED ######### + config_list= [{'path': 'docs/place/glassware_stack_demos_split_normalized.hdf5'}] + print(f"config.train.data: {config.train.data}") + # print(f"config_list: {config_list}") + # if isinstance(config.train.data, str): + # datasets = [] + # datasets.append(config.train.data) + # dataset_config = {"path": datasets} + # elif isinstance(config.train.data, dict): + # datasets = [] + # dataset_config= (config.train.data) + # dataset_config= {'path': 'docs/place/glassware_stack_demos_split_normalized.hdf5',} + # print(f"Dataset Config: {dataset_config}") + + dataset_config = config.train.data[0] shape_meta = FileUtils.get_shape_metadata_from_dataset( - dataset_path=config.train.data, all_obs_keys=config.all_obs_keys, verbose=True + dataset_config=dataset_config, + action_keys=config.train.action_keys, + all_obs_keys=config.all_obs_keys, + verbose=True ) - + print(f"shape_meta: {shape_meta}\n") + if config.experiment.env is not None: env_meta["env_name"] = config.experiment.env print("=" * 30 + "\n" + "Replacing Env to {}\n".format(env_meta["env_name"]) + "=" * 30) @@ -205,7 +227,7 @@ def train(config: Config, device: str, log_dirs: list[str], ckpt_dirs: list[str] print(envs[env.name]) print("") - + print(f"obs keys : {shape_meta['all_obs_keys']}") # load training data trainset, validset = TrainUtils.load_data_for_training(config, obs_keys=shape_meta["all_obs_keys"]) @@ -427,31 +449,22 @@ def keep_data_percentage(trainset, percentage): return trainset -def create_bootstrap_sample(trainset, seed: int): - """Create a bootstrap sample (sampling with replacement) from the training set. - - Each bootstrap sample contains N samples drawn with replacement from the original - dataset of size N. On average, ~63.2% of samples will be unique, with ~36.8% duplicates. - This creates diversity across ensemble members. +def create_bootstrap_sample(trainset, seed: int, percent: float = 0.5): + """Create a random sample (without replacement) of a percentage of the training set. Args: trainset: The original training dataset. seed: Random seed for reproducibility. - + percent: Fraction of the dataset to sample (default 0.75 for 75%). Returns: - A Subset of the training data with bootstrap sampling. + A Subset of the training data with random sampling (no replacement). """ np.random.seed(seed) total_size = len(trainset) - # Sample WITH replacement - same sample can appear multiple times - indices = np.random.choice(total_size, size=total_size, replace=True) - - # Count unique samples - unique_count = len(np.unique(indices)) - unique_percent = (unique_count / total_size) * 100 - - print(f"Bootstrap sample created: {unique_count}/{total_size} unique samples ({unique_percent:.1f}%)") - + sample_size = int(total_size * percent) + indices = np.random.choice(total_size, size=sample_size, replace=False) + indices = np.sort(indices) + print(f"Random sample created: {sample_size}/{total_size} samples ({percent*100:.1f}%)") return Subset(trainset, indices) @@ -515,7 +528,7 @@ def main(args: argparse.Namespace): new_output_dir = f"{original_output_dir}/model{i}/" config.train.output_dir = new_output_dir - log_dir, ckpt_dir, video_dir = TrainUtils.get_exp_dir(config) + log_dir, ckpt_dir, video_dir, _ = TrainUtils.get_exp_dir(config) log_dirs.append(log_dir) ckpt_dirs.append(ckpt_dir) diff --git a/scripts/tools/check_demos.py b/scripts/tools/check_demos.py new file mode 100644 index 000000000000..5ef552833b5a --- /dev/null +++ b/scripts/tools/check_demos.py @@ -0,0 +1,38 @@ +import h5py +import numpy as np + +f = h5py.File('/workspace/isaaclab/docs/place/glassware_stack_demos_split.hdf5', 'r') + +# Analyze gripper action (dim 6) over normalized time +demos = list(f['data'].keys()) +print(f"Number of demos: {len(demos)}") + +# Get a few demo lengths +lengths = [f['data'][d]['actions'].shape[0] for d in demos[:10]] +print(f"Demo lengths (first 10): {lengths}") + +# Analyze gripper at different phases +early_gripper = [] # First 10% of each demo +mid_gripper = [] # 30-50% of each demo +late_gripper = [] # Last 20% of each demo + +for demo in demos: + actions = f['data'][demo]['actions'][:] + n = len(actions) + gripper = actions[:, 6] # Gripper is dim 6 + + early_gripper.extend(gripper[:int(n*0.1)]) + mid_gripper.extend(gripper[int(n*0.3):int(n*0.5)]) + late_gripper.extend(gripper[int(n*0.8):]) + +print(f"\nGripper action analysis by phase:") +print(f"Early (0-10%): mean={np.mean(early_gripper):.3f}, std={np.std(early_gripper):.3f}") +print(f"Mid (30-50%): mean={np.mean(mid_gripper):.3f}, std={np.std(mid_gripper):.3f}") +print(f"Late (80-100%): mean={np.mean(late_gripper):.3f}, std={np.std(late_gripper):.3f}") + +# Check first few steps of first demo +first_demo = f['data'][demos[0]]['actions'][:] +print(f"\nFirst demo - first 20 gripper actions:") +print(first_demo[:20, 6]) + +f.close() \ No newline at end of file diff --git a/scripts/tools/copy_metadata_hdf5.py b/scripts/tools/copy_metadata_hdf5.py index 85fbcc84bf1b..3b57e1804b70 100644 --- a/scripts/tools/copy_metadata_hdf5.py +++ b/scripts/tools/copy_metadata_hdf5.py @@ -6,10 +6,11 @@ DATA_GROUP = "data" # Metadata to add +# robomimic v0.5 requires 'env_kwargs' key (not 'sim_args') metadata = { "env_name": "Cube-Mimic-v0", "type": 2, - "sim_args": { + "env_kwargs": { "dt": 0.01, "decimation": 2, "render_interval": 2, diff --git a/source/isaaclab/isaaclab/managers/recorder_manager.py b/source/isaaclab/isaaclab/managers/recorder_manager.py index 855c975f2a91..2f8829880807 100644 --- a/source/isaaclab/isaaclab/managers/recorder_manager.py +++ b/source/isaaclab/isaaclab/managers/recorder_manager.py @@ -429,8 +429,9 @@ def get_ep_meta(self) -> dict: """Get the episode metadata.""" if not hasattr(self._env.cfg, "get_ep_meta"): # Add basic episode metadata + # robomimic v0.5 requires 'env_kwargs' key (not 'sim_args') ep_meta = dict() - ep_meta["sim_args"] = { + ep_meta["env_kwargs"] = { "dt": self._env.cfg.sim.dt, "decimation": self._env.cfg.decimation, "render_interval": self._env.cfg.sim.render_interval, diff --git a/source/isaaclab/isaaclab/safety/switchingLogic.py b/source/isaaclab/isaaclab/safety/switchingLogic.py index bfcd265adbdf..92c188d98426 100644 --- a/source/isaaclab/isaaclab/safety/switchingLogic.py +++ b/source/isaaclab/isaaclab/safety/switchingLogic.py @@ -30,7 +30,7 @@ def halfway_check(self, step:int=0, grasp_state:bool=True): return True return False -s + def check(self, uncertainties, step) -> bool: # 1. Update windows for joint_num, unc in enumerate(uncertainties): diff --git a/source/isaaclab/isaaclab/utils/datasets/hdf5_dataset_file_handler.py b/source/isaaclab/isaaclab/utils/datasets/hdf5_dataset_file_handler.py index 13f07dddf378..882eac423d05 100644 --- a/source/isaaclab/isaaclab/utils/datasets/hdf5_dataset_file_handler.py +++ b/source/isaaclab/isaaclab/utils/datasets/hdf5_dataset_file_handler.py @@ -56,8 +56,9 @@ def create(self, file_path: str, env_name: str = None): # set environment arguments # the environment type (we use gym environment type) is set to be compatible with robomimic # Ref: https://github.com/ARISE-Initiative/robomimic/blob/master/robomimic/envs/env_base.py#L15 + # robomimic v0.5 requires 'env_kwargs' key (not 'sim_args') env_name = env_name if env_name is not None else "" - self.add_env_args({"env_name": env_name, "type": 2}) + self.add_env_args({"env_name": env_name, "type": 2, "env_kwargs": {}}) def __del__(self): """Destructor for the file handler.""" @@ -177,6 +178,16 @@ def create_dataset_helper(group, key, value): for key, value in episode.data.items(): create_dataset_helper(h5_episode_group, key, value) + # robomimic v0.5: generate next_obs by shifting obs by 1 timestep + if "obs" in episode.data and isinstance(episode.data["obs"], dict): + next_obs_group = h5_episode_group.create_group("next_obs") + for obs_key, obs_value in episode.data["obs"].items(): + if isinstance(obs_value, torch.Tensor): + obs_np = obs_value.cpu().numpy() + # next_obs[t] = obs[t+1], last timestep duplicated + next_obs_np = np.concatenate([obs_np[1:], obs_np[-1:]], axis=0) + next_obs_group.create_dataset(obs_key, data=next_obs_np, compression="gzip") + # increment total step counts self._hdf5_data_group.attrs["total"] += h5_episode_group.attrs["num_samples"] diff --git a/source/isaaclab_mimic/isaaclab_mimic/envs/__init__.py b/source/isaaclab_mimic/isaaclab_mimic/envs/__init__.py index 544d80075af4..c612ed504c9b 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/envs/__init__.py +++ b/source/isaaclab_mimic/isaaclab_mimic/envs/__init__.py @@ -19,6 +19,7 @@ from .cube_blueprint_mimic_env_cfg import CubeBlueprintMimicEnvCfg from .cube_mimic_env_cfg import CubeMimicEnvCfg from .cube_mimic_env import CubeMimicEnv +from .cube_mimic_env_skillgen_cfg import CubeMimicEnvSkillgenCfg ## # Inverse Kinematics - Relative Pose Control ## @@ -84,6 +85,15 @@ disable_env_checker=True, ) +gym.register( + id="Cube-IK-Rel-Skillgen-v0", + entry_point="isaaclab_mimic.envs:CubeMimicEnv", + kwargs={ + "env_cfg_entry_point": cube_mimic_env_skillgen_cfg.CubeMimicEnvSkillgenCfg, + }, + disable_env_checker=True, +) + gym.register( id="Isaac-Stack-Cube-Bin-Franka-IK-Rel-Mimic-v0", entry_point="isaaclab_mimic.envs:FrankaCubeStackIKRelMimicEnv", diff --git a/source/isaaclab_mimic/isaaclab_mimic/envs/cube_blueprint_mimic_env_cfg.py b/source/isaaclab_mimic/isaaclab_mimic/envs/cube_blueprint_mimic_env_cfg.py index 0d3e07572fd7..d85e778b4889 100755 --- a/source/isaaclab_mimic/isaaclab_mimic/envs/cube_blueprint_mimic_env_cfg.py +++ b/source/isaaclab_mimic/isaaclab_mimic/envs/cube_blueprint_mimic_env_cfg.py @@ -79,62 +79,39 @@ def __post_init__(self): ) ) #lift ? - subtask_configs.append( - SubTaskConfig( - # Each subtask involves manipulation with respect to a single object frame. - object_ref="object", - # End of final subtask does not need to be detected - subtask_term_signal="lift", - # No time offsets for the final subtask - subtask_term_offset_range=(5, 10), - # Selection strategy for source subtask segment - selection_strategy="nearest_neighbor_robot_distance", - # Optional parameters for the selection strategy function - selection_strategy_kwargs={"nn_k": 3}, - # Amount of action noise to apply during this subtask - action_noise=0.03, - # Number of interpolation steps to bridge to this subtask segment - num_interpolation_steps=5, - # Additional fixed steps for the robot to reach the necessary pose - num_fixed_steps=0, - # If True, apply action noise during the interpolation phase and execution - apply_noise_during_interpolation=False, - ) - ) - subtask_configs.append( - SubTaskConfig( - # Each subtask involves manipulation with respect to a single object frame. - object_ref="object", - # End of final subtask does not need to be detected - # LOL YOU DO - subtask_term_signal="appr_goal", - # No time offsets for the final subtask - subtask_term_offset_range=(0, 0), - # Selection strategy for source subtask segment - selection_strategy="nearest_neighbor_robot_distance", - # Optional parameters for the selection strategy function - selection_strategy_kwargs={"nn_k": 3}, - # Amount of action noise to apply during this subtask - action_noise=0.03, - # Number of interpolation steps to bridge to this subtask segment - num_interpolation_steps=5, - # Additional fixed steps for the robot to reach the necessary pose - num_fixed_steps=0, - # If True, apply action noise during the interpolation phase and execution - apply_noise_during_interpolation=False, - ) - ) + # subtask_configs.append( + # SubTaskConfig( + # # Each subtask involves manipulation with respect to a single object frame. + # object_ref="object", + # # End of final subtask does not need to be detected + # subtask_term_signal="lift", + # # No time offsets for the final subtask + # subtask_term_offset_range=(5, 10), + # # Selection strategy for source subtask segment + # selection_strategy="nearest_neighbor_robot_distance", + # # Optional parameters for the selection strategy function + # selection_strategy_kwargs={"nn_k": 3}, + # # Amount of action noise to apply during this subtask + # action_noise=0.03, + # # Number of interpolation steps to bridge to this subtask segment + # num_interpolation_steps=5, + # # Additional fixed steps for the robot to reach the necessary pose + # num_fixed_steps=0, + # # If True, apply action noise during the interpolation phase and execution + # apply_noise_during_interpolation=False, + # ) + # ) # subtask_configs.append( # SubTaskConfig( # # Each subtask involves manipulation with respect to a single object frame. # object_ref="object", # # End of final subtask does not need to be detected # # LOL YOU DO - # subtask_term_signal="release_object", + # subtask_term_signal="appr_goal", # # No time offsets for the final subtask # subtask_term_offset_range=(0, 0), # # Selection strategy for source subtask segment - # selection_strategy="nearest_neighbor_object", + # selection_strategy="nearest_neighbor_robot_distance", # # Optional parameters for the selection strategy function # selection_strategy_kwargs={"nn_k": 3}, # # Amount of action noise to apply during this subtask @@ -147,4 +124,27 @@ def __post_init__(self): # apply_noise_during_interpolation=False, # ) # ) + subtask_configs.append( + SubTaskConfig( + # Each subtask involves manipulation with respect to a single object frame. + object_ref="object", + # End of final subtask does not need to be detected + # LOL YOU DO + subtask_term_signal=None, + # No time offsets for the final subtask + subtask_term_offset_range=(0, 0), + # Selection strategy for source subtask segment + selection_strategy="nearest_neighbor_object", + # Optional parameters for the selection strategy function + selection_strategy_kwargs={"nn_k": 3}, + # Amount of action noise to apply during this subtask + action_noise=0.03, + # Number of interpolation steps to bridge to this subtask segment + num_interpolation_steps=5, + # Additional fixed steps for the robot to reach the necessary pose + num_fixed_steps=0, + # If True, apply action noise during the interpolation phase and execution + apply_noise_during_interpolation=False, + ) + ) self.subtask_configs["franka"] = subtask_configs \ No newline at end of file diff --git a/source/isaaclab_mimic/isaaclab_mimic/envs/cube_mimic_env.py b/source/isaaclab_mimic/isaaclab_mimic/envs/cube_mimic_env.py index a295df1c7fce..195e1abf19ab 100755 --- a/source/isaaclab_mimic/isaaclab_mimic/envs/cube_mimic_env.py +++ b/source/isaaclab_mimic/isaaclab_mimic/envs/cube_mimic_env.py @@ -113,10 +113,10 @@ def action_to_target_eef_pose(self, action: torch.Tensor) -> dict[str, torch.Ten delta_rotation_axis = delta_rotation / delta_rotation_angle # Handle invalid division for the case when delta_rotation_angle is close to zero - is_close_to_zero_angle = torch.isclose(delta_rotation_angle, torch.zeros_like(delta_rotation_angle)).squeeze(1) + is_close_to_zero_angle = torch.isclose(delta_rotation_angle, torch.zeros_like(delta_rotation_angle)).squeeze(-1) delta_rotation_axis[is_close_to_zero_angle] = torch.zeros_like(delta_rotation_axis)[is_close_to_zero_angle] - delta_quat = PoseUtils.quat_from_angle_axis(delta_rotation_angle.squeeze(1), delta_rotation_axis).squeeze(0) + delta_quat = PoseUtils.quat_from_angle_axis(delta_rotation_angle.squeeze(-1), delta_rotation_axis) delta_rot_mat = PoseUtils.matrix_from_quat(delta_quat) target_rot = torch.matmul(delta_rot_mat, curr_rot) @@ -158,8 +158,32 @@ def get_subtask_term_signals(self, env_ids: Sequence[int] | None = None) -> dict subtask_terms = self.obs_buf["subtask_terms"] #signals["appr"] = subtask_terms["appr"][env_ids] signals["grasp"] = subtask_terms["grasp"][env_ids] + signals["stacked"] = subtask_terms["stacked"][env_ids] #signals["lift"] = subtask_terms["lift"][env_ids] - signals["appr_goal"] = subtask_terms["appr_goal"][env_ids] - # signals["release_object"] = subtask_terms["release_object"][env_ids] - # final subtask is placing cubeC on cubeA (motion relative to cubeA) - but final subtask signal is not needed + #signals["appr_goal"] = subtask_terms["appr_goal"][env_ids] + #signals["release_object"] = subtask_terms["release_object"][env_ids] + # final subtask is placing cube on goal - signal needed when using --annotate_subtask_start_signals return signals + + def get_expected_attached_object(self, eef_name: str, subtask_index: int, env_cfg) -> str | None: + """ + (SkillGen) Return the expected attached object for the given EEF/subtask. + + Assumes 'stack' subtasks place the object grasped in the preceding 'grasp' subtask. + Returns None for 'grasp' (or others) at subtask start. + """ + if eef_name not in env_cfg.subtask_configs: + return None + + subtask_configs = env_cfg.subtask_configs[eef_name] + if not (0 <= subtask_index < len(subtask_configs)): + return None + + current_cfg = subtask_configs[subtask_index] + # If stacking, expect we are holding the object grasped in the prior subtask + if "stack" in str(current_cfg.subtask_term_signal).lower(): + if subtask_index > 0: + prev_cfg = subtask_configs[subtask_index - 1] + if "grasp" in str(prev_cfg.subtask_term_signal).lower(): + return prev_cfg.object_ref + return None diff --git a/source/isaaclab_mimic/isaaclab_mimic/envs/cube_mimic_env_cfg.py b/source/isaaclab_mimic/isaaclab_mimic/envs/cube_mimic_env_cfg.py index 20181d3aa4a2..71ce2096b510 100755 --- a/source/isaaclab_mimic/isaaclab_mimic/envs/cube_mimic_env_cfg.py +++ b/source/isaaclab_mimic/isaaclab_mimic/envs/cube_mimic_env_cfg.py @@ -26,7 +26,7 @@ def __post_init__(self): self.datagen_config.name = "beaker_rand_scale_D0" self.datagen_config.generation_guarantee = True self.datagen_config.generation_keep_failed = False - self.datagen_config.generation_num_trials = 100 + self.datagen_config.generation_num_trials = 10 self.datagen_config.generation_select_src_per_subtask = True self.datagen_config.generation_transform_first_robot_pose = False self.datagen_config.generation_interpolate_from_last_target_pose = True @@ -81,6 +81,8 @@ def __post_init__(self): num_fixed_steps=0, # If True, apply action noise during the interpolation phase and execution apply_noise_during_interpolation=False, + description="Grasp Glassware", + next_subtask_description="Stack glassware on scale", ) ) #lift ? @@ -106,29 +108,29 @@ def __post_init__(self): # apply_noise_during_interpolation=False, # ) # ) - subtask_configs.append( - SubTaskConfig( - # Each subtask involves manipulation with respect to a single object frame. - object_ref="object", - # End of final subtask does not need to be detected - # This is actually the lift - subtask_term_signal="appr_goal", - # No time offsets for the final subtask - subtask_term_offset_range=(10, 20), - # Selection strategy for source subtask segment - selection_strategy="nearest_neighbor_robot_distance", - # Optional parameters for the selection strategy function - selection_strategy_kwargs={"nn_k": 3}, - # Amount of action noise to apply during this subtask - action_noise=0.03, - # Number of interpolation steps to bridge to this subtask segment - num_interpolation_steps=5, - # Additional fixed steps for the robot to reach the necessary pose - num_fixed_steps=0, - # If True, apply action noise during the interpolation phase and execution - apply_noise_during_interpolation=False, - ) - ) + # subtask_configs.append( + # SubTaskConfig( + # # Each subtask involves manipulation with respect to a single object frame. + # object_ref="object", + # # End of final subtask does not need to be detected + # # This is actually the lift + # subtask_term_signal="appr_goal", + # # No time offsets for the final subtask + # subtask_term_offset_range=(10, 20), + # # Selection strategy for source subtask segment + # selection_strategy="nearest_neighbor_robot_distance", + # # Optional parameters for the selection strategy function + # selection_strategy_kwargs={"nn_k": 3}, + # # Amount of action noise to apply during this subtask + # action_noise=0.03, + # # Number of interpolation steps to bridge to this subtask segment + # num_interpolation_steps=5, + # # Additional fixed steps for the robot to reach the necessary pose + # num_fixed_steps=0, + # # If True, apply action noise during the interpolation phase and execution + # apply_noise_during_interpolation=False, + # ) + # ) subtask_configs.append( SubTaskConfig( @@ -136,11 +138,11 @@ def __post_init__(self): object_ref="object", # End of final subtask - when using --annotate_subtask_start_signals, # the last subtask must also have a signal name - subtask_term_signal="success", + subtask_term_signal="stacked", # No time offsets for the final subtask subtask_term_offset_range=(0, 0), # Selection strategy for source subtask segment - selection_strategy="nearest_neighbor_robot_distance", + selection_strategy="nearest_neighbor_object_distance", # Optional parameters for the selection strategy function selection_strategy_kwargs={"nn_k": 3}, # Amount of action noise to apply during this subtask @@ -151,6 +153,7 @@ def __post_init__(self): num_fixed_steps=0, # If True, apply action noise during the interpolation phase and execution apply_noise_during_interpolation=False, + description="Stack glassware on scale", ) ) diff --git a/source/isaaclab_mimic/isaaclab_mimic/envs/cube_mimic_env_skillgen_cfg.py b/source/isaaclab_mimic/isaaclab_mimic/envs/cube_mimic_env_skillgen_cfg.py new file mode 100755 index 000000000000..196a87c6e47e --- /dev/null +++ b/source/isaaclab_mimic/isaaclab_mimic/envs/cube_mimic_env_skillgen_cfg.py @@ -0,0 +1,90 @@ +# Copyright (c) 2024-2025, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from isaaclab.envs.mimic_env_cfg import MimicEnvCfg, SubTaskConfig +from isaaclab.utils import configclass + +from isaaclab_tasks.manager_based.manipulation.cube_lift.config.franka.dev_ik_rel_skillgen import FrankaDevEnvSkillgenCfg + + +@configclass +class CubeMimicEnvSkillgenCfg(FrankaDevEnvSkillgenCfg, MimicEnvCfg): + """ + Isaac Lab Mimic environment config class for Franka Cube Stack IK Rel env. + """ + + def __post_init__(self): + # post init of parents + super().__post_init__() + # # TODO: Figure out how we can move this to the MimicEnvCfg class + # # The __post_init__() above only calls the init for FrankaCubeStackEnvCfg and not MimicEnvCfg + # # https://stackoverflow.com/questions/59986413/achieving-multiple-inheritance-using-python-dataclasses + + # Override the existing values + self.datagen_config.name = "beaker_rand_scale_D0" + self.datagen_config.generation_guarantee = True + self.datagen_config.generation_keep_failed = False + self.datagen_config.generation_num_trials = 10 + self.datagen_config.generation_select_src_per_subtask = True + self.datagen_config.generation_transform_first_robot_pose = False + self.datagen_config.generation_interpolate_from_last_target_pose = True + self.datagen_config.generation_relative = True + self.datagen_config.max_num_failures = 100 + self.datagen_config.seed = 101 + + # The following are the subtask configurations for the stack task. + # The following are the subtask configurations for the stack task. + subtask_configs = [] + + subtask_configs.append( + SubTaskConfig( + # Each subtask involves manipulation with respect to a single object frame. + object_ref="object", + # Corresponding key for the binary indicator in "datagen_info" for completion + subtask_term_signal="grasp", + # Time offsets for data generation when splitting a trajectory + subtask_term_offset_range=(0, 0), + # Selection strategy for source subtask segment + selection_strategy="nearest_neighbor_object", + # Optional parameters for the selection strategy function + selection_strategy_kwargs={"nn_k": 3}, + # Amount of action noise to apply during this subtask + action_noise=0.03, + # Number of interpolation steps to bridge to this subtask segment + num_interpolation_steps=0, + # Additional fixed steps for the robot to reach the necessary pose + num_fixed_steps=0, + # If True, apply action noise during the interpolation phase and execution + apply_noise_during_interpolation=False, + description="Grasp Glassware", + next_subtask_description="Stack glassware on scale", + ) + ) + + subtask_configs.append( + SubTaskConfig( + # Each subtask involves manipulation with respect to a single object frame. + object_ref="object", + # End of final subtask - when using --annotate_subtask_start_signals, + # the last subtask must also have a signal name + subtask_term_signal="stacked", + # No time offsets for the final subtask + subtask_term_offset_range=(0, 0), + # Selection strategy for source subtask segment + selection_strategy="nearest_neighbor_object_distance", + # Optional parameters for the selection strategy function + selection_strategy_kwargs={"nn_k": 3}, + # Amount of action noise to apply during this subtask + action_noise=0.03, + # Number of interpolation steps to bridge to this subtask segment + num_interpolation_steps=0, + # Additional fixed steps for the robot to reach the necessary pose + num_fixed_steps=0, + # If True, apply action noise during the interpolation phase and execution + apply_noise_during_interpolation=False, + description="Stack glassware on scale", + ) + ) + self.subtask_configs["franka"] = subtask_configs \ No newline at end of file diff --git a/source/isaaclab_mimic/isaaclab_mimic/motion_planners/curobo/curobo_planner.py b/source/isaaclab_mimic/isaaclab_mimic/motion_planners/curobo/curobo_planner.py index f9c6cf69cbdb..70b7272540aa 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/motion_planners/curobo/curobo_planner.py +++ b/source/isaaclab_mimic/isaaclab_mimic/motion_planners/curobo/curobo_planner.py @@ -181,7 +181,7 @@ def __init__( self.n_repeat: int | None = self.config.n_repeat self.step_size: float | None = self.config.motion_step_size self.visualize_plan: bool = self.config.visualize_plan - self.visualize_spheres: bool = self.config.visualize_spheres + self.visualize_spheres: bool = False # Log the config parameter values self.logger.info(f"Config parameter values: {self.config}") diff --git a/source/isaaclab_mimic/isaaclab_mimic/motion_planners/curobo/curobo_planner_cfg.py b/source/isaaclab_mimic/isaaclab_mimic/motion_planners/curobo/curobo_planner_cfg.py index eaf28e0ed153..3290a42f693d 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/motion_planners/curobo/curobo_planner_cfg.py +++ b/source/isaaclab_mimic/isaaclab_mimic/motion_planners/curobo/curobo_planner_cfg.py @@ -126,7 +126,7 @@ class CuroboPlannerCfg: time_dilation_factor: float = 1.0 """Time dilation factor for planning.""" - surface_sphere_radius: float = 0.005 + surface_sphere_radius: float = 0.001 """Radius of surface spheres for collision checking.""" # Debug and visualization @@ -136,13 +136,13 @@ class CuroboPlannerCfg: motion_step_size: float | None = None """Step size (in radians) for retiming motion plans. If None, no retiming.""" - visualize_spheres: bool = False + visualize_spheres: bool = True """Visualize robot collision spheres. Note: only works for env 0.""" visualize_plan: bool = False """Visualize motion plan in Rerun. Note: only works for env 0.""" - debug_planner: bool = False + debug_planner: bool = True """Enable detailed motion planning debug information.""" sphere_update_freq: int = 5 @@ -442,14 +442,14 @@ def franka_stack_cube_config(cls) -> "CuroboPlannerCfg": def franka_beaker_config(cls) -> "CuroboPlannerCfg": """Create configuration for Franka stacking a normal cube.""" config = cls.franka_config() - config.static_objects = ["table"] + config.static_objects = ["table", "scale"] config.visualize_plan = True config.debug_planner = True config.motion_noise_scale = 0.02 config.collision_activation_distance = 0.0 config.approach_distance = 0.05 config.retreat_distance = 0.05 - config.surface_sphere_radius = 0.005 + config.surface_sphere_radius = 0.001 config.get_world_config = lambda: config._get_world_config_with_table_adjustment() return config @@ -469,7 +469,7 @@ def from_task_name(cls, task_name: str) -> "CuroboPlannerCfg": return cls.franka_stack_cube_bin_config() elif "stack-cube" in task_lower: return cls.franka_stack_cube_config() - elif "Cube-Mimic" in task_lower: + elif "cube-ik-rel" in task_lower: return cls.franka_beaker_config() else: # Default to Franka configuration diff --git a/source/isaaclab_mimic/setup.py b/source/isaaclab_mimic/setup.py index 658aed9ee807..84f14e1f2cb1 100644 --- a/source/isaaclab_mimic/setup.py +++ b/source/isaaclab_mimic/setup.py @@ -29,7 +29,7 @@ # Check if the platform is Linux and add the dependency if platform.system() == "Linux": - EXTRAS_REQUIRE["robomimic"].append("robomimic@git+https://github.com/ARISE-Initiative/robomimic.git@v0.4.0") + EXTRAS_REQUIRE["robomimic"].append("robomimic@git+https://github.com/ARISE-Initiative/robomimic.git@v0.5.0") # Cumulation of all extra-requires EXTRAS_REQUIRE["all"] = list(itertools.chain.from_iterable(EXTRAS_REQUIRE.values())) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/beaker_pour/config/franka/agents/robomimic/diffusion_policy.json b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/beaker_pour/config/franka/agents/robomimic/diffusion_policy.json index 47fafeea8cb1..27353a78c1ea 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/beaker_pour/config/franka/agents/robomimic/diffusion_policy.json +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/beaker_pour/config/franka/agents/robomimic/diffusion_policy.json @@ -40,6 +40,7 @@ "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", @@ -70,7 +71,10 @@ "initial": 0.0001, "decay_factor": 1.0, "epoch_schedule": [], - "scheduler_type": "constant_with_warmup" + "scheduler_type": "cosine", + "step_every_batch": true, + "warmup_steps": 500, + "num_cycles": 0.5 }, "regularization": { "L2": 0.0001 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 7f7b0e7bae8d..0c07f8dcb01b 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 @@ -147,13 +147,22 @@ entry_point="isaaclab.envs:ManagerBasedRLEnv", kwargs={ "env_cfg_entry_point": f"{__name__}.dev_ik_rel_env_place: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-Place-Skillgen-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + kwargs={ + "env_cfg_entry_point": f"{__name__}.dev_ik_rel_skillgen:FrankaDevEnvSkillgenCfg", + "robomimic_bc_cfg_entry_point": os.path.join(agents.__path__[0], "robomimic/bc_rnn_low_dim.json"), }, disable_env_checker=True, ) \ No newline at end of file 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 d5071d660c5d..8125682666ea 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": "pick_place_new_demos_med_qual" + "wandb_proj_name": "pick_place_bc_rnn_0.5_mix" }, "save": { "enabled": true, @@ -20,7 +20,7 @@ }, "epoch_every_n_steps": 100, "validation_epoch_every_n_steps": 10, - "env": "Dev-IK-Rel-v1", + "env": "Dev-IK-Rel-place-v0", "additional_envs": null, "render": false, "render_video": true, @@ -28,9 +28,9 @@ "video_skip": 5, "rollout": { "enabled": false, - "terminate_on_success" : true, - "n": 50, - "horizon": 400, + "terminate_on_success": true, + "n": 50, + "horizon": 400, "rate": 50 } }, @@ -43,13 +43,15 @@ "hdf5_normalize_obs": false, "hdf5_filter_key": "train", "hdf5_validation_filter_key": "valid", - "seq_length": 100, + "seq_length": 16, "dataset_keys": [ "actions", "rewards", "dones" ], - + "action_keys": [ + "actions" + ], "goal_mode": null, "cuda": true, "batch_size": 100, @@ -76,7 +78,6 @@ "l1_weight": 0.0, "cos_weight": 0.0 }, - "gmm": { "enabled": true, "num_modes": 10, @@ -84,7 +85,6 @@ "std_activation": "softplus", "low_noise_eval": true }, - "rnn": { "enabled": true, "horizon": 100, @@ -96,7 +96,6 @@ "bidirectional": false } } - }, "observation": { "modalities": { diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/agents/robomimic/diffusion_policy.json b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/agents/robomimic/diffusion_policy.json index ae7cffd9ea58..9d6e437ec9b7 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/agents/robomimic/diffusion_policy.json +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/agents/robomimic/diffusion_policy.json @@ -12,7 +12,7 @@ "save": { "enabled": true, "every_n_seconds": null, - "every_n_epochs": 50, + "every_n_epochs": null, "epochs": [], "on_best_validation": true, "on_best_rollout_return": false, @@ -43,11 +43,10 @@ "hdf5_normalize_obs": false, "hdf5_filter_key": "train", "hdf5_validation_filter_key": "valid", - "seq_length": 10, - "frame_stack": 1, - "pad_frame_stack": true, + "seq_length": 16, "pad_seq_length": true, - "get_pad_mask": false, + "frame_stack": 2, + "pad_frame_stack": true, "goal_mode": null, "dataset_keys": [ "actions", @@ -58,8 +57,8 @@ "actions" ], "cuda": true, - "batch_size": 256, - "num_epochs": 3000, + "batch_size": 128, + "num_epochs": 1000, "seed": 101 }, "algo": { @@ -68,9 +67,12 @@ "optimizer_type": "adamw", "learning_rate": { "initial": 0.0001, - "decay_factor": 1.0, + "decay_factor": 0.1, "epoch_schedule": [], - "scheduler_type": "constant_with_warmup" + "scheduler_type": "cosine", + "step_every_batch": true, + "warmup_steps": 500, + "num_cycles": 0.5 }, "regularization": { "L2": 0.0001 @@ -79,7 +81,7 @@ }, "horizon": { "observation_horizon": 2, - "action_horizon": 8, + "action_horizon": 4, "prediction_horizon": 16 }, "unet": { diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/dev_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/dev_env_cfg.py index 896391c9b6dd..2c2eba3b3b98 100755 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/dev_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/dev_env_cfg.py @@ -67,6 +67,9 @@ def __post_init__(self): open_command_expr={"panda_finger_.*": 0.04}, close_command_expr={"panda_finger_.*": 0.0}, ) + self.gripper_joint_names = ["panda_finger_.*"] + self.gripper_open_val = 0.04 + self.gripper_threshold = 0.005 # Set the body name for the end effector self.commands.object_pose.body_name = "panda_hand" self.commands.object_pose.body_name = "panda_hand" diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/dev_ik_rel_env_place.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/dev_ik_rel_env_place.py index fd754ccfbc58..223211fd0669 100755 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/dev_ik_rel_env_place.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/dev_ik_rel_env_place.py @@ -76,10 +76,10 @@ def __post_init__(self): scale=0.5, body_offset=DifferentialInverseKinematicsActionCfg.OffsetCfg(pos=[0.0, 0.0, 0.0]), ) - self.terminations.success=DoneTerm(func=mdp.object_stacked_upright, params={"lower_object_cfg": SceneEntityCfg("scale")}) + #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} - ) + #self.observations.subtask_terms.appr_goal=ObsTerm(func=mdp.is_object_lifted, params={"threshold":0.15} + #) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/dev_ik_rel_skillgen.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/dev_ik_rel_skillgen.py new file mode 100755 index 000000000000..0872b185059a --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/dev_ik_rel_skillgen.py @@ -0,0 +1,87 @@ +# 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 +from source.isaaclab_assets.isaaclab_assets.robots.universal_robots import UR10_CFG +from . import dev_env_cfg +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.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 + + + +@configclass +class FrankaDevEnvSkillgenCfg(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.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", body_names="Object"), + "asset2_cfg" : SceneEntityCfg("stirplate"), + "asset3_cfg" : SceneEntityCfg("scale") + + }, + ) + #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_stacked_upright, params={"lower_object_cfg": SceneEntityCfg("scale")}) + + #self.observations.subtask_terms.appr_goal=ObsTerm(func=mdp.is_object_lifted, params={"threshold":0.15} + #) + # for f in self.scene.ee_frame.target_frames: + # if f.name == "end_effector": + # f.offset.pos = [0.0, 0.0, 0.0] + # break + 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 new file mode 100644 index 000000000000..40a0566b68d5 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/config/franka/dev_ik_rel_vial_insert.py @@ -0,0 +1,97 @@ +# 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 +from source.isaaclab_assets.isaaclab_assets.robots.universal_robots import UR10_CFG +from . import dev_env_cfg +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.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 + + + +@configclass +class FrankaDevEnvCfg(dev_env_cfg.FrankaDevEnvCfg): + def __post_init__(self): + # post init of parent + super().__post_init__() + # get the vial from the stirplate and place in vial rack + glassware = ChemistryGlassware() + self.scene.object=glassware.vial(pos=[0.5, 0.0, 0.02],name="Object") + self.scene.stirplate = glassware.stirplate(pos=[0.5, 0.0, 0.01]) + + self.scene.vialrack = glassware.vialrack(pos=[0.3, -0.3, 0.01]) + + 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", body_names="Object"), + "asset2_cfg" : SceneEntityCfg("stirplate"), + "asset3_cfg" : SceneEntityCfg("scale") + + }, + ) + #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")}) + + #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/lift_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/lift_env_cfg.py index eb2138fa5ddf..9b74ea1e1364 100755 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/lift_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/cube_lift/lift_env_cfg.py @@ -182,6 +182,14 @@ class SubtaskCfg(ObsGroup): "object_cfg": SceneEntityCfg("object"), }, ) + stacked = ObsTerm( + func=mdp.object_stacked, + params={ + "robot_cfg": SceneEntityCfg("robot"), + "upper_object_cfg": SceneEntityCfg("object"), + "lower_object_cfg": SceneEntityCfg("scale"), + }, + ) # lift = ObsTerm( # func=mdp.is_object_lifted, # params={ @@ -189,13 +197,13 @@ class SubtaskCfg(ObsGroup): # "threshold" : 0.1 # } # ) - appr_goal = ObsTerm( - func=mdp.object_near_goal, - params={ - "threshold": 0.05, - "command_name": "object_pose", - }, - ) + # appr_goal = ObsTerm( + # func=mdp.object_near_goal, + # params={ + # "threshold": 0.05, + # "command_name": "object_pose", + # }, + # ) # release_object = ObsTerm( # func=mdp.object_released, # params={ @@ -308,7 +316,8 @@ class TerminationsCfg: ) object_tipped = DoneTerm(func=mdp.object_knocked) #fix object dropping - success = DoneTerm(func=mdp.object_near_goal) + #success = DoneTerm(func=mdp.object_near_goal) + success= DoneTerm(func=mdp.object_stacked_upright, params={"lower_object_cfg": SceneEntityCfg("scale")}) @configclass 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 75355c70d87a..119f5ef00407 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 @@ -151,7 +151,7 @@ def is_object_lifted( ): #return true when object z coord above a threshold value object = env.scene[obj_cfg.name] - print(f"Object z pos : {object.data.root_pos_w[:, 2]}") + # print(f"Object z pos : {object.data.root_pos_w[:, 2]}") # if object.data.root_pos_w[:, 2].item() > threshold : # print(f"Observed Object Lifted : {object.data.root_pos_w[:, 2].item()}") # loghelper.logsubtask(LogType.LIFT) @@ -304,10 +304,8 @@ def object_tilt (env: ManagerBasedRLEnv, upper_object_cfg: SceneEntityCfg = Scen def object_knocked(env: ManagerBasedRLEnv, upper_object_cfg: SceneEntityCfg = SceneEntityCfg("object"), max_tilt: float = 45): 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]) - if tilt_deg.cpu().detach().numpy()[0] > max_tilt: - # print("object knocked") - return torch.tensor([True], device='cuda:0') - else :return torch.tensor([False], device='cuda:0') + # Return a bool tensor with shape (num_envs,) for proper batching + return tilt_deg > max_tilt def target_position( env: ManagerBasedRLEnv, @@ -318,4 +316,28 @@ def target_position( object: RigidObject = env.scene[object_cfg.name] object_pos_w = object.data.root_pos_w[:, :3] #print(f"target position : {object_pos_w}") - return object_pos_w \ No newline at end of file + return object_pos_w + + +def target_pose( + env: ManagerBasedRLEnv, + object_cfg: SceneEntityCfg = SceneEntityCfg("scale"), +) -> torch.Tensor: + """The pose (position + quaternion) of the target object in world frame. + + Returns: + Tensor of shape (num_envs, 7) containing [x, y, z, qw, qx, qy, qz]. + """ + object: RigidObject = env.scene[object_cfg.name] + object_pos_w = object.data.root_pos_w[:, :3] # (num_envs, 3) + object_quat_w = object.data.root_quat_w # (num_envs, 4) - already wxyz format + return torch.cat([object_pos_w, object_quat_w], dim=-1) # (num_envs, 7) + +# def object_distance_to_goal( +# env: ManagerBasedRLEnv, +# object_cfg: SceneEntityCfg = SceneEntityCfg("object"), +# goal_object_location: SceneEntityCfg = SceneEntityCfg("vialrack"), +# ) -> torch.Tensor: + + + 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 6065810915a6..fa2ab92788f9 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 @@ -99,16 +99,39 @@ def placed( def object_stacked_upright(env: ManagerBasedRLEnv, robot_cfg: SceneEntityCfg = SceneEntityCfg("robot"), upper_object_cfg: SceneEntityCfg = SceneEntityCfg("object"),lower_object_cfg: SceneEntityCfg = SceneEntityCfg("goal"), - xy_threshold: float = 0.1, height_threshold: float = 0.1, height_diff: float = 0.05, + xy_threshold: float = 0.05, height_threshold: float = 0.1, height_diff: float = 0.05, + 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: - """Stacked AND sufficiently upright (ā‰¤ upright_good_deg).""" + + """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]) 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) stacked_upright = stacked & upright_good + gripper_joint_ids, _ = robot.find_joints(env.cfg.gripper_joint_names) + stacked = torch.logical_and( + torch.isclose( + robot.data.joint_pos[:, gripper_joint_ids[0]], + torch.tensor(env.cfg.gripper_open_val, dtype=torch.float32).to(env.device), + atol=atol, + rtol=rtol, + ), + stacked_upright, + ) + stacked = torch.logical_and( + torch.isclose( + robot.data.joint_pos[:, gripper_joint_ids[1]], + torch.tensor(env.cfg.gripper_open_val, dtype=torch.float32).to(env.device), + atol=atol, + rtol=rtol, + ), + stacked, + ) #if logging: #rising = log(env, "object_stacked_upright", stacked_upright, f"Stacked (upright ā‰¤ {upright_good_deg}Ā°)") @@ -116,7 +139,7 @@ def object_stacked_upright(env: ManagerBasedRLEnv, robot_cfg: SceneEntityCfg = S #if idx.numel() > 0: # print(f"[stacked-upright tilt] n={int(idx.numel())} sample={tilt_deg[idx][:5].tolist()}") - return stacked_upright + return stacked def upright_tilt_deg(q_cur, q_init) -> torch.Tensor: @@ -127,7 +150,7 @@ def upright_tilt_deg(q_cur, q_init) -> torch.Tensor: def object_stacked(env: ManagerBasedRLEnv, robot_cfg: SceneEntityCfg = SceneEntityCfg("robot"), upper_object_cfg: SceneEntityCfg = SceneEntityCfg("object"), lower_object_cfg: SceneEntityCfg = SceneEntityCfg("goal"), - xy_threshold: float = 0.1, height_threshold: float = 0.1, height_diff: float = 0.05, + xy_threshold: float = 0.05, height_threshold: float = 0.1, height_diff: float = 0.05, gripper_open_val: torch.Tensor = torch.tensor([0.04]),logging=False) -> torch.Tensor: """Check if an object is stacked by the specified robot.""" @@ -142,5 +165,49 @@ def object_stacked(env: ManagerBasedRLEnv, robot_cfg: SceneEntityCfg = SceneEnti # gripper_open_val.to(env.device), atol=1e-4, rtol=1e-4), stacked) # stacked = torch.logical_and(torch.isclose(robot.data.joint_pos[:, -2], # gripper_open_val.to(env.device), atol=1e-4, rtol=1e-4), stacked) - print(f"For DEBUG : STACKED STATUS : {stacked}") - return stacked \ No newline at end of file + #print(f"For DEBUG : STACKED STATUS : {stacked}") + return stacked + +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.01, height_diff: float = 0.05, + 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: + + + """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]) + 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) + stacked_upright = stacked & upright_good + gripper_joint_ids, _ = robot.find_joints(env.cfg.gripper_joint_names) + stacked = torch.logical_and( + torch.isclose( + robot.data.joint_pos[:, gripper_joint_ids[0]], + torch.tensor(env.cfg.gripper_open_val, dtype=torch.float32).to(env.device), + atol=atol, + rtol=rtol, + ), + stacked_upright, + ) + stacked = torch.logical_and( + torch.isclose( + robot.data.joint_pos[:, gripper_joint_ids[1]], + torch.tensor(env.cfg.gripper_open_val, dtype=torch.float32).to(env.device), + atol=atol, + rtol=rtol, + ), + stacked, + ) + + #if logging: + #rising = log(env, "object_stacked_upright", stacked_upright, f"Stacked (upright ā‰¤ {upright_good_deg}Ā°)") + # idx = torch.nonzero(rising, as_tuple=False).squeeze(-1) + #if idx.numel() > 0: + # print(f"[stacked-upright tilt] n={int(idx.numel())} sample={tilt_deg[idx][:5].tolist()}") + + return inserted \ No newline at end of file diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/config/franka/stack_ik_rel_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/config/franka/stack_ik_rel_env_cfg.py index f173ee644cef..ad944c137153 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/config/franka/stack_ik_rel_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/config/franka/stack_ik_rel_env_cfg.py @@ -40,30 +40,3 @@ def __post_init__(self): body_offset=DifferentialInverseKinematicsActionCfg.OffsetCfg(pos=[0.0, 0.0, 0.107]), ) - self.teleop_devices = DevicesCfg( - devices={ - "handtracking": OpenXRDeviceCfg( - retargeters=[ - Se3RelRetargeterCfg( - bound_hand=OpenXRDevice.TrackingTarget.HAND_RIGHT, - zero_out_xy_rotation=True, - use_wrist_rotation=False, - use_wrist_position=True, - delta_pos_scale_factor=10.0, - delta_rot_scale_factor=10.0, - sim_device=self.sim.device, - ), - GripperRetargeterCfg( - bound_hand=OpenXRDevice.TrackingTarget.HAND_RIGHT, sim_device=self.sim.device - ), - ], - sim_device=self.sim.device, - xr_cfg=self.xr, - ), - "keyboard": Se3KeyboardCfg( - pos_sensitivity=0.05, - rot_sensitivity=0.05, - sim_device=self.sim.device, - ), - } - ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/config/franka/stack_ik_rel_instance_randomize_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/config/franka/stack_ik_rel_instance_randomize_env_cfg.py index cee2530ee4f7..d647f76f3d94 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/config/franka/stack_ik_rel_instance_randomize_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/config/franka/stack_ik_rel_instance_randomize_env_cfg.py @@ -6,7 +6,7 @@ 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.glassware.glassware import ChemistryGlassware from . import stack_joint_pos_instance_randomize_env_cfg ## @@ -28,7 +28,7 @@ def __post_init__(self): self.scene.robot = FRANKA_PANDA_HIGH_PD_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") # Reduce the number of environments due to camera resources - self.scene.num_envs = 2 + self.scene.num_envs = 1 # Set actions for the specific robot type (franka) self.actions.arm_action = DifferentialInverseKinematicsActionCfg( diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/config/franka/stack_joint_pos_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/config/franka/stack_joint_pos_env_cfg.py index 6b8561684fe0..a73d35ee5f61 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/config/franka/stack_joint_pos_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/config/franka/stack_joint_pos_env_cfg.py @@ -21,8 +21,8 @@ # Pre-defined configs ## from isaaclab.markers.config import FRAME_MARKER_CFG # isort: skip -from isaaclab_assets.robots.franka import FRANKA_PANDA_CFG # isort: skip - +from isaaclab_assets.robots.franka import FRANKA_PANDA_CFG +from isaaclab_assets.glassware.glassware import ChemistryGlassware @configclass class EventCfg: @@ -32,7 +32,8 @@ class EventCfg: func=franka_stack_events.set_default_joint_pose, mode="reset", params={ - "default_pose": [0.0444, -0.1894, -0.1107, -2.5148, 0.0044, 2.3775, 0.6952, 0.0400, 0.0400], + "default_pose": [0.3281, -0.3684, -0.2787, -2.6138, -2.7527, 2.4999, 0.3331, 0.0400, 0.0400], + #"default_pose": [0.0444, -0.1894, -0.1107, -2.5148, 0.0044, 2.3775, 0.6952, 0.0400, 0.0400], }, ) @@ -50,8 +51,8 @@ class EventCfg: func=franka_stack_events.randomize_object_pose, mode="reset", params={ - "pose_range": {"x": (0.4, 0.6), "y": (-0.10, 0.10), "z": (0.0203, 0.0203), "yaw": (-1.0, 1, 0)}, - "min_separation": 0.1, + "pose_range": {"x": (0.4, 0.6), "y": (-0.10, 0.10), "z": (0.0203, 0.0203)}, + "min_separation": 0.4, "asset_cfgs": [SceneEntityCfg("cube_1"), SceneEntityCfg("cube_2"), SceneEntityCfg("cube_3")], }, ) @@ -101,38 +102,11 @@ def __post_init__(self): max_depenetration_velocity=5.0, disable_gravity=False, ) - + glassware = ChemistryGlassware() # Set each stacking cube deterministically - self.scene.cube_1 = RigidObjectCfg( - prim_path="{ENV_REGEX_NS}/Cube_1", - init_state=RigidObjectCfg.InitialStateCfg(pos=[0.4, 0.0, 0.0203], rot=[1, 0, 0, 0]), - spawn=UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/blue_block.usd", - scale=(1.0, 1.0, 1.0), - rigid_props=cube_properties, - semantic_tags=[("class", "cube_1")], - ), - ) - self.scene.cube_2 = RigidObjectCfg( - prim_path="{ENV_REGEX_NS}/Cube_2", - init_state=RigidObjectCfg.InitialStateCfg(pos=[0.55, 0.05, 0.0203], rot=[1, 0, 0, 0]), - spawn=UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/red_block.usd", - scale=(1.0, 1.0, 1.0), - rigid_props=cube_properties, - semantic_tags=[("class", "cube_2")], - ), - ) - self.scene.cube_3 = RigidObjectCfg( - prim_path="{ENV_REGEX_NS}/Cube_3", - init_state=RigidObjectCfg.InitialStateCfg(pos=[0.60, -0.1, 0.0203], rot=[1, 0, 0, 0]), - spawn=UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/green_block.usd", - scale=(1.0, 1.0, 1.0), - rigid_props=cube_properties, - semantic_tags=[("class", "cube_3")], - ), - ) + self.scene.cube_1 = glassware.scale(pos=[0.4, -0.35, 0.01], name="cube_1"), + self.scene.cube_2 = glassware.beaker(pos=[0.45, 0.25, 0.02], name="cube_2"), + # Listens to the required transforms marker_cfg = FRAME_MARKER_CFG.copy() diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/config/franka/stack_joint_pos_instance_randomize_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/config/franka/stack_joint_pos_instance_randomize_env_cfg.py index 360f1f6482e1..dc105e2d59bf 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/config/franka/stack_joint_pos_instance_randomize_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/config/franka/stack_joint_pos_instance_randomize_env_cfg.py @@ -5,15 +5,12 @@ import torch -from isaaclab.assets import RigidObjectCfg, RigidObjectCollectionCfg +from isaaclab.assets import RigidObjectCollectionCfg from isaaclab.managers import EventTermCfg as EventTerm from isaaclab.managers import SceneEntityCfg from isaaclab.sensors import FrameTransformerCfg from isaaclab.sensors.frame_transformer.frame_transformer_cfg import OffsetCfg -from isaaclab.sim.schemas.schemas_cfg import RigidBodyPropertiesCfg -from isaaclab.sim.spawners.from_files.from_files_cfg import UsdFileCfg from isaaclab.utils import configclass -from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR from isaaclab_tasks.manager_based.manipulation.stack_glassware import mdp from isaaclab_tasks.manager_based.manipulation.stack_glassware.mdp import franka_stack_events @@ -21,6 +18,8 @@ StackInstanceRandomizeEnvCfg, ) +from isaaclab_assets.glassware.glassware import ChemistryGlassware + ## # Pre-defined configs ## @@ -36,7 +35,7 @@ class EventCfg: func=franka_stack_events.set_default_joint_pose, mode="startup", params={ - "default_pose": [0.0444, -0.1894, -0.1107, -2.5148, 0.0044, 2.3775, 0.6952, 0.0400, 0.0400], + "default_pose": [0.3281, -0.3684, -0.2787, -2.6138, -2.7527, 2.4999, 0.3331, 0.0400, 0.0400], }, ) @@ -54,10 +53,10 @@ class EventCfg: func=franka_stack_events.randomize_rigid_objects_in_focus, mode="reset", params={ - "asset_cfgs": [SceneEntityCfg("cube_1"), SceneEntityCfg("cube_2"), SceneEntityCfg("cube_3")], + "asset_cfgs": [SceneEntityCfg("cube_1"), SceneEntityCfg("cube_2")], "out_focus_state": torch.tensor([10.0, 10.0, 10.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), - "pose_range": {"x": (0.4, 0.6), "y": (-0.10, 0.10), "z": (0.0203, 0.0203), "yaw": (-1.0, 1, 0)}, - "min_separation": 0.1, + "pose_range": {"x": (0.3, 0.65), "y": (-0.35, 0.30), "z": (0.02, 0.02)}, + "min_separation": 0.15, }, ) @@ -75,7 +74,7 @@ def __post_init__(self): self.scene.robot = FRANKA_PANDA_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") # Reduce the number of environments due to camera resources - self.scene.num_envs = 2 + self.scene.num_envs = 1 # Set actions for the specific robot type (franka) self.actions.arm_action = mdp.JointPositionActionCfg( @@ -92,83 +91,23 @@ def __post_init__(self): self.gripper_open_val = 0.04 self.gripper_threshold = 0.005 - # Rigid body properties of each cube - cube_properties = RigidBodyPropertiesCfg( - solver_position_iteration_count=16, - solver_velocity_iteration_count=1, - max_angular_velocity=1000.0, - max_linear_velocity=1000.0, - max_depenetration_velocity=5.0, - disable_gravity=False, - ) + # Initialize ChemistryGlassware for creating glassware objects + glassware = ChemistryGlassware() - # Set each stacking cube to be a collection of rigid objects + # cube_1: Scale (target placement surface) - two variants at different positions cube_1_config_dict = { - "blue_cube": RigidObjectCfg( - prim_path="{ENV_REGEX_NS}/Cube_1_Blue", - init_state=RigidObjectCfg.InitialStateCfg(pos=[0.4, 0.0, 0.0203], rot=[1, 0, 0, 0]), - spawn=UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/blue_block.usd", - scale=(1.0, 1.0, 1.0), - rigid_props=cube_properties, - ), - ), - "red_cube": RigidObjectCfg( - prim_path="{ENV_REGEX_NS}/Cube_1_Red", - init_state=RigidObjectCfg.InitialStateCfg(pos=[0.4, 0.0, 0.0403], rot=[1, 0, 0, 0]), - spawn=UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/red_block.usd", - scale=(1.0, 1.0, 1.0), - rigid_props=cube_properties, - ), - ), + "scale_1": glassware.scale(pos=[0.4, -0.35, 0.01], name="Scale_1"), + "scale_2": glassware.scale(pos=[10.0, 10.0, 0.01], name="Scale_2"), # Out of focus position } + # cube_2: Beaker and Conical Flask - positioned far right (positive y) cube_2_config_dict = { - "red_cube": RigidObjectCfg( - prim_path="{ENV_REGEX_NS}/Cube_2_Red", - init_state=RigidObjectCfg.InitialStateCfg(pos=[0.55, 0.05, 0.0203], rot=[1, 0, 0, 0]), - spawn=UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/red_block.usd", - scale=(1.0, 1.0, 1.0), - rigid_props=cube_properties, - ), - ), - "yellow_cube": RigidObjectCfg( - prim_path="{ENV_REGEX_NS}/Cube_2_Yellow", - init_state=RigidObjectCfg.InitialStateCfg(pos=[0.55, 0.05, 0.0403], rot=[1, 0, 0, 0]), - spawn=UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/yellow_block.usd", - scale=(1.0, 1.0, 1.0), - rigid_props=cube_properties, - ), - ), - } - - cube_3_config_dict = { - "yellow_cube": RigidObjectCfg( - prim_path="{ENV_REGEX_NS}/Cube_3_Yellow", - init_state=RigidObjectCfg.InitialStateCfg(pos=[0.60, -0.1, 0.0203], rot=[1, 0, 0, 0]), - spawn=UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/yellow_block.usd", - scale=(1.0, 1.0, 1.0), - rigid_props=cube_properties, - ), - ), - "green_cube": RigidObjectCfg( - prim_path="{ENV_REGEX_NS}/Cube_2_Green", - init_state=RigidObjectCfg.InitialStateCfg(pos=[0.60, -0.1, 0.0403], rot=[1, 0, 0, 0]), - spawn=UsdFileCfg( - usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/green_block.usd", - scale=(1.0, 1.0, 1.0), - rigid_props=cube_properties, - ), - ), + "beaker": glassware.beaker(pos=[0.45, 0.25, 0.02], name="Beaker_1"), + "flask": glassware.flask(pos=[10.0, 10.0, 0.02], name="Flask_1", scale=(0.7, 0.7, 0.7)), # Out of focus position } self.scene.cube_1 = RigidObjectCollectionCfg(rigid_objects=cube_1_config_dict) self.scene.cube_2 = RigidObjectCollectionCfg(rigid_objects=cube_2_config_dict) - self.scene.cube_3 = RigidObjectCollectionCfg(rigid_objects=cube_3_config_dict) # Listens to the required transforms marker_cfg = FRAME_MARKER_CFG.copy() diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/mdp/franka_stack_events.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/mdp/franka_stack_events.py index 009a44b1b372..5dc97b3e05b2 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/mdp/franka_stack_events.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/mdp/franka_stack_events.py @@ -231,9 +231,13 @@ def randomize_rigid_objects_in_focus( object_states = torch.stack([out_focus_state] * asset.num_objects).to(device=env.device) pose_tensor = torch.tensor([pose_list[asset_idx]], device=env.device) positions = pose_tensor[:, 0:3] + env.scene.env_origins[cur_env, 0:3] - orientations = math_utils.quat_from_euler_xyz(pose_tensor[:, 3], pose_tensor[:, 4], pose_tensor[:, 5]) + + # Keep the current/default orientation instead of randomizing + # default_object_state shape is [num_envs, num_objects, 13] - get orientation (indices 3:7) + current_orientation = asset.data.default_object_state[cur_env, object_id, 3:7] + object_states[object_id, 0:3] = positions - object_states[object_id, 3:7] = orientations + object_states[object_id, 3:7] = current_orientation asset.write_object_state_to_sim( object_state=object_states, env_ids=torch.tensor([cur_env], device=env.device) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/mdp/observations.py index 2f65cd916ee9..d1e127d4a4c6 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/mdp/observations.py @@ -21,91 +21,76 @@ def cube_positions_in_world_frame( env: ManagerBasedRLEnv, cube_1_cfg: SceneEntityCfg = SceneEntityCfg("cube_1"), cube_2_cfg: SceneEntityCfg = SceneEntityCfg("cube_2"), - cube_3_cfg: SceneEntityCfg = SceneEntityCfg("cube_3"), ) -> torch.Tensor: """The position of the cubes in the world frame.""" cube_1: RigidObject = env.scene[cube_1_cfg.name] cube_2: RigidObject = env.scene[cube_2_cfg.name] - cube_3: RigidObject = env.scene[cube_3_cfg.name] - return torch.cat((cube_1.data.root_pos_w, cube_2.data.root_pos_w, cube_3.data.root_pos_w), dim=1) + return torch.cat((cube_1.data.root_pos_w, cube_2.data.root_pos_w), dim=1) def instance_randomize_cube_positions_in_world_frame( env: ManagerBasedRLEnv, cube_1_cfg: SceneEntityCfg = SceneEntityCfg("cube_1"), cube_2_cfg: SceneEntityCfg = SceneEntityCfg("cube_2"), - cube_3_cfg: SceneEntityCfg = SceneEntityCfg("cube_3"), ) -> torch.Tensor: """The position of the cubes in the world frame.""" if not hasattr(env, "rigid_objects_in_focus"): - return torch.full((env.num_envs, 9), fill_value=-1) + return torch.full((env.num_envs, 6), fill_value=-1) cube_1: RigidObjectCollection = env.scene[cube_1_cfg.name] cube_2: RigidObjectCollection = env.scene[cube_2_cfg.name] - cube_3: RigidObjectCollection = env.scene[cube_3_cfg.name] cube_1_pos_w = [] cube_2_pos_w = [] - cube_3_pos_w = [] for env_id in range(env.num_envs): cube_1_pos_w.append(cube_1.data.object_pos_w[env_id, env.rigid_objects_in_focus[env_id][0], :3]) cube_2_pos_w.append(cube_2.data.object_pos_w[env_id, env.rigid_objects_in_focus[env_id][1], :3]) - cube_3_pos_w.append(cube_3.data.object_pos_w[env_id, env.rigid_objects_in_focus[env_id][2], :3]) cube_1_pos_w = torch.stack(cube_1_pos_w) cube_2_pos_w = torch.stack(cube_2_pos_w) - cube_3_pos_w = torch.stack(cube_3_pos_w) - return torch.cat((cube_1_pos_w, cube_2_pos_w, cube_3_pos_w), dim=1) + return torch.cat((cube_1_pos_w, cube_2_pos_w), dim=1) def cube_orientations_in_world_frame( env: ManagerBasedRLEnv, cube_1_cfg: SceneEntityCfg = SceneEntityCfg("cube_1"), cube_2_cfg: SceneEntityCfg = SceneEntityCfg("cube_2"), - cube_3_cfg: SceneEntityCfg = SceneEntityCfg("cube_3"), ): """The orientation of the cubes in the world frame.""" cube_1: RigidObject = env.scene[cube_1_cfg.name] cube_2: RigidObject = env.scene[cube_2_cfg.name] - cube_3: RigidObject = env.scene[cube_3_cfg.name] - return torch.cat((cube_1.data.root_quat_w, cube_2.data.root_quat_w, cube_3.data.root_quat_w), dim=1) + return torch.cat((cube_1.data.root_quat_w, cube_2.data.root_quat_w), dim=1) def instance_randomize_cube_orientations_in_world_frame( env: ManagerBasedRLEnv, cube_1_cfg: SceneEntityCfg = SceneEntityCfg("cube_1"), cube_2_cfg: SceneEntityCfg = SceneEntityCfg("cube_2"), - cube_3_cfg: SceneEntityCfg = SceneEntityCfg("cube_3"), ) -> torch.Tensor: """The orientation of the cubes in the world frame.""" if not hasattr(env, "rigid_objects_in_focus"): - return torch.full((env.num_envs, 9), fill_value=-1) + return torch.full((env.num_envs, 8), fill_value=-1) cube_1: RigidObjectCollection = env.scene[cube_1_cfg.name] cube_2: RigidObjectCollection = env.scene[cube_2_cfg.name] - cube_3: RigidObjectCollection = env.scene[cube_3_cfg.name] cube_1_quat_w = [] cube_2_quat_w = [] - cube_3_quat_w = [] for env_id in range(env.num_envs): cube_1_quat_w.append(cube_1.data.object_quat_w[env_id, env.rigid_objects_in_focus[env_id][0], :4]) cube_2_quat_w.append(cube_2.data.object_quat_w[env_id, env.rigid_objects_in_focus[env_id][1], :4]) - cube_3_quat_w.append(cube_3.data.object_quat_w[env_id, env.rigid_objects_in_focus[env_id][2], :4]) cube_1_quat_w = torch.stack(cube_1_quat_w) cube_2_quat_w = torch.stack(cube_2_quat_w) - cube_3_quat_w = torch.stack(cube_3_quat_w) - return torch.cat((cube_1_quat_w, cube_2_quat_w, cube_3_quat_w), dim=1) + return torch.cat((cube_1_quat_w, cube_2_quat_w), dim=1) def object_obs( env: ManagerBasedRLEnv, cube_1_cfg: SceneEntityCfg = SceneEntityCfg("cube_1"), cube_2_cfg: SceneEntityCfg = SceneEntityCfg("cube_2"), - cube_3_cfg: SceneEntityCfg = SceneEntityCfg("cube_3"), ee_frame_cfg: SceneEntityCfg = SceneEntityCfg("ee_frame"), ): """ @@ -114,18 +99,12 @@ def object_obs( cube_1 quat, cube_2 pos, cube_2 quat, - cube_3 pos, - cube_3 quat, gripper to cube_1, gripper to cube_2, - gripper to cube_3, cube_1 to cube_2, - cube_2 to cube_3, - cube_1 to cube_3, """ cube_1: RigidObject = env.scene[cube_1_cfg.name] cube_2: RigidObject = env.scene[cube_2_cfg.name] - cube_3: RigidObject = env.scene[cube_3_cfg.name] ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] cube_1_pos_w = cube_1.data.root_pos_w @@ -134,17 +113,17 @@ def object_obs( cube_2_pos_w = cube_2.data.root_pos_w cube_2_quat_w = cube_2.data.root_quat_w - cube_3_pos_w = cube_3.data.root_pos_w - cube_3_quat_w = cube_3.data.root_quat_w - ee_pos_w = ee_frame.data.target_pos_w[:, 0, :] gripper_to_cube_1 = cube_1_pos_w - ee_pos_w gripper_to_cube_2 = cube_2_pos_w - ee_pos_w - gripper_to_cube_3 = cube_3_pos_w - ee_pos_w cube_1_to_2 = cube_1_pos_w - cube_2_pos_w - cube_2_to_3 = cube_2_pos_w - cube_3_pos_w - cube_1_to_3 = cube_1_pos_w - cube_3_pos_w + + # Debug print: distance between cube_1 and cube_2 + xy_dist = torch.linalg.vector_norm(cube_1_to_2[:, :2], dim=1) + h_dist = torch.abs(cube_1_to_2[:, 2]) + print(f"[object_obs] cube_1_to_2 xy_dist: {xy_dist[0].item():.4f}, h_dist: {h_dist[0].item():.4f}, " + f"z_diff: {cube_1_to_2[0, 2].item():.4f}", flush=True) return torch.cat( ( @@ -152,14 +131,9 @@ def object_obs( cube_1_quat_w, cube_2_pos_w - env.scene.env_origins, cube_2_quat_w, - cube_3_pos_w - env.scene.env_origins, - cube_3_quat_w, gripper_to_cube_1, gripper_to_cube_2, - gripper_to_cube_3, cube_1_to_2, - cube_2_to_3, - cube_1_to_3, ), dim=1, ) @@ -169,7 +143,6 @@ def instance_randomize_object_obs( env: ManagerBasedRLEnv, cube_1_cfg: SceneEntityCfg = SceneEntityCfg("cube_1"), cube_2_cfg: SceneEntityCfg = SceneEntityCfg("cube_2"), - cube_3_cfg: SceneEntityCfg = SceneEntityCfg("cube_3"), ee_frame_cfg: SceneEntityCfg = SceneEntityCfg("ee_frame"), ): """ @@ -178,51 +151,36 @@ def instance_randomize_object_obs( cube_1 quat, cube_2 pos, cube_2 quat, - cube_3 pos, - cube_3 quat, gripper to cube_1, gripper to cube_2, - gripper to cube_3, cube_1 to cube_2, - cube_2 to cube_3, - cube_1 to cube_3, """ if not hasattr(env, "rigid_objects_in_focus"): - return torch.full((env.num_envs, 9), fill_value=-1) + return torch.full((env.num_envs, 23), fill_value=-1) cube_1: RigidObjectCollection = env.scene[cube_1_cfg.name] cube_2: RigidObjectCollection = env.scene[cube_2_cfg.name] - cube_3: RigidObjectCollection = env.scene[cube_3_cfg.name] ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] cube_1_pos_w = [] cube_2_pos_w = [] - cube_3_pos_w = [] cube_1_quat_w = [] cube_2_quat_w = [] - cube_3_quat_w = [] for env_id in range(env.num_envs): cube_1_pos_w.append(cube_1.data.object_pos_w[env_id, env.rigid_objects_in_focus[env_id][0], :3]) cube_2_pos_w.append(cube_2.data.object_pos_w[env_id, env.rigid_objects_in_focus[env_id][1], :3]) - cube_3_pos_w.append(cube_3.data.object_pos_w[env_id, env.rigid_objects_in_focus[env_id][2], :3]) cube_1_quat_w.append(cube_1.data.object_quat_w[env_id, env.rigid_objects_in_focus[env_id][0], :4]) cube_2_quat_w.append(cube_2.data.object_quat_w[env_id, env.rigid_objects_in_focus[env_id][1], :4]) - cube_3_quat_w.append(cube_3.data.object_quat_w[env_id, env.rigid_objects_in_focus[env_id][2], :4]) cube_1_pos_w = torch.stack(cube_1_pos_w) cube_2_pos_w = torch.stack(cube_2_pos_w) - cube_3_pos_w = torch.stack(cube_3_pos_w) cube_1_quat_w = torch.stack(cube_1_quat_w) cube_2_quat_w = torch.stack(cube_2_quat_w) - cube_3_quat_w = torch.stack(cube_3_quat_w) ee_pos_w = ee_frame.data.target_pos_w[:, 0, :] gripper_to_cube_1 = cube_1_pos_w - ee_pos_w gripper_to_cube_2 = cube_2_pos_w - ee_pos_w - gripper_to_cube_3 = cube_3_pos_w - ee_pos_w cube_1_to_2 = cube_1_pos_w - cube_2_pos_w - cube_2_to_3 = cube_2_pos_w - cube_3_pos_w - cube_1_to_3 = cube_1_pos_w - cube_3_pos_w return torch.cat( ( @@ -230,14 +188,9 @@ def instance_randomize_object_obs( cube_1_quat_w, cube_2_pos_w - env.scene.env_origins, cube_2_quat_w, - cube_3_pos_w - env.scene.env_origins, - cube_3_quat_w, gripper_to_cube_1, gripper_to_cube_2, - gripper_to_cube_3, cube_1_to_2, - cube_2_to_3, - cube_1_to_3, ), dim=1, ) @@ -355,6 +308,11 @@ def object_stacked( height_dist = torch.linalg.vector_norm(pos_diff[:, 2:], dim=1) xy_dist = torch.linalg.vector_norm(pos_diff[:, :2], dim=1) + # Debug print: distance between upper and lower objects + print(f"[object_stacked] xy_dist: {xy_dist[0].item():.4f} (thresh: {xy_threshold}), " + f"h_dist: {height_dist[0].item():.4f} (expected: {height_diff}, thresh: {height_threshold}), " + f"z_diff: {pos_diff[0, 2].item():.4f}", flush=True) + stacked = torch.logical_and(xy_dist < xy_threshold, (height_dist - height_diff) < height_threshold) if hasattr(env.scene, "surface_grippers") and len(env.scene.surface_grippers) > 0: @@ -395,7 +353,6 @@ def cube_poses_in_base_frame( env: ManagerBasedRLEnv, cube_1_cfg: SceneEntityCfg = SceneEntityCfg("cube_1"), cube_2_cfg: SceneEntityCfg = SceneEntityCfg("cube_2"), - cube_3_cfg: SceneEntityCfg = SceneEntityCfg("cube_3"), robot_cfg: SceneEntityCfg = SceneEntityCfg("robot"), return_key: Literal["pos", "quat", None] = None, ) -> torch.Tensor: @@ -403,15 +360,12 @@ def cube_poses_in_base_frame( cube_1: RigidObject = env.scene[cube_1_cfg.name] cube_2: RigidObject = env.scene[cube_2_cfg.name] - cube_3: RigidObject = env.scene[cube_3_cfg.name] pos_cube_1_world = cube_1.data.root_pos_w pos_cube_2_world = cube_2.data.root_pos_w - pos_cube_3_world = cube_3.data.root_pos_w quat_cube_1_world = cube_1.data.root_quat_w quat_cube_2_world = cube_2.data.root_quat_w - quat_cube_3_world = cube_3.data.root_quat_w robot: Articulation = env.scene[robot_cfg.name] root_pos_w = robot.data.root_pos_w @@ -423,12 +377,9 @@ def cube_poses_in_base_frame( pos_cube_2_base, quat_cube_2_base = math_utils.subtract_frame_transforms( root_pos_w, root_quat_w, pos_cube_2_world, quat_cube_2_world ) - pos_cube_3_base, quat_cube_3_base = math_utils.subtract_frame_transforms( - root_pos_w, root_quat_w, pos_cube_3_world, quat_cube_3_world - ) - pos_cubes_base = torch.cat((pos_cube_1_base, pos_cube_2_base, pos_cube_3_base), dim=1) - quat_cubes_base = torch.cat((quat_cube_1_base, quat_cube_2_base, quat_cube_3_base), dim=1) + pos_cubes_base = torch.cat((pos_cube_1_base, pos_cube_2_base), dim=1) + quat_cubes_base = torch.cat((quat_cube_1_base, quat_cube_2_base), dim=1) if return_key == "pos": return pos_cubes_base @@ -442,7 +393,6 @@ def object_abs_obs_in_base_frame( env: ManagerBasedRLEnv, cube_1_cfg: SceneEntityCfg = SceneEntityCfg("cube_1"), cube_2_cfg: SceneEntityCfg = SceneEntityCfg("cube_2"), - cube_3_cfg: SceneEntityCfg = SceneEntityCfg("cube_3"), ee_frame_cfg: SceneEntityCfg = SceneEntityCfg("ee_frame"), robot_cfg: SceneEntityCfg = SceneEntityCfg("robot"), ): @@ -452,14 +402,11 @@ def object_abs_obs_in_base_frame( cube_1 quat, cube_2 pos, cube_2 quat, - cube_3 pos, - cube_3 quat, gripper pos, gripper quat, """ cube_1: RigidObject = env.scene[cube_1_cfg.name] cube_2: RigidObject = env.scene[cube_2_cfg.name] - cube_3: RigidObject = env.scene[cube_3_cfg.name] ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] robot: Articulation = env.scene[robot_cfg.name] @@ -472,18 +419,12 @@ def object_abs_obs_in_base_frame( cube_2_pos_w = cube_2.data.root_pos_w cube_2_quat_w = cube_2.data.root_quat_w - cube_3_pos_w = cube_3.data.root_pos_w - cube_3_quat_w = cube_3.data.root_quat_w - pos_cube_1_base, quat_cube_1_base = math_utils.subtract_frame_transforms( root_pos_w, root_quat_w, cube_1_pos_w, cube_1_quat_w ) pos_cube_2_base, quat_cube_2_base = math_utils.subtract_frame_transforms( root_pos_w, root_quat_w, cube_2_pos_w, cube_2_quat_w ) - pos_cube_3_base, quat_cube_3_base = math_utils.subtract_frame_transforms( - root_pos_w, root_quat_w, cube_3_pos_w, cube_3_quat_w - ) ee_pos_w = ee_frame.data.target_pos_w[:, 0, :] ee_quat_w = ee_frame.data.target_quat_w[:, 0, :] @@ -495,8 +436,6 @@ def object_abs_obs_in_base_frame( quat_cube_1_base, pos_cube_2_base, quat_cube_2_base, - pos_cube_3_base, - quat_cube_3_base, ee_pos_base, ee_quat_base, ), diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/mdp/terminations.py index e306f9eb4a0a..f988a1f14dfe 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/mdp/terminations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/mdp/terminations.py @@ -26,35 +26,35 @@ def cubes_stacked( robot_cfg: SceneEntityCfg = SceneEntityCfg("robot"), cube_1_cfg: SceneEntityCfg = SceneEntityCfg("cube_1"), cube_2_cfg: SceneEntityCfg = SceneEntityCfg("cube_2"), - cube_3_cfg: SceneEntityCfg = SceneEntityCfg("cube_3"), xy_threshold: float = 0.04, height_threshold: float = 0.005, height_diff: float = 0.0468, atol=0.0001, rtol=0.0001, ): + """Check if cube_2 is stacked on cube_1.""" robot: Articulation = env.scene[robot_cfg.name] cube_1: RigidObject = env.scene[cube_1_cfg.name] cube_2: RigidObject = env.scene[cube_2_cfg.name] - cube_3: RigidObject = env.scene[cube_3_cfg.name] + # Position difference: cube_1 (lower) - cube_2 (upper) pos_diff_c12 = cube_1.data.root_pos_w - cube_2.data.root_pos_w - pos_diff_c23 = cube_2.data.root_pos_w - cube_3.data.root_pos_w # Compute cube position difference in x-y plane xy_dist_c12 = torch.norm(pos_diff_c12[:, :2], dim=1) - xy_dist_c23 = torch.norm(pos_diff_c23[:, :2], dim=1) # Compute cube height difference h_dist_c12 = torch.norm(pos_diff_c12[:, 2:], dim=1) - h_dist_c23 = torch.norm(pos_diff_c23[:, 2:], dim=1) - # Check cube positions - stacked = torch.logical_and(xy_dist_c12 < xy_threshold, xy_dist_c23 < xy_threshold) + # Debug print: distance between cube_2 and cube_1 + print(f"[cubes_stacked] xy_dist: {xy_dist_c12[0].item():.4f} (thresh: {xy_threshold}), " + f"h_dist: {h_dist_c12[0].item():.4f} (expected: {height_diff}, thresh: {height_threshold}), " + f"z_diff: {pos_diff_c12[0, 2].item():.4f}") + + # Check cube positions: cube_2 should be on top of cube_1 + stacked = xy_dist_c12 < xy_threshold stacked = torch.logical_and(h_dist_c12 - height_diff < height_threshold, stacked) - stacked = torch.logical_and(pos_diff_c12[:, 2] < 0.0, stacked) - stacked = torch.logical_and(h_dist_c23 - height_diff < height_threshold, stacked) - stacked = torch.logical_and(pos_diff_c23[:, 2] < 0.0, stacked) + stacked = torch.logical_and(pos_diff_c12[:, 2] < 0.0, stacked) # cube_1.z < cube_2.z means cube_2 is on top # Check gripper positions if hasattr(env.scene, "surface_grippers") and len(env.scene.surface_grippers) > 0: diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/stack_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/stack_env_cfg.py index 3af83de3a19a..8a6e27071031 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/stack_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack_glassware/stack_env_cfg.py @@ -104,7 +104,7 @@ def __post_init__(self): class SubtaskCfg(ObsGroup): """Observations for subtask group.""" - grasp_1 = ObsTerm( + grasp = ObsTerm( func=mdp.object_grasped, params={ "robot_cfg": SceneEntityCfg("robot"), @@ -112,7 +112,7 @@ class SubtaskCfg(ObsGroup): "object_cfg": SceneEntityCfg("cube_2"), }, ) - stack_1 = ObsTerm( + stacked = ObsTerm( func=mdp.object_stacked, params={ "robot_cfg": SceneEntityCfg("robot"), @@ -120,14 +120,6 @@ class SubtaskCfg(ObsGroup): "lower_object_cfg": SceneEntityCfg("cube_1"), }, ) - grasp_2 = ObsTerm( - func=mdp.object_grasped, - params={ - "robot_cfg": SceneEntityCfg("robot"), - "ee_frame_cfg": SceneEntityCfg("ee_frame"), - "object_cfg": SceneEntityCfg("cube_3"), - }, - ) def __post_init__(self): self.enable_corruption = False @@ -153,12 +145,10 @@ class TerminationsCfg: func=mdp.root_height_below_minimum, params={"minimum_height": -0.05, "asset_cfg": SceneEntityCfg("cube_2")} ) - cube_3_dropping = DoneTerm( - func=mdp.root_height_below_minimum, params={"minimum_height": -0.05, "asset_cfg": SceneEntityCfg("cube_3")} + success = DoneTerm( + func=mdp.cubes_stacked ) - success = DoneTerm(func=mdp.cubes_stacked) - @configclass class StackEnvCfg(ManagerBasedRLEnvCfg):