diff --git a/.gitignore b/.gitignore index 0889c69a..f775f0fb 100644 --- a/.gitignore +++ b/.gitignore @@ -71,3 +71,6 @@ tests/ # Docker history .uw-lab-docker-history + +# local reset datasets +Datasets/ diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 00000000..af380be5 --- /dev/null +++ b/.gitmodules @@ -0,0 +1,3 @@ +[submodule "diffusion_policy"] + path = diffusion_policy + url = git@github.com:sriyash421/diffusion_policy.git diff --git a/README.md b/README.md index a97c857f..f1705aec 100644 --- a/README.md +++ b/README.md @@ -36,6 +36,7 @@ Follow the [installation guide](https://uw-lab.github.io/UWLab/main/source/setup - **Train Your First Policy** — Train an ant to run in minutes → [Quickstart](https://uw-lab.github.io/UWLab/main/source/setup/installation/pip_installation.html#train-a-robot) - **OmniReset** — RL for manipulation without reward engineering or demos → [Quickstart](https://uw-lab.github.io/UWLab/main/source/publications/omnireset/index.html#quick-start) +- **ASTEROID** — Iterative in-context exploration + distillation for cube pick-up → [Quickstart](https://uw-lab.github.io/UWLab/main/source/publications/asteroid/index.html#asteroid-quick-start) See [all available environments](https://uw-lab.github.io/UWLab/main/source/overview/uw_environments.html) and [full documentation](https://uw-lab.github.io/UWLab) for details. diff --git a/diffusion_policy b/diffusion_policy new file mode 160000 index 00000000..7e790f68 --- /dev/null +++ b/diffusion_policy @@ -0,0 +1 @@ +Subproject commit 7e790f681646e4bf1f3fabe0dafeacdd07add61e diff --git a/docs/index.rst b/docs/index.rst index 222a3971..eec37907 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -77,6 +77,7 @@ Table of Contents source/publications/pg1 source/publications/omnireset/index + source/publications/asteroid/index .. toctree:: :maxdepth: 3 diff --git a/docs/source/publications/asteroid/index.rst b/docs/source/publications/asteroid/index.rst new file mode 100644 index 00000000..6f9a1646 --- /dev/null +++ b/docs/source/publications/asteroid/index.rst @@ -0,0 +1,161 @@ +ASTEROID +======== + +| **Code:** ``scripts/ASTEROID`` and ``source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid`` + +ASTEROID trains a proprioceptive student policy for cube pick-up by iterating +**in-context exploration** and **distillation**: a state-based RL expert (trained with +the OmniReset recipe) supervises a diffusion-policy student, and from the second +iteration on the previous student acts as an *explorer* for the first part of each +data-collection episode before the expert takes over. + +Every iteration runs three stages: + +1. **collect** -- roll out the expert (plus the previous student as explorer) in the + data-collection env and record proprioceptive observations, actions and a per-step + expert mask (``scripts/ASTEROID/collect_demos_asteroid.py``). +2. **train** -- fit a diffusion-policy student on every dataset collected so far, with + a per-iteration sampling curriculum (``diffusion_policy/train.py``). +3. **eval** -- roll out the student in the eval env + (``scripts/ASTEROID/eval_asteroid_policy.py``). + +The orchestrator ``scripts/ASTEROID/run_asteroid.py`` organises hyperparameters as a +hierarchy of dataclasses: a ``RunCfg`` holds the run-level settings plus an ordered list +of ``IterationCfg``, each owning the ``CollectCfg`` / ``TrainCfg`` / ``EvalCfg`` for that +iteration. Curricula are functions that build the iteration list (``--schedule``). + +---- + +.. _asteroid-quick-start: + +Quick Start +----------- + +.. important:: + + Make sure you have completed the `installation `_ + before running these commands. The distillation stages additionally need the + ``diffusion_policy`` submodule. + +Environments +^^^^^^^^^^^^ + +All ASTEROID environments are pick-only variants of the OmniReset UR5e + Robotiq 2F-85 +tasks (no receptive object; success = object lifted with the gripper pointing down): + +.. list-table:: + :header-rows: 1 + :widths: 55 45 + + * - Task + - Purpose + * - ``Asteroid-UR5eRobotiq2f85-ObjectAnywhereEEAnywhere-v0`` + - Record reset states: object on the table, EE above it + * - ``Asteroid-UR5eRobotiq2f85-ObjectRestingEEGrasped-v0`` + - Record reset states: object resting, EE grasping it + * - ``Asteroid-UR5eRobotiq2f85-ObjectAnywhereEEGrasped-v0`` + - Record reset states: object anywhere, EE grasping it + * - ``Asteroid-Ur5eRobotiq2f85-RelCartesianOSC-State-v0`` + - Train the state expert (Stage 1) + * - ``Asteroid-Ur5eRobotiq2f85-RelCartesianOSC-State-Finetune-v0`` + - Finetune the expert with sysid / gain curriculum (Stage 2) + * - ``Asteroid-Ur5eRobotiq2f85-RelCartesianOSC-State-Play-v0`` + - Evaluate a Stage 1 expert + * - ``Asteroid-Ur5eRobotiq2f85-RelCartesianOSC-State-Finetune-Play-v0`` + - Evaluate a Stage 2 expert + * - ``Asteroid-Ur5eRobotiq2f85-RelCartesianOSC-Tactile-DataCollection-v0`` + - Collect student demos with a Stage 1 expert + * - ``Asteroid-Ur5eRobotiq2f85-RelCartesianOSC-Tactile-Finetune-DataCollection-v0`` + - Collect student demos with a Stage 2 expert + * - ``Asteroid-Ur5eRobotiq2f85-RelCartesianOSC-Tactile-Play-v0`` + - Evaluate a student (Stage 1 gains, front camera video) + * - ``Asteroid-Ur5eRobotiq2f85-RelCartesianOSC-Tactile-Finetune-Play-v0`` + - Evaluate a student (Stage 2 gains, front camera video) + +Reset-state and grasp datasets are read from a local directory (default ``Datasets/CubePick``, +override with the ``ASTEROID_DATASETS_DIR`` environment variable) keyed by the insertive +object only:: + + Datasets/CubePick/Resets/InsertiveCube/resets_ObjectAnywhereEEAnywhere.pt + Datasets/CubePick/Resets/InsertiveCube/resets_ObjectRestingEEGrasped.pt + Datasets/CubePick/Resets/InsertiveCube/resets_ObjectAnywhereEEGrasped.pt + Datasets/CubePick/Grasps/InsertiveCube/grasps.pt + +1. Record reset states +^^^^^^^^^^^^^^^^^^^^^^ + +.. code:: bash + + python scripts_v2/tools/record_reset_states.py \ + --task Asteroid-UR5eRobotiq2f85-ObjectAnywhereEEAnywhere-v0 \ + --dataset_dir Datasets/CubePick \ + --num_envs 64 --num_reset_states 1000 --headless \ + env.scene.insertive_object=cube + +Repeat for ``ObjectRestingEEGrasped`` and ``ObjectAnywhereEEGrasped`` (these two need the +``ObjectAnywhereEEAnywhere`` resets and a grasp dataset; see the OmniReset +:doc:`../omnireset/rl_training` page for grasp sampling). + +2. Train the state expert +^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. code:: bash + + python scripts/reinforcement_learning/rsl_rl/train.py \ + --task Asteroid-Ur5eRobotiq2f85-RelCartesianOSC-State-v0 \ + --num_envs 4096 --headless \ + env.scene.insertive_object=cube + + # evaluate + python scripts/reinforcement_learning/rsl_rl/play.py \ + --task Asteroid-Ur5eRobotiq2f85-RelCartesianOSC-State-Play-v0 \ + --num_envs 1 --checkpoint logs/rsl_rl/ur5e_robotiq_2f85_asteroid_agent//model_.pt \ + env.scene.insertive_object=cube + +Export the expert to TorchScript (``logs/rsl_rl/.../exported/policy.pt``) as in the OmniReset +:doc:`../omnireset/distillation` page; the data-collection stage loads it with ``torch.jit.load``. + +3. Run ASTEROID +^^^^^^^^^^^^^^^ + +.. code:: bash + + python scripts/ASTEROID/run_asteroid.py \ + --data_task Asteroid-Ur5eRobotiq2f85-RelCartesianOSC-Tactile-DataCollection-v0 \ + --eval_task Asteroid-Ur5eRobotiq2f85-RelCartesianOSC-Tactile-Play-v0 \ + --expert_policy_checkpoint logs/rsl_rl/ur5e_robotiq_2f85_asteroid_agent//exported/policy.pt \ + --config_name in_context_exploration_tactile_base.yaml \ + --num_demos 32768 --num_data_envs 512 \ + --num_eval_envs 32 --num_eval_episodes 100 \ + --max_iterations 4 --exp_name cube_asteroid --no_video + +Useful flags: + +- ``--dry_run`` prints the stage commands without launching Isaac Sim. +- ``--schedule`` selects a curriculum from ``CURRICULA`` in ``run_asteroid.py``. +- ``--start_iteration N --checkpoint_dir `` resumes an interrupted run. +- ``--initial_dataset_path`` reuses an existing iteration-0 dataset. + +Each run writes ``run_cfg.json`` (the full hyperparameter tree), one +``dataset-iteration-{i}/`` and one ``iteration_{i}/`` (student checkpoints) per iteration. + +---- + +Package layout +-------------- + +``uwlab_tasks.manager_based.manipulation.asteroid`` mirrors ``omnireset`` and subclasses +it; only the pick-specific deltas live here: + +- ``mdp/commands*.py`` -- ``PickTaskCommand``: task command without a receptive object. +- ``mdp/rewards.py`` -- ``ProgressContextPickOnly`` (lift height + gripper-down success) and + the matching dense / sparse rewards. +- ``mdp/events.py`` -- ``SingleObjectMultiResetManager`` (resets keyed by one object), + ``randomize_env_cfg_unified`` (coupled sysid / OSC-gain / action-scale DR), + ``randomize_gripper_pos_affine`` (gripper-reading calibration drift), + ``reset_root_states_discrete_grid``. +- ``mdp/observations.py`` -- ``gripper_pos_normalized`` (real-robot POS register analogue), + ``fingertip_contact_force_b``. +- ``mdp/recorders/`` -- per-step expert mask recorder for DAgger-style datasets. +- ``mdp/actions/`` -- position-only (3-DOF + gripper) Cartesian OSC action. +- ``config/ur5e_robotiq_2f85/`` -- reset-state, RL-state and tactile data-collection configs. diff --git a/scripts/ASTEROID/collect_cubepick_expert5.sh b/scripts/ASTEROID/collect_cubepick_expert5.sh new file mode 100755 index 00000000..197f6b8f --- /dev/null +++ b/scripts/ASTEROID/collect_cubepick_expert5.sh @@ -0,0 +1,52 @@ +#!/bin/sh + +ckpts=( + model_voff_curr + # model_voff_nocurr + model_von_curr + # model_von_nocurr +) + +envs=( + Asteroid-Ur5eRobotiq2f85-RelCartesianOSC-Tactile-Finetune-DataCollection-v0 + # Asteroid-Ur5eRobotiq2f85-RelCartesianOSC-Tactile-DataCollection-v0 + Asteroid-Ur5eRobotiq2f85-RelCartesianOSC-Tactile-Finetune-DataCollection-v0 + # Asteroid-Ur5eRobotiq2f85-RelCartesianOSC-Tactile-DataCollection-v0 +) + +for i in "${!ckpts[@]}"; do + ckpt="${ckpts[$i]}" + env="${envs[$i]}" + python scripts/ASTEROID/collect_demos_asteroid.py \ + --task $env \ + --dataset_file "logs/debug/dataset-iteration-0-$ckpt/data.zarr" \ + --num_envs 1 \ + --num_demos 10 \ + --headless \ + --seed 0 \ + --min_exploration_horizon 0.0 \ + --max_exploration_horizon 0.0 \ + --episode_length_s 10.0 \ + --expert_noise 0.0 \ + --video \ + --video_length 2000 \ + --video_dir "logs/$ckpt/debug_video" \ + env.scene.insertive_object=cube \ + agent.algorithm.offline_algorithm_cfg.behavior_cloning_cfg.experts_path=["logs/$ckpt/exported/policy.pt"] +done +# python scripts/ASTEROID/collect_demos_asteroid.py \ +# --task Asteroid-Ur5eRobotiq2f85-RelCartesianOSC-Tactile-DataCollection-v0 \ +# --dataset_file "logs/dataset-iteration-0-debug/data.zarr" \ +# --num_envs 1 \ +# --num_demos 5 \ +# --headless \ +# --seed 0 \ +# --min_exploration_horizon 0.0 \ +# --max_exploration_horizon 0.0 \ +# --episode_length_s 10.0 \ +# --expert_noise 0.0 \ +# --video \ +# --video_length 2000 \ +# --video_dir "logs/debug_videos" \ +# env.scene.insertive_object=cube \ +# 'agent.algorithm.offline_algorithm_cfg.behavior_cloning_cfg.experts_path=["logs/exported/policy.pt"]' \ No newline at end of file diff --git a/scripts/ASTEROID/collect_demos_asteroid.py b/scripts/ASTEROID/collect_demos_asteroid.py new file mode 100644 index 00000000..730d6dac --- /dev/null +++ b/scripts/ASTEROID/collect_demos_asteroid.py @@ -0,0 +1,583 @@ +# Copyright (c) 2024-2025, The Octi Lab Project Developers. +# Proprietary and Confidential - All Rights Reserved. +# +# Unauthorized copying of this file, via any medium is strictly prohibited + +"""Script to collect demonstrations using exploration policy + expert policy.""" + +"""Launch Isaac Sim Simulator first.""" + +import argparse +import contextlib +import os +import gymnasium as gym +import torch +from tqdm import tqdm +from typing import Sequence + +from isaaclab.app import AppLauncher + + +# add argparse arguments +parser = argparse.ArgumentParser(description="Collect demonstrations with exploration + expert policy.") +parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.") +parser.add_argument("--task", type=str, default=None, help="Name of the task.") +parser.add_argument( + "--dataset_file", type=str, default="./datasets/dataset.zarr", help="File path to export recorded demos." +) +parser.add_argument( + "--num_demos", type=int, default=10, help="Number of demonstrations to record." +) +parser.add_argument( + "--exploration_checkpoint", type=str, default=None, help="Path to exploration policy checkpoint." +) +parser.add_argument( + "--min_exploration_horizon", type=float, default=0.02, help="Minimum exploration horizon for exploration policy." +) +parser.add_argument( + "--max_exploration_horizon", type=float, default=0.3, help="Maximum exploration horizon for exploration policy." +) +parser.add_argument( + "--episode_length_s", type=float, default=5.0, help="Episode length in seconds." +) +parser.add_argument("--expert_noise", type=float, default=0.0, help="Noise level for the expert policy.") + +parser.add_argument("--render", action="store_true", help="Whether to render the environment.") +parser.add_argument("--seed", type=int, default=0, help="Random seed for environment.") +parser.add_argument("--video", action="store_true", default=False, help="Record video of env 0 during collection.") +parser.add_argument("--video_length", type=int, default=400, help="Video length in env steps (only used when --video).") +parser.add_argument("--video_interval", type=int, default=0, + help="Restart video recording every N env steps. 0 (default) = single video at step 0.") +parser.add_argument("--video_dir", type=str, default=None, + help="Directory to write videos. Defaults to _dir/videos.") + +# append AppLauncher cli args +AppLauncher.add_app_launcher_args(parser) +args_cli, remaining_args = parser.parse_known_args() + +# Mirror play.py: cameras must be enabled for the RecordVideo wrapper to render. +if args_cli.video: + args_cli.enable_cameras = True + +# launch omniverse app +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import dill +import hydra + +from isaaclab.envs import ( + DirectRLEnvCfg, + ManagerBasedRLEnvCfg +) + +from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlVecEnvWrapper + +import isaaclab_tasks # noqa: F401 +import uwlab_tasks # noqa: F401 + +# Import dataset handlers +from isaaclab.utils.datasets import HDF5DatasetFileHandler +from isaaclab.managers.recorder_manager import DatasetExportMode + +from uwlab.utils.datasets import ZarrDatasetFileHandler +from uwlab_tasks.utils.hydra import hydra_task_compose +from uwlab_tasks.manager_based.manipulation.asteroid.mdp.recorders.recorders_cfg import ( + AsteroidActionStateRecorderManagerCfg as ActionStateRecorderManagerCfg, +) + +# Diffusion policy imports +from diffusion_policy.workspace.base_workspace import BaseWorkspace +from diffusion_policy.policy.base_image_policy import BaseImagePolicy +from uwlab_rl.wrappers.diffusion import DiffusionPolicyWrapper + +torch.backends.cuda.matmul.allow_tf32 = True +torch.backends.cudnn.allow_tf32 = True +torch.backends.cudnn.deterministic = False +torch.backends.cudnn.benchmark = False + +def record_pre_reset(self, env_ids: Sequence[int] | None, force_export_or_skip=None) -> None: + """Trigger recorder terms for pre-reset functions. + + Args: + env_ids: The environment ids in which a reset is triggered. + """ + # Do nothing if no active recorder terms are provided + if len(self.active_terms) == 0: + return + + if env_ids is None: + env_ids = list(range(self._env.num_envs)) + if isinstance(env_ids, torch.Tensor): + env_ids = env_ids.tolist() + + for term in self._terms.values(): + key, value = term.record_pre_reset(env_ids) + self.add_to_episodes(key, value, env_ids) + + # Set task success values for the relevant episodes + success_results = torch.zeros(len(env_ids), dtype=bool, device=self._env.device) + # Check success indicator from termination terms + if hasattr(self._env, "termination_manager"): + if "success" in self._env.termination_manager.active_terms: + success_results |= self._env.termination_manager.get_term("success")[env_ids] + + # Check episode length condition for success + episode_lengths = self._env.episode_length_buf[env_ids] + exploration_lengths = self.exploration_lengths[env_ids] + exploration_ratios = exploration_lengths / episode_lengths + max_exploration_horizon_for_save = 0.9 + success_results = success_results & (exploration_ratios < max_exploration_horizon_for_save) + self.set_success_to_episodes(env_ids, success_results) + + if force_export_or_skip or (force_export_or_skip is None and self.cfg.export_in_record_pre_reset): + self.export_episodes(env_ids) + +def process_agent_cfg(env_cfg, agent_cfg): + if hasattr(agent_cfg.algorithm, "behavior_cloning_cfg"): + if agent_cfg.algorithm.behavior_cloning_cfg is None: + del agent_cfg.algorithm.behavior_cloning_cfg + else: + bc_cfg = agent_cfg.algorithm.behavior_cloning_cfg + if bc_cfg.experts_observation_group_cfg is not None: + import importlib + + # resolve path to the module location + mod_name, attr_name = bc_cfg.experts_observation_group_cfg.split(":") + mod = importlib.import_module(mod_name) + cfg_cls = mod + for attr in attr_name.split("."): + cfg_cls = getattr(cfg_cls, attr) + cfg = cfg_cls() + setattr(env_cfg.observations, "expert_obs", cfg) + + if hasattr(agent_cfg.algorithm, "offline_algorithm_cfg"): + if agent_cfg.algorithm.offline_algorithm_cfg is None: + del agent_cfg.algorithm.offline_algorithm_cfg + else: + if agent_cfg.algorithm.offline_algorithm_cfg.behavior_cloning_cfg is None: + del agent_cfg.algorithm.offline_algorithm_cfg.behavior_cloning_cfg + else: + bc_cfg = agent_cfg.algorithm.offline_algorithm_cfg.behavior_cloning_cfg + if bc_cfg.experts_observation_group_cfg is not None: + import importlib + + # resolve path to the module location + mod_name, attr_name = bc_cfg.experts_observation_group_cfg.split(":") + mod = importlib.import_module(mod_name) + cfg_cls = mod + for attr in attr_name.split("."): + cfg_cls = getattr(cfg_cls, attr) + cfg = cfg_cls() + setattr(env_cfg.observations, "expert_obs", cfg) + return agent_cfg + + +def load_exploration_policy(checkpoint_path: str, device: torch.device, num_envs: int): + """Load exploration diffusion policy from checkpoint.""" + payload = torch.load(open(checkpoint_path, 'rb'), pickle_module=dill) + cfg = payload['cfg'] + cls = hydra.utils.get_class(cfg._target_) + workspace = cls(cfg) + workspace: BaseWorkspace + workspace.load_payload(payload, exclude_keys=None, include_keys=None) + + policy: BaseImagePolicy + policy = workspace.model + if cfg.training.use_ema: + policy = workspace.ema_model + + policy.eval().to(device) + + wrapped_policy = DiffusionPolicyWrapper(policy, device, n_obs_steps=policy.n_obs_steps, num_envs=num_envs) + return wrapped_policy + + +def sample_exploration_horizons(num_envs: int, min_horizon: int, max_horizon: int, device: torch.device) -> torch.Tensor: + """Sample exploration horizons for each environment.""" + return torch.randint(min_horizon, max_horizon + 1, (num_envs,), device=device) + + +@hydra_task_compose(args_cli.task, "rsl_rl_cfg_entry_point", hydra_args=remaining_args) +def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg, agent_cfg: RslRlOnPolicyRunnerCfg): + """Collect demonstrations from the environment using exploration + expert policy.""" + # device = torch.device(args_cli.device if args_cli.device else 'cuda' if torch.cuda.is_available() else 'cpu') + device = torch.device('cuda:0') + policy_device = torch.device('cuda:1') if torch.cuda.device_count() > 1 else device + + # get directory path and file name (without extension) from cli arguments + output_dir = os.path.dirname(args_cli.dataset_file) + output_file_name = os.path.basename(args_cli.dataset_file) + + # create directory if it does not exist + if not os.path.exists(output_dir): + os.makedirs(output_dir, exist_ok=True) + + # add recordermanager to save data + use_zarr_format = args_cli.dataset_file.endswith('.zarr') + if use_zarr_format: + dataset_handler = ZarrDatasetFileHandler + else: + dataset_handler = HDF5DatasetFileHandler + + env_cfg.recorders = ActionStateRecorderManagerCfg() + env_cfg.recorders.dataset_export_dir_path = output_dir + env_cfg.recorders.dataset_filename = output_file_name + env_cfg.recorders.dataset_export_mode = DatasetExportMode.EXPORT_SUCCEEDED_ONLY + env_cfg.recorders.dataset_file_handler_class_type = dataset_handler + + # override configurations with non-hydra CLI arguments + env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs + env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device + env_cfg.seed = None + # we want to have the terms in the observations returned as a dictionary + env_cfg.observations.policy.concatenate_terms = False + + # set episode length based on cli argument + env_cfg.episode_length_s = args_cli.episode_length_s + episode_length = int(env_cfg.episode_length_s / (env_cfg.sim.dt * env_cfg.sim.render_interval)) + max_exploration_horizon = int(args_cli.max_exploration_horizon * episode_length) + min_exploration_horizon = int(args_cli.min_exploration_horizon * episode_length) + + # add expert obs into env_cfg + agent_cfg = process_agent_cfg(env_cfg, agent_cfg) + + # create isaac environment + env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array") + + # Optional video recording (mirrors scripts/reinforcement_learning/rsl_rl/play.py). + # Wrapping must happen *before* RslRlVecEnvWrapper so RecordVideo sees env.step() + # at the right level. RecordVideo only renders env 0 of the vec env. + if args_cli.video: + if args_cli.video_dir is None: + args_cli.video_dir = os.path.join( + os.path.dirname(args_cli.dataset_file), "videos" + ) + os.makedirs(args_cli.video_dir, exist_ok=True) + if args_cli.video_interval and args_cli.video_interval > 0: + video_kwargs = { + "video_folder": args_cli.video_dir, + "step_trigger": lambda step: step % args_cli.video_interval == 0, + "video_length": args_cli.video_length, + "disable_logger": True, + } + else: + video_kwargs = { + "video_folder": args_cli.video_dir, + "step_trigger": lambda step: step == 0, + "video_length": args_cli.video_length, + "disable_logger": True, + } + print(f"[collect_demos] Recording video to {args_cli.video_dir}") + env = gym.wrappers.RecordVideo(env, **video_kwargs) + + # wrap around environment for rsl-rl + env = RslRlVecEnvWrapper(env) + num_envs = env.num_envs + + # load exploration policy + exploration_policy = None + if args_cli.exploration_checkpoint: + exploration_policy = load_exploration_policy(args_cli.exploration_checkpoint, policy_device, num_envs) + + # load expert policy + bc = agent_cfg.algorithm.offline_algorithm_cfg.behavior_cloning_cfg + assert len(bc.experts_path) == 1, "Only one expert is supported for now." + expert_obs_fn = bc.experts_observation_func + loader = bc.experts_loader + if not callable(loader): + loader = eval(loader) + expert_policy = loader(bc.experts_path[0]).to(device) + expert_policy.eval() + + # get expert mask recorder + recorder_manager = env.unwrapped.recorder_manager + expert_mask_recorder = recorder_manager._terms.get("record_pre_step_expert_mask") + + # initialize exploration horizons for each env + exploration_horizons = sample_exploration_horizons( + num_envs, min_exploration_horizon, max_exploration_horizon, device + ) + expert_mask_recorder.set_exploration_horizon(exploration_horizons) + + if exploration_policy is not None: + from types import MethodType + recorder_manager.record_pre_reset = MethodType(record_pre_reset, recorder_manager) + + exploration_lengths = torch.zeros((num_envs,), device=device, dtype=torch.int32) + recorder_manager.exploration_lengths = exploration_lengths + + # reset environment and get initial observations + obs_dict, _ = env.reset() + + # reset exploration policy + if exploration_policy is not None: + reset_ids = torch.arange(num_envs, device=device) + exploration_policy.reset(reset_ids) + + rendered_frames = [] + # simulate environment -- run everything in inference mode + current_recorded_demo_count = 0 + # ---- debug instrumentation ---- + import time as _time_dbg + _dbg_step = 0 + _dbg_last_demo_count = 0 + _dbg_last_demo_step = 0 + _dbg_last_demo_wall = _time_dbg.time() + _dbg_step_times: list[float] = [] # rolling step latencies + _dbg_total_dones = 0 + _dbg_total_successes = 0 + _dbg_window_steps = 0 + _dbg_window_dones = 0 + _dbg_window_successes = 0 + _dbg_window_wall = _time_dbg.time() + DBG_PRINT_EVERY = 50 # env steps + DBG_NO_PROGRESS_S = 30.0 # alert when no new demo for this many seconds + # ---- stall recovery (auto-retry inside this process) ---- + STALL_RESET_AFTER_S = 90.0 # full env.reset() if no demo progress for this long + STALL_GIVEUP_AFTER = 3 # break after this many resets in a row that don't help + _stall_reset_count = 0 + _last_stall_reset_demo_count = 0 + _last_stall_reset_step = 0 + # CSV trace of action stats per dbg interval — offline-plottable. + _dbg_action_csv = os.path.join( + os.path.dirname(args_cli.dataset_file), "action_stats.csv" + ) + _dbg_csv_header_written = False + # -------------------------------- + with contextlib.suppress(KeyboardInterrupt) and torch.inference_mode(): + pbar = tqdm(total=args_cli.num_demos, desc="Recording Demonstrations", unit="demo") + + while True: + # get current episode step for each env + episode_steps = env.unwrapped.episode_length_buf + + # determine which policy to use for each env + use_exploration = (episode_steps < exploration_horizons) & (exploration_policy is not None) + exploration_lengths += use_exploration.int() + recorder_manager.exploration_lengths = exploration_lengths + use_expert = ~use_exploration + + # compute actions + actions = torch.zeros((num_envs, env.action_space.shape[-1]), device=device) + expert_mask = torch.zeros((num_envs, 1), dtype=torch.bool, device=device) + # exploration policy actions + if use_exploration.any() and exploration_policy is not None: + exploration_env_ids = use_exploration.nonzero(as_tuple=False).reshape(-1) + exploration_obs = {k: v[use_exploration] for k, v in obs_dict['policy'].items()} + exploration_actions = exploration_policy.predict_action(exploration_obs, exploration_env_ids) + exploration_actions = exploration_actions.to(device) + actions[use_exploration] = exploration_actions + + # expert policy actions + if use_expert.any(): + expert_policy_obs = expert_obs_fn(env) + if isinstance(expert_policy_obs, dict): + expert_policy_obs = torch.cat([v for v in expert_policy_obs.values()], dim=-1) + mean, std = expert_policy.compute_distribution(expert_policy_obs) + expert_actions = torch.normal(mean, std) + if args_cli.expert_noise > 0.0: + noise = torch.randn_like(expert_actions) * args_cli.expert_noise + expert_actions += noise + actions[use_expert] = expert_actions[use_expert] + expert_mask[use_expert] = True + + # mask actions for first step after reset (first image not valid) + first_step_mask = (episode_steps == 0) + if torch.any(first_step_mask): + actions[first_step_mask, :-1] = 0.0 + actions[first_step_mask, -1] = -1.0 # close gripper + + # set expert mask for recorder + expert_mask_recorder.set_mask(expert_mask) + + # env stepping + _dbg_t0 = _time_dbg.time() + obs_dict, _, dones, _ = env.step(actions) + _dbg_step_dt = _time_dbg.time() - _dbg_t0 + _dbg_step_times.append(_dbg_step_dt) + if len(_dbg_step_times) > 100: + _dbg_step_times = _dbg_step_times[-100:] + + # handle resets + if dones.any(): + reset_ids = (dones > 0).nonzero(as_tuple=False).reshape(-1) + # resample exploration horizons for reset envs + exploration_horizons[reset_ids] = sample_exploration_horizons( + len(reset_ids), min_exploration_horizon, max_exploration_horizon, device + ) + if exploration_policy is not None: + exploration_policy.reset(reset_ids) + + exploration_lengths[reset_ids] = 0 + # expert_mask_recorder.set_exploration_horizon(exploration_horizons) + + # update demo count + new_count = env.unwrapped.recorder_manager.exported_successful_episode_count + if new_count > current_recorded_demo_count: + increment = new_count - current_recorded_demo_count + current_recorded_demo_count = new_count + _dbg_window_successes += increment + _dbg_last_demo_step = _dbg_step + _dbg_last_demo_wall = _time_dbg.time() + pbar.update(increment) + + # ---- debug telemetry ---- + _dbg_step += 1 + _dbg_window_steps += 1 + _dbg_n_dones = int(dones.sum().item()) + _dbg_window_dones += _dbg_n_dones + _dbg_total_dones += _dbg_n_dones + + if _dbg_step % DBG_PRINT_EVERY == 0: + _now = _time_dbg.time() + wall_dt = max(_now - _dbg_window_wall, 1e-9) + ep_lens = env.unwrapped.episode_length_buf + ep_min, ep_max, ep_mean = ( + int(ep_lens.min().item()), + int(ep_lens.max().item()), + float(ep_lens.float().mean().item()), + ) + # time_out_buf / terminated_buf if available + tm = env.unwrapped.termination_manager + _to = int(tm.time_outs.sum().item()) if hasattr(tm, "time_outs") else -1 + _term = int(tm.dones.sum().item()) if hasattr(tm, "dones") else -1 + step_med = sorted(_dbg_step_times)[len(_dbg_step_times) // 2] if _dbg_step_times else 0.0 + step_max = max(_dbg_step_times) if _dbg_step_times else 0.0 + no_progress_s = _now - _dbg_last_demo_wall + rate = _dbg_window_successes / wall_dt + done_rate = _dbg_window_dones / wall_dt + # action stats — track per-dim mean / std / saturation. + _act_dim = actions.shape[-1] + _act_mean = actions.mean(dim=0).detach().cpu().tolist() + _act_std = actions.std(dim=0).detach().cpu().tolist() + _act_norm = actions.norm(dim=-1).detach().cpu() + _act_norm_med = float(_act_norm.median().item()) + _act_norm_max = float(_act_norm.max().item()) + # how often does any action dim hit ±1 (the typical clip bound)? + _act_saturated_frac = float(((actions.abs() >= 0.999).any(dim=-1)).float().mean().item()) + print( + f"[dbg] step={_dbg_step} demos={current_recorded_demo_count} " + f"window: dones={_dbg_window_dones} successes={_dbg_window_successes} " + f"({done_rate:.1f} dones/s, {rate:.2f} success/s) | " + f"ep_len min/mean/max={ep_min}/{ep_mean:.1f}/{ep_max} | " + f"step_dt med/max={step_med*1000:.0f}/{step_max*1000:.0f}ms | " + f"this-step time_outs={_to} dones_in_buf={_term} | " + f"act |a| med/max={_act_norm_med:.3f}/{_act_norm_max:.3f} sat={_act_saturated_frac:.2f} | " + f"no_progress={no_progress_s:.0f}s", + flush=True, + ) + # CSV row: step,demos,no_progress_s,act_mean[0..D],act_std[0..D],act_norm_median,act_saturated_frac + try: + if not _dbg_csv_header_written: + os.makedirs(os.path.dirname(_dbg_action_csv), exist_ok=True) + with open(_dbg_action_csv, "w") as _f: + cols = ( + ["step", "demos", "no_progress_s"] + + [f"act_mean_{i}" for i in range(_act_dim)] + + [f"act_std_{i}" for i in range(_act_dim)] + + ["act_norm_med", "act_norm_max", "act_sat_frac"] + ) + _f.write(",".join(cols) + "\n") + _dbg_csv_header_written = True + with open(_dbg_action_csv, "a") as _f: + row = ( + [str(_dbg_step), str(current_recorded_demo_count), f"{no_progress_s:.1f}"] + + [f"{x:.5f}" for x in _act_mean] + + [f"{x:.5f}" for x in _act_std] + + [f"{_act_norm_med:.5f}", f"{_act_norm_max:.5f}", f"{_act_saturated_frac:.4f}"] + ) + _f.write(",".join(row) + "\n") + except Exception as _e: + print(f"[dbg] CSV write failed: {_e}", flush=True) + if no_progress_s > DBG_NO_PROGRESS_S: + # diagnose stall: split between physics-step-stuck vs success-rate=0 + if step_max > 5.0: + print(f"[dbg] STALL CAUSE: step latency exploded ({step_max*1000:.0f} ms). " + "physx solver likely hung on at least one env.", flush=True) + elif _dbg_window_dones == 0 and ep_max > 0: + print(f"[dbg] STALL CAUSE: env steps fast ({step_med*1000:.0f} ms median) " + "but ZERO terminations in window. time_out / abnormal / success " + "all stopped firing — termination buffers may be stuck.", flush=True) + elif _dbg_window_dones > 0 and _dbg_window_successes == 0: + print(f"[dbg] STALL CAUSE: episodes terminating ({_dbg_window_dones} in window) " + "but zero successes — all attempts failing. EXPORT_SUCCEEDED_ONLY " + "drops failures so demo count can't advance.", flush=True) + else: + print(f"[dbg] STALL CAUSE: indeterminate. step_med={step_med*1000:.0f}ms " + f"dones_window={_dbg_window_dones} successes_window={_dbg_window_successes}", + flush=True) + _dbg_window_steps = 0 + _dbg_window_dones = 0 + _dbg_window_successes = 0 + _dbg_window_wall = _now + + # ---- stall recovery: full env.reset() to clear physx state and + # re-sample reset states; break only if multiple consecutive + # resets fail to unstick collection. + if no_progress_s > STALL_RESET_AFTER_S: + if current_recorded_demo_count > _last_stall_reset_demo_count: + # last reset DID help (we got new demos since); reset counter + _stall_reset_count = 0 + _stall_reset_count += 1 + if _stall_reset_count > STALL_GIVEUP_AFTER: + print( + f"[dbg] STALL GIVEUP: {_stall_reset_count} consecutive " + f"force-resets did not unstick collection. Saving " + f"{current_recorded_demo_count} demos and exiting cleanly.", + flush=True, + ) + break + print( + f"[dbg] STALL RECOVERY: forcing env.reset() (attempt " + f"{_stall_reset_count}/{STALL_GIVEUP_AFTER}) at " + f"{current_recorded_demo_count} demos to clear physx state.", + flush=True, + ) + obs_dict, _ = env.reset() + if exploration_policy is not None: + exploration_policy.reset(torch.arange(num_envs, device=device)) + exploration_lengths.zero_() + exploration_horizons[:] = sample_exploration_horizons( + num_envs, min_exploration_horizon, max_exploration_horizon, device + ) + _last_stall_reset_demo_count = current_recorded_demo_count + _last_stall_reset_step = _dbg_step + _dbg_last_demo_wall = _time_dbg.time() # don't immediately re-trigger + # ------------------------- + + if args_cli.num_demos > 0 and new_count >= args_cli.num_demos: + print(f"All {args_cli.num_demos} demonstrations recorded. Exiting the app.") + break + + # check that simulation is stopped or not + if env.unwrapped.sim.is_stopped(): + break + + if args_cli.render: + rendered_frames.append(env.render()) + + pbar.close() + + if args_cli.render: + # save rendered frames as a video + import cv2 + video_path = os.path.join(output_dir, "recorded_demos_video.mp4") + height, width, _ = rendered_frames[0].shape + video_writer = cv2.VideoWriter(video_path, cv2.VideoWriter_fourcc(*'mp4v'), 30, (width, height)) + for frame in rendered_frames: + frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) + video_writer.write(frame_bgr) + video_writer.release() + print(f"Saved recorded demonstrations video at: {video_path}") + # close the simulator + env.close() + + +if __name__ == "__main__": + # run the main function - the decorator handles parameter passing + main() # type: ignore + # close sim app + simulation_app.close() diff --git a/scripts/ASTEROID/eval_asteroid_policy.py b/scripts/ASTEROID/eval_asteroid_policy.py new file mode 100644 index 00000000..19ccf057 --- /dev/null +++ b/scripts/ASTEROID/eval_asteroid_policy.py @@ -0,0 +1,348 @@ +# Copyright (c) 2022-2024, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Script to run a trained diffusion policy.""" + +"""Launch Isaac Sim Simulator first.""" + +import argparse + +from isaaclab.app import AppLauncher + +# add argparse arguments +parser = argparse.ArgumentParser(description="Play policy trained using diffusion policy for Isaac Lab environments.") +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="Path to diffusion policy checkpoint.") +parser.add_argument("--num_envs", type=int, default=1, help="Number of environments to run in parallel.") +parser.add_argument("--num_trajectories", type=int, default=100, help="Number of trajectories to evaluate. If None, run until simulation is stopped.") +parser.add_argument("--seed", type=int, default=42, help="Random seed for reproducibility.") +parser.add_argument("--use_amp", action="store_true", default=False, help="Use automatic mixed precision.") +parser.add_argument("--save_video", action="store_true", default=False, help="Save video of the policy.") +parser.add_argument("--episode_length_s", type=float, default=24.0, help="Episode length in seconds.") +parser.add_argument("--exp_name", type=str, default="diffusion_policy_eval", help="Experiment name for logging.") +parser.add_argument("--wandb_project", type=str, default="diffusion_policy_eval_new", help="WandB project name for logging.") +parser.add_argument("--wandb_group", type=str, default="default_group", help="WandB group name for logging.") +parser.add_argument("--iteration", type=int, default=0, help="Iteration number for logging.") +# append AppLauncher cli args +AppLauncher.add_app_launcher_args(parser) +# parse the arguments +args_cli, remaining_args = parser.parse_known_args() + +# launch omniverse app +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +"""Rest everything follows.""" + +import gymnasium as gym +import torch +import dill +import hydra +from contextlib import nullcontext +from tqdm import tqdm +import random +import numpy as np +import wandb + +import isaaclab_tasks # noqa: F401 +import uwlab_tasks # noqa: F401 +from isaaclab.envs import DirectRLEnvCfg, ManagerBasedRLEnvCfg +from uwlab_tasks.utils.hydra import hydra_task_compose + +# Diffusion policy imports +from diffusion_policy.workspace.base_workspace import BaseWorkspace +from diffusion_policy.policy.base_image_policy import BaseImagePolicy + +# Import the Diffusion policy wrapper +from uwlab_rl.wrappers.diffusion import DiffusionPolicyWrapper + +# import imageio + + +@hydra_task_compose(args_cli.task, "env_cfg_entry_point", hydra_args=remaining_args) +def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg, agent_cfg): + """Run a trained diffusion policy with Isaac Lab environment.""" + # Set seeds for reproducibility + random.seed(args_cli.seed) + np.random.seed(args_cli.seed) + torch.manual_seed(args_cli.seed) + torch.cuda.manual_seed_all(args_cli.seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + + # Check device is available + device = torch.device("cuda:0")#torch.device(args_cli.device if args_cli.device else 'cuda' if torch.cuda.is_available() else 'cpu') + policy_device = torch.device("cuda:1") if torch.cuda.device_count() > 1 else device + torch.backends.cudnn.benchmark = True + torch.backends.cuda.matmul.allow_tf32 = True + + # Override configurations with CLI arguments + env_cfg.scene.num_envs = args_cli.num_envs + env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device + env_cfg.sim.use_fabric = not args_cli.disable_fabric + # Set environment seed + env_cfg.seed = args_cli.seed + # we want to have the terms in the observations returned as a dictionary + # rather than a concatenated tensor + env_cfg.observations.policy.concatenate_terms = False + + # create environment + env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array") + + # ckpt_base = args_cli.checkpoint.split("/")[-1].split(".")[0] if args_cli.checkpoint else "none" + wandb.init( + project=args_cli.wandb_project, + group=args_cli.wandb_group, + config={ + "task": args_cli.task, + "checkpoint": args_cli.checkpoint, + "num_envs": args_cli.num_envs, + "num_trajectories": args_cli.num_trajectories, + "seed": args_cli.seed, + "use_amp": args_cli.use_amp, + "save_video": args_cli.save_video, + "episode_length_s": args_cli.episode_length_s, + "exp_name": args_cli.exp_name, + "iteration": args_cli.iteration, + }, + ) + + # Load diffusion policy checkpoint + ckpt_path = args_cli.checkpoint + payload = torch.load(open(ckpt_path, 'rb'), pickle_module=dill) + cfg = payload['cfg'] + cls = hydra.utils.get_class(cfg._target_) + workspace = cls(cfg) + workspace: BaseWorkspace + workspace.load_payload(payload, exclude_keys=None, include_keys=None) + + # Load policy based on configuration + policy: BaseImagePolicy + policy = workspace.model + if cfg.training.use_ema: + policy = workspace.ema_model + + # print policy summary + print(policy) + + policy.eval().to(policy_device) + + # Wrap policy to handle Isaac Lab observations + wrapped_policy = DiffusionPolicyWrapper(policy, policy_device, n_obs_steps=policy.n_obs_steps, num_envs=args_cli.num_envs) + + # reset environment + obs_dict, _ = env.reset() + dones = torch.ones(args_cli.num_envs, dtype=torch.bool, device=device) + reset_ids = (dones > 0).nonzero(as_tuple=False).reshape(-1) + wrapped_policy.reset(reset_ids) + + # Get termination term names to identify success + term_names = env.unwrapped.termination_manager._term_names # type: ignore + assert "success" in term_names, "Success term not found in termination manager" + + episodes = 0 + steps = 0 + successful_episodes = 0 # Track successful episodes + + # Track all episode metrics + episode_metrics = {} + + ep_returns = torch.zeros(args_cli.num_envs, dtype=torch.float32, device=device) + success_video_count = 0 + fail_video_count = 0 + render_frames = [] + successes = [] + + # Initialize progress bar if num_trajectories is specified + pbar = None + if args_cli.num_trajectories is not None: + pbar = tqdm(total=args_cli.num_trajectories, desc="Evaluating trajectories (Success: 0.00%)") + + # simulate environment + if args_cli.save_video: + env_frames = [[] for _ in range(args_cli.num_envs)] + frames_to_save = [] + cam_keys = sorted([key for key in obs_dict['policy'].keys() if 'rgb' in key]) + + while simulation_app.is_running(): + # Check if we've reached the desired number of trajectories + if args_cli.num_trajectories is not None and episodes >= args_cli.num_trajectories: + if pbar is not None: + pbar.close() + print(f"\nReached target number of trajectories ({args_cli.num_trajectories}). Stopping evaluation.") + break + + # run everything in inference mode + with torch.inference_mode(), torch.autocast(device_type=device.type) if args_cli.use_amp else nullcontext(): + # compute actions using wrapped diffusion policy + episode_steps = env.unwrapped.episode_length_buf + # first_step_mask = (episode_steps == 0).to(device) + # action_ids = ~first_step_mask.nonzero(as_tuple=False).reshape(-1).to(device) + # actions = torch.zeros((args_cli.num_envs, env.action_space.shape[0]), device=device) + # actions[:, -1] = -1.0 # default noop for last action dimension + + # if len(action_ids) > 0: + # actions[action_ids] = wrapped_policy.predict_action(obs_dict, action_ids.tolist()).to(device) + actions = wrapped_policy.predict_action(obs_dict).to(device) + + first_step_mask = (episode_steps == 0) + if torch.any(first_step_mask): + actions[first_step_mask, :-1] = 0.0 + actions[first_step_mask, -1] = -1.0 # close gripper + + if args_cli.save_video: + if len(cam_keys) == 0: + frame = env.render() + env_frames[0].append(frame) + else: + for i in range(args_cli.num_envs): + imgs = [] + + for cam in cam_keys: + img = obs_dict['policy'][cam][i].detach().cpu().permute(1, 2, 0).numpy() + img = (img * 255).clip(0, 255).astype('uint8') + imgs.append(img) + frame = np.concatenate(imgs, axis=1) + env_frames[i].append(frame) + + # apply actions using environment + step_result = env.step(actions) + if len(step_result) == 4: + obs_dict, rewards, dones, infos = step_result + else: + # Handle gymnasium v0.26+ format with 5 return values + obs_dict, rewards, terminated, truncated, infos = step_result + dones = terminated | truncated + + steps += 1 + + rewards_t = rewards if isinstance(rewards, torch.Tensor) else torch.as_tensor(rewards, device=device) + ep_returns += rewards_t.to(device) + + # Clear data for completed episodes + new_ids = [] + if isinstance(dones, torch.Tensor): + new_ids = (dones > 0).nonzero(as_tuple=False) + episodes += len(new_ids) + else: + # Handle scalar done value + if dones: + episodes += 1 + new_ids = [0] # Single episode done + + if isinstance(dones, torch.Tensor) and dones.any(): + reset_ids = (dones > 0).nonzero(as_tuple=False).reshape(-1) + num_new_successes = 0 + + term_dones = env.unwrapped.termination_manager._term_dones[reset_ids] # type: ignore + ep_success_flags = [] + for env_idx, term_row in enumerate(term_dones): + active_term_idx = term_row.nonzero(as_tuple=False) + is_success = False + if active_term_idx.numel() > 0: + # Handle multiple active termination conditions + active_term_indices = active_term_idx.flatten().cpu().tolist() + for term_idx in active_term_indices: + if term_names[term_idx] == "success": + num_new_successes += 1 + is_success = True + break # Count each environment only once + ep_success_flags.append(is_success) + successes.append(is_success) + + successful_episodes += num_new_successes + + for k, env_id in enumerate(reset_ids.detach().cpu().tolist()): + video_log = {} + ep_ret = float(ep_returns[env_id].detach().cpu().item()) + is_success = float(ep_success_flags[k]) + video_log["eval/episode_return"] = ep_ret + video_log["eval/success"] = is_success + # video_log["eval/success_rate_running"] = (successful_episodes / episodes) if episodes > 0 else 0.0 + if args_cli.save_video and len(env_frames[env_id]) > 0: + v = np.asarray(env_frames[env_id]) # T,H,W,C uint8 + v = v.transpose(0, 3, 1, 2) # T,C,H,W + if is_success > 0: + video_log[f"eval/success_video_{success_video_count}"] = wandb.Video(v, fps=10, format="mp4") + success_video_count += 1 + else: + video_log[f"eval/fail_video_{fail_video_count}"] = wandb.Video(v, fps=10, format="mp4") + fail_video_count += 1 + # video_log[f"eval/video_{video_count}"] = wandb.Video(v, fps=10, format="mp4") + # video_count += 1 + wandb.log(video_log) + + ep_returns[reset_ids] = 0.0 + + wrapped_policy.reset(reset_ids) + + # Store metrics for completed episodes + if "log" in infos: + # Store all metrics from this episode + for key, value in infos["log"].items(): + if key.startswith("Metrics/") or key.startswith("Episode_Reward/"): + if key not in episode_metrics: + episode_metrics[key] = [] + episode_metrics[key].append(value) + + steps = 0 + + if args_cli.save_video: + for i in reset_ids: + frames_to_save.extend(env_frames[i]) + env_frames[i] = [] + # imageio.mimsave("logs/policy_cameras.mp4", frames_to_save, fps=10, codec='libx264') + + # Update progress bar with success rate + if pbar is not None: + pbar.update(len(new_ids)) + success_rate = (successful_episodes / episodes * 100) if episodes > 0 else 0.0 + pbar.set_description(f"Evaluating trajectories (Success: {success_rate:.2f}%)") + + # Print final statistics + print("\nFinal Statistics:") + print(f"Total trajectories evaluated: {episodes}") + if successful_episodes > 0 or "Episode_Termination/success" in episode_metrics: + print(f"Successful trajectories: {successful_episodes}") + print(f"Success rate: {successful_episodes/episodes*100:.2f}%") + else: + print("Success rate: Not calculable (success metric not found in environment)") + + # Print metrics statistics + if episode_metrics: + print("\nAverage Metrics:") + for metric_name, values in sorted(episode_metrics.items()): + if values: # Only print if we have values + values = [float(v) if isinstance(v, torch.Tensor) else v for v in values] + mean = sum(values) / len(values) + print(f"{metric_name}: {mean:.4f}") + + final_stats = { + "eval/num_episodes": episodes, + "eval/successful_episodes": successful_episodes, + "eval/success_rate": (successful_episodes / episodes) if episodes > 0 else 0.0, + } + if episode_metrics: + for metric_name, values in sorted(episode_metrics.items()): + if values: + vals = [float(v.detach().cpu()) if isinstance(v, torch.Tensor) else float(v) for v in values] + final_stats[f"eval/{metric_name}_mean"] = float(np.mean(vals)) + wandb.log(final_stats) + wandb.finish() + + # Cleanup + if pbar is not None: + pbar.close() + env.close() + + +if __name__ == "__main__": + # run the main function - the decorator handles parameter passing + main() # type: ignore + # close sim app + simulation_app.close() diff --git a/scripts/ASTEROID/play_log_obs.py b/scripts/ASTEROID/play_log_obs.py new file mode 100644 index 00000000..1ba66e5e --- /dev/null +++ b/scripts/ASTEROID/play_log_obs.py @@ -0,0 +1,283 @@ +"""Play a trained RSL-RL state policy in the tactile data-collection env, log a +single observation term from the `data_collection` group every step, and save a +matplotlib plot. + +Mirrors play.py's preamble (AppLauncher boot, ckpt +load via OnPolicyRunner) but reads obs via the env's ObservationManager so we +can pull from any group, not just the wrapped policy obs. +""" + +from __future__ import annotations + +import argparse +import sys + +from isaaclab.app import AppLauncher + +parser = argparse.ArgumentParser(description="Play policy and log a data_collection obs term.") +parser.add_argument("--task", type=str, required=True, help="Gym task ID (tactile data-collection env).") +parser.add_argument("--checkpoint", type=str, required=True, help="Path to model_*.pt to load.") +parser.add_argument("--num_envs", type=int, default=1) +parser.add_argument("--num_steps", type=int, default=200, help="Env steps to roll out and log.") +parser.add_argument("--obs_group", type=str, default="data_collection") +parser.add_argument("--obs_term", type=str, nargs="+", default=["left_knuckle_pos"], + help="One or more obs term names to log; plotted on the same figure.") +parser.add_argument("--out_plot", type=str, default="./gripper_pos.png") +parser.add_argument("--video", action="store_true", default=False, help="Record video of the env.") +parser.add_argument("--video_length", type=int, default=None, + help="Number of env steps to record. Defaults to --num_steps.") +parser.add_argument("--video_dir", type=str, default="./play_videos", + help="Where gym.RecordVideo will write the rl-video-step-0.mp4.") +parser.add_argument("--out_synced_video", type=str, default="./synced_play.mp4", + help="Output path for the merged camera+plot video.") +parser.add_argument("--seed", type=int, default=None) +AppLauncher.add_app_launcher_args(parser) +args_cli, hydra_args = parser.parse_known_args() + +# Cameras must be enabled for video capture; mirror play.py. +if args_cli.video: + args_cli.enable_cameras = True + if args_cli.video_length is None: + args_cli.video_length = args_cli.num_steps + +# Hydra reads from sys.argv; isolate its overrides. +sys.argv = [sys.argv[0]] + hydra_args + +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +import gymnasium as gym # noqa: E402 +import os # noqa: E402 +import cv2 # noqa: E402 +import matplotlib # noqa: E402 +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 +import numpy as np # noqa: E402 +import torch # noqa: E402 + +from rsl_rl.runners import OnPolicyRunner # noqa: E402 + +from isaaclab.envs import ManagerBasedRLEnvCfg # noqa: E402 +from isaaclab_rl.rsl_rl import RslRlVecEnvWrapper # noqa: E402 + +import inspect # noqa: E402 + +import isaaclab_tasks # noqa: E402,F401 +import uwlab_tasks # noqa: E402,F401 +from uwlab_tasks.utils.hydra import hydra_task_config # noqa: E402 + + +def _drop_unknown_algorithm_keys(agent_cfg) -> None: + """Mirror cli_args.sanitize_rsl_rl_cfg: strip alg keys the installed PPO class can't accept.""" + alg_cfg = agent_cfg.algorithm + class_name = getattr(alg_cfg, "class_name", None) + if class_name is None: + return + from rsl_rl import algorithms + alg_class = getattr(algorithms, class_name, None) + if alg_class is None: + return + accepted = set(inspect.signature(alg_class.__init__).parameters.keys()) + for key in list(vars(alg_cfg)): + if key != "class_name" and key not in accepted: + delattr(alg_cfg, key) + + +@hydra_task_config(args_cli.task, "rsl_rl_cfg_entry_point") +def main(env_cfg: ManagerBasedRLEnvCfg, agent_cfg) -> None: + env_cfg.scene.num_envs = args_cli.num_envs + env_cfg.sim.device = args_cli.device or env_cfg.sim.device + env_cfg.seed = args_cli.seed if args_cli.seed is not None else agent_cfg.seed + + _drop_unknown_algorithm_keys(agent_cfg) + + env = gym.make( + args_cli.task, + cfg=env_cfg, + render_mode="rgb_array" if args_cli.video else None, + ) + + if args_cli.video: + os.makedirs(args_cli.video_dir, exist_ok=True) + env = gym.wrappers.RecordVideo( + env, + video_folder=args_cli.video_dir, + step_trigger=lambda step: step == 0, + video_length=args_cli.video_length, + disable_logger=True, + ) + print(f"[INFO] Recording video to {args_cli.video_dir}") + + base_env = env.unwrapped + wrapped = RslRlVecEnvWrapper(env, clip_actions=getattr(agent_cfg, "clip_actions", None)) + + runner = OnPolicyRunner(wrapped, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device) + print(f"[INFO] Loading actor-only weights from: {args_cli.checkpoint}") + # Tactile env's critic obs (188 dims at training) differs from current critic dim + # (155). strict=False does NOT tolerate size mismatches — it only ignores missing + # / unexpected keys. So drop everything not on the actor side before loading. + ckpt = torch.load(args_cli.checkpoint, weights_only=False, map_location=agent_cfg.device) + full_sd = ckpt["model_state_dict"] + actor_sd = {k: v for k, v in full_sd.items() if k.startswith(("actor.", "actor_obs_normalizer.", "std", "log_std"))} + # rsl_rl's ActorCritic.load_state_dict wraps nn.Module's and returns a bool, not + # _IncompatibleKeys, so call super() (nn.Module) directly to get key reports. + import torch.nn as _nn + incompat = _nn.Module.load_state_dict(runner.alg.policy, actor_sd, strict=False) + if any(k.startswith(("actor.", "actor_obs_normalizer.")) for k in incompat.missing_keys): + raise RuntimeError(f"Actor weights missing after partial load: {incompat.missing_keys}") + print(f"[INFO] Loaded {len(actor_sd)} actor tensors; " + f"skipped {len(full_sd) - len(actor_sd)} critic / RND / etc. tensors. " + f"({len(incompat.missing_keys)} missing, {len(incompat.unexpected_keys)} unexpected.)") + policy = runner.get_inference_policy(device=base_env.device) + + obs_buf: dict[str, list[np.ndarray]] = {term: [] for term in args_cli.obs_term} + + obs = wrapped.get_observations() + for step in range(args_cli.num_steps): + with torch.inference_mode(): + actions = policy(obs) + obs, _, _, _ = wrapped.step(actions) + # Pull the dict-form group directly from the obs manager + # (update_history=False so we don't double-tick history buffers). + group_obs = base_env.observation_manager.compute_group(args_cli.obs_group, update_history=False) + if not isinstance(group_obs, dict): + raise RuntimeError( + f"Group '{args_cli.obs_group}' is not dict-form (concatenate_terms=True?); " + f"set concatenate_terms=False to capture individual terms." + ) + for term in args_cli.obs_term: + if term not in group_obs: + raise KeyError( + f"Term '{term}' not in group '{args_cli.obs_group}'. " + f"Available: {list(group_obs.keys())}" + ) + obs_buf[term].append(group_obs[term].detach().cpu().numpy().copy()) + + dt = base_env.step_dt + fig, ax = plt.subplots(figsize=(10, 4)) + for term in args_cli.obs_term: + series = np.stack(obs_buf[term], axis=0).squeeze() + t = np.arange(series.shape[0]) * dt + if series.ndim == 1: + ax.plot(t, series, label=term) + else: + for i in range(series.shape[-1]): + ax.plot(t, series[..., i], label=f"{term}[{i}]") + print(f"[INFO] {term}: shape={series.shape}, min={series.min():.4f} max={series.max():.4f}") + ax.set_xlabel(f"time (s, env step_dt={dt:.3f})") + ax.set_ylabel("joint angle (rad)") + ax.set_title(f"{args_cli.obs_group} terms during play (ckpt={os.path.basename(args_cli.checkpoint)})") + ax.grid(True, alpha=0.3) + ax.legend() + fig.tight_layout() + out_path = os.path.abspath(args_cli.out_plot) + fig.savefig(out_path, dpi=120) + print(f"[INFO] Saved plot to: {out_path}") + + env.close() + + if args_cli.video: + # Locate the file gym.RecordVideo just wrote. + vid_files = sorted(f for f in os.listdir(args_cli.video_dir) if f.endswith(".mp4")) + if not vid_files: + print(f"[WARN] No mp4 found in {args_cli.video_dir}; skipping synced video.") + return + cam_path = os.path.join(args_cli.video_dir, vid_files[-1]) + out_synced = os.path.abspath(args_cli.out_synced_video) + _make_synced_video( + cam_path=cam_path, + obs_buf=obs_buf, + obs_terms=args_cli.obs_term, + step_dt=dt, + out_path=out_synced, + ckpt_name=os.path.basename(args_cli.checkpoint), + ) + + +def _make_synced_video(cam_path, obs_buf, obs_terms, step_dt, out_path, ckpt_name): + """Stack the recorded env video on top of a matplotlib plot of `obs_terms`, + drawing a moving vertical cursor on the plot in lock-step with the video. + Mirrors diffusion_policy/visualize_episode.py.""" + cap = cv2.VideoCapture(cam_path) + vid_fps = cap.get(cv2.CAP_PROP_FPS) or (1.0 / step_dt) + n_vid_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + ret, frame0 = cap.read() + cap.set(cv2.CAP_PROP_POS_FRAMES, 0) + if not ret: + print(f"[WARN] Could not read first frame from {cam_path}; skipping synced video.") + cap.release() + return + cam_h, cam_w = frame0.shape[:2] + + series = {term: np.stack(obs_buf[term], axis=0).squeeze() for term in obs_terms} + N = len(next(iter(series.values()))) + ratio = vid_fps / (1.0 / step_dt) # video frames per env step + + print(f"[INFO] Synced video: {n_vid_frames} video frames @ {vid_fps:.1f} fps, " + f"{N} env steps @ {1/step_dt:.1f} Hz (ratio {ratio:.2f}x)") + + plot_h = 280 + plot_dpi = 100 + fig, ax = plt.subplots(figsize=(cam_w / plot_dpi, plot_h / plot_dpi), dpi=plot_dpi) + fig.patch.set_facecolor("#1e1e1e") + ax.set_facecolor("#2a2a2a") + colors = ["#e67e22", "#3498db", "#2ecc71", "#9b59b6", "#e74c3c"] + t = np.arange(N) + for i, term in enumerate(obs_terms): + s = series[term] + if s.ndim == 1: + ax.plot(t, s, color=colors[i % len(colors)], lw=1.0, label=term) + else: + for j in range(s.shape[-1]): + ax.plot(t, s[:, j], color=colors[(i + j) % len(colors)], lw=1.0, label=f"{term}[{j}]") + ax.set_xlim(0, max(N - 1, 1)) + ax.set_ylabel("obs", color="#cccccc", fontsize=8) + ax.set_xlabel(f"env step ({step_dt*1000:.0f} ms)", color="#cccccc", fontsize=8) + ax.tick_params(colors="#aaaaaa", labelsize=7) + for spine in ax.spines.values(): + spine.set_edgecolor("#555555") + ax.grid(True, alpha=0.2, color="#888888") + ax.legend(loc="upper right", fontsize=7, facecolor="#333333", edgecolor="#555555", labelcolor="white") + ax.set_title(f"{ckpt_name}", color="#cccccc", fontsize=8) + fig.tight_layout(pad=0.6) + + fig.canvas.draw() + # matplotlib >=3.10 dropped tostring_rgb; buffer_rgba works on all recent versions. + plot_bg = np.asarray(fig.canvas.buffer_rgba())[..., :3].copy() # H x W x 3 (RGB) + plot_bg_bgr = cv2.cvtColor(plot_bg, cv2.COLOR_RGB2BGR) + + bbox = ax.get_window_extent() + img_h = plot_bg.shape[0] + ax_x0 = int(bbox.x0) + ax_x1 = int(bbox.x1) + ax_y0_img = img_h - int(bbox.y1) + ax_y1_img = img_h - int(bbox.y0) + plt.close(fig) + + out_h = cam_h + plot_h + out_w = cam_w + fourcc = cv2.VideoWriter_fourcc(*"mp4v") + writer = cv2.VideoWriter(out_path, fourcc, vid_fps, (out_w, out_h)) + print(f"[INFO] Writing {out_path} ({out_w}x{out_h} @ {vid_fps:.1f} fps)") + + for frame_idx in range(n_vid_frames): + env_idx = min(int(round(frame_idx / max(ratio, 1e-6))), N - 1) + ret, frame = cap.read() + if not ret: + frame = np.zeros((cam_h, cam_w, 3), dtype=np.uint8) + plot_frame = plot_bg_bgr.copy() + x_norm = env_idx / max(N - 1, 1) + x_px = ax_x0 + int(x_norm * (ax_x1 - ax_x0)) + cv2.line(plot_frame, (x_px, ax_y0_img), (x_px, ax_y1_img), (255, 255, 255), 1) + if plot_frame.shape[1] != out_w: + plot_frame = cv2.resize(plot_frame, (out_w, plot_h)) + writer.write(np.vstack([frame, plot_frame])) + + writer.release() + cap.release() + print(f"[INFO] Saved synced video to: {out_path}") + + +if __name__ == "__main__": + main() + simulation_app.close() diff --git a/scripts/ASTEROID/record_cube_object_anywhere_resets.sh b/scripts/ASTEROID/record_cube_object_anywhere_resets.sh new file mode 100755 index 00000000..51db8036 --- /dev/null +++ b/scripts/ASTEROID/record_cube_object_anywhere_resets.sh @@ -0,0 +1,10 @@ +#!/bin/sh + +python scripts_v2/tools/record_reset_states.py \ + --task Asteroid-UR5eRobotiq2f85-ObjectAnywhereEEAnywhere-v0 \ + --dataset_dir ./Datasets/CubePick \ + --reset_type ObjectAnywhereEEAnywhere \ + --num_envs 64 \ + --num_reset_states 1000 \ + --headless \ + env.scene.insertive_object=cube diff --git a/scripts/ASTEROID/run_asteroid.py b/scripts/ASTEROID/run_asteroid.py new file mode 100644 index 00000000..dac5da35 --- /dev/null +++ b/scripts/ASTEROID/run_asteroid.py @@ -0,0 +1,509 @@ +"""ASTEROID orchestrator: iterative in-context exploration + distillation. + +Each iteration runs three stages, all as subprocesses: + + 1. collect -- roll out the expert (plus, from iteration 1 on, the previous + student as an explorer) in the data-collection env + (``scripts/ASTEROID/collect_demos_asteroid.py``) + 2. train -- fit a diffusion-policy student on every dataset collected so far + (``diffusion_policy/train.py``) + 3. eval -- roll out the student in the eval env + (``scripts/ASTEROID/eval_asteroid_policy.py``) + +Hyperparameters form a hierarchy of dataclasses: a :class:`RunCfg` holds the +run-level settings plus an ordered list of :class:`IterationCfg`, and each +iteration owns the :class:`CollectCfg` / :class:`TrainCfg` / :class:`EvalCfg` +for that iteration. Curricula are functions that build the iteration list; add +a new one to :data:`CURRICULA` and select it with ``--schedule``. + +Run from the repository root, e.g.:: + + python scripts/ASTEROID/run_asteroid.py \ + --data_task Asteroid-Ur5eRobotiq2f85-RelCartesianOSC-Tactile-DataCollection-v0 \ + --eval_task Asteroid-Ur5eRobotiq2f85-RelCartesianOSC-Tactile-Play-v0 \ + --expert_policy_checkpoint logs/exported/policy.pt \ + --max_iterations 4 --exp_name my_run +""" + +from __future__ import annotations + +import argparse +import dataclasses +import datetime +import glob +import json +import os +import re +import subprocess +import sys +from dataclasses import dataclass, field + +# --------------------------------------------------------------------------- +# Hyperparameter hierarchy +# --------------------------------------------------------------------------- + + +@dataclass +class CollectCfg: + """Data collection for one iteration (``collect_demos_asteroid.py``). + + Exploration horizons are fractions of the episode during which the explorer + (previous iteration's student) acts before the expert takes over. Iteration 0 + has no student yet, so its horizons must be ``0.0`` (expert only). + """ + + num_demos: int = 10 + num_envs: int = 2 + episode_length_s: float = 8.0 + min_exploration_horizon: float = 0.0 + max_exploration_horizon: float = 0.0 + expert_noise: float = 0.0 + + +@dataclass +class TrainCfg: + """Student training for one iteration (``diffusion_policy/train.py``).""" + + lr: float = 1e-4 + #: Sampling weight for each dataset collected so far (index 0 = iteration 0's + #: dataset). Must have ``iteration + 1`` entries summing to 1. + sampling_ratios: tuple[float, ...] = (1.0,) + #: Warm-start from the previous iteration's student checkpoint. + init_from_previous: bool = True + #: Training step whose checkpoint is handed to eval / the next iteration. + checkpoint_step: int = 40_000 + + +@dataclass +class EvalCfg: + """Student evaluation for one iteration (``eval_distilled_policy.py``).""" + + num_trajectories: int = 10 + num_envs: int = 2 + episode_length_s: float = 10.0 + + +@dataclass +class IterationCfg: + """Everything that varies per iteration.""" + + collect: CollectCfg = field(default_factory=CollectCfg) + train: TrainCfg = field(default_factory=TrainCfg) + eval: EvalCfg = field(default_factory=EvalCfg) + + +@dataclass +class RunCfg: + """Run-level settings shared by every iteration, plus the iteration schedule.""" + + exp_name: str + output_dir: str + wandb_project: str + data_task: str + eval_task: str + expert_policy_checkpoint: str + config_dir: str + config_name: str + insertive_object: str = "cube" + receptive_object: str | None = None + seed: int = 0 + video: bool = True + iterations: list[IterationCfg] = field(default_factory=list) + + @property + def num_iterations(self) -> int: + return len(self.iterations) + + def validate(self) -> None: + assert self.num_iterations > 0, "RunCfg needs at least one iteration" + for i, it in enumerate(self.iterations): + c, t = it.collect, it.train + assert 0.0 <= c.min_exploration_horizon <= c.max_exploration_horizon <= 1.0, ( + f"iteration {i}: exploration horizons must satisfy 0 <= min <= max <= 1, " + f"got ({c.min_exploration_horizon}, {c.max_exploration_horizon})" + ) + if i == 0: + assert c.max_exploration_horizon == 0.0, ( + "iteration 0 has no explorer yet; its exploration horizons must be 0.0" + ) + assert len(t.sampling_ratios) == i + 1, ( + f"iteration {i}: sampling_ratios must have {i + 1} entries, got {len(t.sampling_ratios)}" + ) + assert abs(sum(t.sampling_ratios) - 1.0) < 1e-6, ( + f"iteration {i}: sampling_ratios must sum to 1.0, got {sum(t.sampling_ratios)}" + ) + + def to_dict(self) -> dict: + return dataclasses.asdict(self) + + def dump(self, path: str) -> None: + with open(path, "w") as f: + json.dump(self.to_dict(), f, indent=2) + + +# --------------------------------------------------------------------------- +# Curricula: functions that build the per-iteration schedule +# --------------------------------------------------------------------------- + + +def default_curriculum( + num_iterations: int, + *, + num_demos: int, + num_data_envs: int, + num_eval_envs: int, + num_eval_episodes: int, + expert_noise: float, + init_from_previous: bool, +) -> list[IterationCfg]: + """Curriculum used for the cube-pick runs. + + Index ``i`` of each table is iteration ``i``. ``HORIZONS[i]`` is the explorer + horizon used when collecting iteration ``i``'s dataset (so ``HORIZONS[0]`` is + expert-only); ``SAMPLING_RATIOS[i]`` weights datasets ``0..i`` when training + iteration ``i``'s student. + """ + COLLECT_EPISODE_LENGTH_S = 8.0 + EVAL_EPISODE_LENGTH_S = 10.0 + HORIZONS = [ + (0.00, 0.00), # expert only + (0.20, 0.50), # 1.6s - 4.0s of an 8s episode + (0.30, 0.70), + (0.40, 0.90), + (0.50, 0.95), + (0.60, 0.95), + ] + LRS = [1e-4, 1e-5, 1e-5, 1e-5, 1e-5, 1e-5] + SAMPLING_RATIOS = [ + (1.0,), + (0.25, 0.75), + (0.2, 0.3, 0.5), + (0.1, 0.2, 0.3, 0.4), + (0.05, 0.1, 0.2, 0.25, 0.4), + (0.05, 0.1, 0.15, 0.15, 0.2, 0.35), + ] + max_supported = min(len(HORIZONS), len(LRS), len(SAMPLING_RATIOS)) + assert 1 <= num_iterations <= max_supported, ( + f"default curriculum supports 1..{max_supported} iterations, got {num_iterations}" + ) + + iterations = [] + for i in range(num_iterations): + iterations.append( + IterationCfg( + collect=CollectCfg( + num_demos=num_demos, + num_envs=num_data_envs, + episode_length_s=COLLECT_EPISODE_LENGTH_S, + min_exploration_horizon=HORIZONS[i][0], + max_exploration_horizon=HORIZONS[i][1], + expert_noise=expert_noise, + ), + train=TrainCfg( + lr=LRS[i], + sampling_ratios=SAMPLING_RATIOS[i], + init_from_previous=init_from_previous, + ), + eval=EvalCfg( + num_trajectories=num_eval_episodes, + num_envs=num_eval_envs, + episode_length_s=EVAL_EPISODE_LENGTH_S, + ), + ) + ) + return iterations + + +CURRICULA = { + "default": default_curriculum, +} + + +# --------------------------------------------------------------------------- +# Stage runners +# --------------------------------------------------------------------------- + +_STEP_CKPT_RE = re.compile(r"step_(\d+)\.ckpt$") + + +def expected_train_checkpoint(train_output_dir: str, step: int) -> str: + """Resolve the checkpoint written by a training iteration. + + Selection order: + 1. ``step_{step:07d}.ckpt`` for the requested step. + 2. The highest-numbered ``step_*.ckpt`` in the checkpoints dir. + 3. ``latest.ckpt`` -- the final-state snapshot written by the workspace. + """ + ckpt_dir = os.path.join(train_output_dir, "checkpoints") + + preferred = os.path.join(ckpt_dir, f"step_{step:07d}.ckpt") + if os.path.exists(preferred): + return preferred + + candidates: list[tuple[int, str]] = [] + for path in glob.glob(os.path.join(ckpt_dir, "step_*.ckpt")): + m = _STEP_CKPT_RE.search(os.path.basename(path)) + if m is not None: + candidates.append((int(m.group(1)), path)) + if candidates: + best_step, best_path = max(candidates) + print( + f"[asteroid] step_{step:07d}.ckpt missing under {ckpt_dir}; " + f"falling back to {os.path.basename(best_path)} (step {best_step})." + ) + return best_path + + latest = os.path.join(ckpt_dir, "latest.ckpt") + if os.path.exists(latest): + print(f"[asteroid] no step_*.ckpt under {ckpt_dir}; falling back to latest.ckpt.") + return latest + + return preferred + + +class AsteroidRun: + """Executes a :class:`RunCfg` under ``base_output_dir``. + + Layout:: + + base_output_dir/ + run_cfg.json + dataset-iteration-{i}/data.zarr collected by iteration i + iteration_{i}/checkpoints/*.ckpt student trained in iteration i + """ + + def __init__(self, cfg: RunCfg, base_output_dir: str, dry_run: bool = False): + cfg.validate() + self.cfg = cfg + self.base_output_dir = base_output_dir + self.dry_run = dry_run + + # -- paths ------------------------------------------------------------- + + def dataset_dir(self, iteration: int) -> str: + return os.path.join(self.base_output_dir, f"dataset-iteration-{iteration}") + + def train_output_dir(self, iteration: int) -> str: + return os.path.join(self.base_output_dir, f"iteration_{iteration}") + + def student_checkpoint(self, iteration: int) -> str: + step = self.cfg.iterations[iteration].train.checkpoint_step + return expected_train_checkpoint(self.train_output_dir(iteration), step) + + # -- stages ------------------------------------------------------------ + + def _run(self, stage: str, command: list[str]) -> None: + print(f"[asteroid] {stage}: {' '.join(command)}", flush=True) + if self.dry_run: + return + result = subprocess.run(command) + if result.returncode != 0: + print(f"[asteroid] {stage} failed with return code {result.returncode}") + sys.exit(1) + print(f"[asteroid] {stage} finished") + + def _scene_overrides(self) -> list[str]: + overrides = [f"env.scene.insertive_object={self.cfg.insertive_object}"] + if self.cfg.receptive_object is not None: + overrides.append(f"env.scene.receptive_object={self.cfg.receptive_object}") + return overrides + + def collect(self, iteration: int, explorer_checkpoint: str | None) -> str: + """Collect iteration ``iteration``'s dataset; returns the dataset dir.""" + cfg, c = self.cfg, self.cfg.iterations[iteration].collect + dataset_dir = self.dataset_dir(iteration) + command = [ + "python", "scripts/ASTEROID/collect_demos_asteroid.py", + "--task", cfg.data_task, + "--dataset_file", os.path.join(dataset_dir, "data.zarr"), + "--num_envs", str(c.num_envs), + "--num_demos", str(c.num_demos), + "--episode_length_s", str(c.episode_length_s), + "--min_exploration_horizon", str(c.min_exploration_horizon), + "--max_exploration_horizon", str(c.max_exploration_horizon), + "--expert_noise", str(c.expert_noise), + "--seed", str(cfg.seed), + "--headless", + f'agent.algorithm.offline_algorithm_cfg.behavior_cloning_cfg.experts_path=["{cfg.expert_policy_checkpoint}"]', + *self._scene_overrides(), + ] + if explorer_checkpoint is not None: + command += ["--exploration_checkpoint", explorer_checkpoint] + if cfg.video: + command += ["--video", "--video_dir", os.path.join(dataset_dir, "videos")] + self._run(f"collect[{iteration}]", command) + return dataset_dir + + def train(self, iteration: int, dataset_dirs: list[str], pretrained_checkpoint: str | None) -> str: + """Train iteration ``iteration``'s student; returns its checkpoint path.""" + cfg, t = self.cfg, self.cfg.iterations[iteration].train + assert len(dataset_dirs) == len(t.sampling_ratios), ( + f"iteration {iteration}: have {len(dataset_dirs)} datasets but {len(t.sampling_ratios)} sampling ratios" + ) + dataset_config = ",".join( + f"{{dataset_dir: {d}, sampling_ratio: {r}}}" for d, r in zip(dataset_dirs, t.sampling_ratios) + ) + output_dir = self.train_output_dir(iteration) + if not self.dry_run: + os.makedirs(output_dir, exist_ok=True) + command = [ + "python", "diffusion_policy/train.py", + "--config-name", cfg.config_name, + "--config-dir", cfg.config_dir, + f"output_dir={output_dir}", + f"task.dataset.dataset_config=[{dataset_config}]", + f"name={cfg.exp_name}", + f"exp_name={cfg.exp_name}", + f"logging.project={cfg.wandb_project}", + "logging.group=train", + f"optimizer.lr={t.lr}", + f"seed={cfg.seed}", + f"iteration={iteration}", + ] + if pretrained_checkpoint is not None: + command.append(f"checkpoint.pretrained_ckpt_path={pretrained_checkpoint}") + self._run(f"train[{iteration}]", command) + return self.student_checkpoint(iteration) + + def eval(self, iteration: int, checkpoint: str) -> None: + cfg, e = self.cfg, self.cfg.iterations[iteration].eval + command = [ + "python", "scripts/ASTEROID/eval_asteroid_policy.py", + "--task", cfg.eval_task, + "--checkpoint", checkpoint, + "--num_trajectories", str(e.num_trajectories), + "--num_envs", str(e.num_envs), + "--episode_length_s", str(e.episode_length_s), + "--seed", str(cfg.seed), + "--exp_name", cfg.exp_name, + "--wandb_project", cfg.wandb_project, + "--wandb_group", "eval", + "--iteration", str(iteration), + "--headless", + *self._scene_overrides(), + ] + if cfg.video: + command += ["--save_video", "--enable_cameras"] + self._run(f"eval[{iteration}]", command) + + # -- driver ------------------------------------------------------------ + + def run(self, start_iteration: int = 0, initial_dataset_dir: str | None = None) -> None: + """Run iterations ``start_iteration .. num_iterations-1``. + + When ``start_iteration > 0`` the datasets and student checkpoint of the + earlier iterations are expected to exist under ``base_output_dir``. + ``initial_dataset_dir`` replaces iteration 0's collection with an + existing dataset. + """ + cfg = self.cfg + assert 0 <= start_iteration < cfg.num_iterations + if not self.dry_run: + os.makedirs(self.base_output_dir, exist_ok=True) + cfg.dump(os.path.join(self.base_output_dir, "run_cfg.json")) + + dataset_dirs = [self.dataset_dir(i) for i in range(start_iteration)] + student_checkpoint = self.student_checkpoint(start_iteration - 1) if start_iteration > 0 else None + + for i in range(start_iteration, cfg.num_iterations): + print(f"[asteroid] ===== iteration {i} / {cfg.num_iterations - 1} =====") + if i == 0 and initial_dataset_dir is not None: + dataset_dirs.append(initial_dataset_dir) + else: + dataset_dirs.append(self.collect(i, explorer_checkpoint=student_checkpoint)) + + pretrained = student_checkpoint if cfg.iterations[i].train.init_from_previous else None + student_checkpoint = self.train(i, dataset_dirs, pretrained_checkpoint=pretrained) + self.eval(i, student_checkpoint) + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + # tasks / policies + p.add_argument("--data_task", default="Asteroid-Ur5eRobotiq2f85-RelCartesianOSC-Tactile-DataCollection-v0") + p.add_argument("--eval_task", default="Asteroid-Ur5eRobotiq2f85-RelCartesianOSC-Tactile-Play-v0") + p.add_argument("--expert_policy_checkpoint", default="logs/policy_cube_final_v4.pt") + p.add_argument("--insertive_object", default="cube") + p.add_argument("--receptive_object", default=None) + # student training config + p.add_argument("--config_dir", default="diffusion_policy/diffusion_policy/config") + p.add_argument("--config_name", default="incontext_exploration_debug.yaml") + # logging + p.add_argument("--output_dir", default="logs/incontext_exploration_debug") + p.add_argument("--exp_name", default="incontext_exploration_debug") + p.add_argument("--wandb_project", default="incontext_exploration") + p.add_argument("--no_video", action="store_true", help="Disable video recording in collect and eval.") + p.add_argument("--seed", type=int, default=0) + # curriculum + p.add_argument("--schedule", choices=sorted(CURRICULA), default="default", help="Which curriculum to run.") + p.add_argument("--max_iterations", type=int, default=3, help="Number of iterations to run.") + p.add_argument("--num_demos", type=int, default=10, help="Demos collected per iteration.") + p.add_argument("--num_data_envs", type=int, default=2) + p.add_argument("--num_eval_envs", type=int, default=2) + p.add_argument("--num_eval_episodes", type=int, default=10) + p.add_argument("--expert_noise", type=float, default=0.0) + p.add_argument("--not_use_pretrained_checkpoint", action="store_true", + help="Train every iteration's student from scratch instead of warm-starting.") + # resume + p.add_argument("--initial_dataset_path", default=None, help="Use this dataset for iteration 0 instead of collecting.") + p.add_argument("--start_iteration", type=int, default=None, help="Resume from this iteration (requires --checkpoint_dir).") + p.add_argument("--checkpoint_dir", default=None, help="Existing run directory to resume from.") + p.add_argument("--dry_run", action="store_true", help="Print the stage commands without running them.") + return p.parse_args(argv) + + +def build_run_cfg(args: argparse.Namespace) -> RunCfg: + iterations = CURRICULA[args.schedule]( + args.max_iterations, + num_demos=args.num_demos, + num_data_envs=args.num_data_envs, + num_eval_envs=args.num_eval_envs, + num_eval_episodes=args.num_eval_episodes, + expert_noise=args.expert_noise, + init_from_previous=not args.not_use_pretrained_checkpoint, + ) + return RunCfg( + exp_name=args.exp_name, + output_dir=args.output_dir, + wandb_project=args.wandb_project, + data_task=args.data_task, + eval_task=args.eval_task, + expert_policy_checkpoint=args.expert_policy_checkpoint, + config_dir=args.config_dir, + config_name=args.config_name, + insertive_object=args.insertive_object, + receptive_object=args.receptive_object, + seed=args.seed, + video=not args.no_video, + iterations=iterations, + ) + + +def main(argv: list[str] | None = None) -> None: + args = parse_args(argv) + cfg = build_run_cfg(args) + print("[asteroid] run config:\n" + json.dumps(cfg.to_dict(), indent=2)) + + if args.start_iteration is not None: + assert args.checkpoint_dir is not None, "--start_iteration requires --checkpoint_dir" + assert args.start_iteration > 0, "--start_iteration must be > 0 (use a fresh run for iteration 0)" + assert args.initial_dataset_path is None, "--initial_dataset_path only applies to a fresh run" + base_output_dir = args.checkpoint_dir + start_iteration = args.start_iteration + else: + stamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + base_output_dir = os.path.join(cfg.output_dir, cfg.exp_name, stamp) + start_iteration = 0 + + AsteroidRun(cfg, base_output_dir, dry_run=args.dry_run).run( + start_iteration=start_iteration, initial_dataset_dir=args.initial_dataset_path + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/reinforcement_learning/rsl_rl/play.py b/scripts/reinforcement_learning/rsl_rl/play.py index 611e5047..e2312b9d 100644 --- a/scripts/reinforcement_learning/rsl_rl/play.py +++ b/scripts/reinforcement_learning/rsl_rl/play.py @@ -34,6 +34,13 @@ help="Use the pre-trained checkpoint from Nucleus.", ) parser.add_argument("--real-time", action="store_true", default=False, help="Run in real-time, if possible.") +parser.add_argument("--max_episodes", type=int, default=None, + help="Stop after this many completed episodes and write a JSON.") +parser.add_argument("--eval_output", type=str, default=None, help="Path to write per-eval JSON.") +parser.add_argument("--discretize_actions", action="store_true", default=False, + help="Snap arm dims to bin grid + binarize gripper sign before env.step.") +parser.add_argument("--num_bins", type=int, default=51, help="Bins per arm dim for --discretize_actions.") +parser.add_argument("--action_bound", type=float, default=25.0, help="Symmetric clip range for arm dims.") # append RSL-RL cli arguments cli_args.add_rsl_rl_args(parser) # append AppLauncher cli args @@ -179,6 +186,32 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen dt = env.unwrapped.step_dt + # Discretizer setup (UWLab-ICL collect_demos_asteroid scheme): arm dims 0-5 + # snap to evenly-spaced bin centers in [-bound, +bound]; gripper dim 6 sign- + # thresholded to {-1, +1}. Mirrors the inference-side discretization used by + # the categorical-head distillation student. + bin_centers = None + if args_cli.discretize_actions: + bin_centers = torch.linspace( + -args_cli.action_bound, args_cli.action_bound, args_cli.num_bins, + device=env.unwrapped.device, dtype=torch.float32, + ) + print(f"[eval] discretize ON: arm bins={args_cli.num_bins} ±{args_cli.action_bound}, gripper sign-thresh") + + # Termination bookkeeping (success vs timeout vs other) + eval_episodes = 0 + eval_term_success = 0 + eval_term_timeout = 0 + eval_term_other = 0 + has_success_term = False + has_timeout_term = False + if args_cli.max_episodes is not None: + term_mgr = env.unwrapped.termination_manager + active = list(term_mgr.active_terms) + has_success_term = "success" in active + has_timeout_term = "time_out" in active + print(f"[eval] termination terms: {active}") + # reset environment obs = env.get_observations() timestep = 0 @@ -189,10 +222,45 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen with torch.inference_mode(): # agent stepping actions = policy(obs) + if bin_centers is not None: + a = actions.clamp(-args_cli.action_bound, args_cli.action_bound) + # arm dims (0..5): snap to nearest bin center + for d in range(min(6, actions.shape[-1])): + diff = (a[:, d:d+1] - bin_centers.unsqueeze(0)).abs() + actions[:, d] = bin_centers[diff.argmin(dim=-1).squeeze(-1)] + # gripper dim (6): hard sign-threshold + if actions.shape[-1] >= 7: + actions[:, 6] = torch.where( + actions[:, 6] >= 0, + torch.ones_like(actions[:, 6]), + -torch.ones_like(actions[:, 6]), + ) # env stepping obs, _, dones, _ = env.step(actions) # reset recurrent states for episodes that have terminated policy_nn.reset(dones) + + if args_cli.max_episodes is not None and isinstance(dones, torch.Tensor) and dones.any(): + done_ids = (dones > 0).nonzero(as_tuple=False).reshape(-1) + term_mgr = env.unwrapped.termination_manager + if has_success_term: + eval_term_success += int(term_mgr.get_term("success")[done_ids].sum().item()) + if has_timeout_term: + eval_term_timeout += int(term_mgr.get_term("time_out")[done_ids].sum().item()) + acc = ( + (int(term_mgr.get_term("success")[done_ids].sum().item()) if has_success_term else 0) + + (int(term_mgr.get_term("time_out")[done_ids].sum().item()) if has_timeout_term else 0) + ) + eval_term_other += int(done_ids.numel()) - acc + eval_episodes += int(done_ids.numel()) + rate_s = eval_term_success / max(eval_episodes, 1) + rate_t = eval_term_timeout / max(eval_episodes, 1) + print(f"[eval] episodes={eval_episodes}/{args_cli.max_episodes} " + f"term_success={eval_term_success} ({rate_s:.4f}) " + f"term_timeout={eval_term_timeout} ({rate_t:.4f}) " + f"term_other={eval_term_other}", flush=True) + if eval_episodes >= args_cli.max_episodes: + break if args_cli.video: timestep += 1 # Exit the play loop after recording one video @@ -204,6 +272,29 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen if args_cli.real_time and sleep_time > 0: time.sleep(sleep_time) + if args_cli.max_episodes is not None and args_cli.eval_output is not None: + import json + rate_s = eval_term_success / max(eval_episodes, 1) + rate_t = eval_term_timeout / max(eval_episodes, 1) + result = { + "task": args_cli.task, + "checkpoint": resume_path, + "discretize_actions": bool(args_cli.discretize_actions), + "num_bins": args_cli.num_bins, + "action_bound": args_cli.action_bound, + "num_envs": env_cfg.scene.num_envs, + "num_episodes": eval_episodes, + "term_success": eval_term_success, + "term_success_rate": rate_s, + "term_timeout": eval_term_timeout, + "term_timeout_rate": rate_t, + "term_other": eval_term_other, + } + os.makedirs(os.path.dirname(args_cli.eval_output) or ".", exist_ok=True) + with open(args_cli.eval_output, "w") as f: + json.dump(result, f, indent=2) + print(f"[eval] wrote {args_cli.eval_output}: term_success={rate_s:.4f}") + # close the simulator env.close() diff --git a/scripts_v2/tools/record_reset_states.py b/scripts_v2/tools/record_reset_states.py index c656bba1..0ef961ca 100644 --- a/scripts_v2/tools/record_reset_states.py +++ b/scripts_v2/tools/record_reset_states.py @@ -35,10 +35,17 @@ parser.add_argument( "--num_reset_states", type=int, default=100, help="Number of reset states to record. Set to 0 for infinite." ) +parser.add_argument("--video", action="store_true", default=False, help="Record video of the env.") +parser.add_argument("--video_length", type=int, default=200, help="Length of recorded video (env steps).") +parser.add_argument("--video_dir", type=str, default=None, help="Output dir for the video (defaults under --dataset_dir).") AppLauncher.add_app_launcher_args(parser) args_cli, remaining_args = parser.parse_known_args() +# Cameras must be enabled for video capture; mirror play.py. +if args_cli.video: + args_cli.enable_cameras = True + # Launch omniverse app app_launcher = AppLauncher(args_cli) simulation_app = app_launcher.app @@ -80,8 +87,14 @@ def main(env_cfg, agent_cfg) -> None: # Derive pair directory and reset type for output path insertive_usd_path = env_cfg.scene.insertive_object.spawn.usd_path - receptive_usd_path = env_cfg.scene.receptive_object.spawn.usd_path - pair = task_mdp.utils.compute_pair_dir(insertive_usd_path, receptive_usd_path) + receptive_object_cfg = getattr(env_cfg.scene, "receptive_object", None) + if receptive_object_cfg is not None: + receptive_usd_path = receptive_object_cfg.spawn.usd_path + pair = task_mdp.utils.compute_pair_dir(insertive_usd_path, receptive_usd_path) + else: + # single-object tasks (e.g. ASTEROID pick) are keyed by the insertive object alone + receptive_usd_path = None + pair = task_mdp.utils.compute_pair_dir(insertive_usd_path) # Auto-infer reset_type from task name if not provided reset_type = args_cli.reset_type @@ -100,7 +113,8 @@ def main(env_cfg, agent_cfg) -> None: print(f"Recording reset states for: {pair} / {reset_type}") print(f"Insertive: {insertive_usd_path}") - print(f"Receptive: {receptive_usd_path}") + if receptive_usd_path is not None: + print(f"Receptive: {receptive_usd_path}") # Setup recording configuration output_dir = os.path.join(args_cli.dataset_dir, "Resets", pair) @@ -114,18 +128,38 @@ def main(env_cfg, agent_cfg) -> None: env_cfg.recorders.dataset_file_handler_class_type = TorchDatasetFileHandler # create environment - env = cast(ManagerBasedRLEnv, gym.make(args_cli.task, cfg=env_cfg)).unwrapped + env = gym.make( + args_cli.task, + cfg=env_cfg, + render_mode="rgb_array" if args_cli.video else None, + ) + + if args_cli.video: + video_dir = args_cli.video_dir or os.path.join(args_cli.dataset_dir, "videos", reset_type) + os.makedirs(video_dir, exist_ok=True) + video_kwargs = { + "video_folder": video_dir, + "step_trigger": lambda step: step == 0, + "video_length": args_cli.video_length, + "disable_logger": True, + } + print(f"[INFO] Recording video to {video_dir}") + env = gym.wrappers.RecordVideo(env, **video_kwargs) + + # Step through the wrapped env so RecordVideo intercepts each step; keep a + # handle on the unwrapped ManagerBasedRLEnv for action_space / recorder_manager. + base_env = cast(ManagerBasedRLEnv, env.unwrapped) env.reset() # Run reset state sampling num_reset_conditions_evaluated = 0 current_successful_reset_conditions = 0 - actions = torch.zeros(env.action_space.shape, device=env.device, dtype=torch.float32) + actions = torch.zeros(base_env.action_space.shape, device=base_env.device, dtype=torch.float32) if "ObjectAnywhereEEGrasped" in args_cli.task or "ObjectRestingEEGrasped" in args_cli.task: actions[:, -1] = -1.0 else: actions[:, -1] = ( - torch.randint(0, 2, (env.num_envs,), device=env.device, dtype=torch.float32) * 2 - 1 + torch.randint(0, 2, (base_env.num_envs,), device=base_env.device, dtype=torch.float32) * 2 - 1 ) # Randomly choose between -1 and 1 # Create progress bar @@ -133,7 +167,14 @@ def main(env_cfg, agent_cfg) -> None: start_time = time.time() - while current_successful_reset_conditions < args_cli.num_reset_states: + # When recording a video, ensure we run long enough to capture video_length steps + # even if the requested --num_reset_states would otherwise be hit sooner. + video_steps_done = 0 + while current_successful_reset_conditions < args_cli.num_reset_states or ( + args_cli.video and video_steps_done < args_cli.video_length + ): + if args_cli.video: + video_steps_done += 1 # Step environment (this will evaluate grasps in parallel across environments) _, _, terminated, truncated, _ = env.step(actions) dones = terminated | truncated @@ -144,11 +185,11 @@ def main(env_cfg, agent_cfg) -> None: "ObjectAnywhereEEGrasped" in args_cli.task or "ObjectRestingEEGrasped" in args_cli.task ): actions[done_idx, -1] = ( - torch.randint(0, 2, (done_idx.numel(),), device=env.device, dtype=torch.float32) * 2 - 1 + torch.randint(0, 2, (done_idx.numel(),), device=base_env.device, dtype=torch.float32) * 2 - 1 ) # Update progress based on successful reset conditions - new_successful_count = env.recorder_manager.exported_successful_episode_count + new_successful_count = base_env.recorder_manager.exported_successful_episode_count if new_successful_count > current_successful_reset_conditions: increment = new_successful_count - current_successful_reset_conditions current_successful_reset_conditions = new_successful_count @@ -157,13 +198,13 @@ def main(env_cfg, agent_cfg) -> None: # Count total reset conditions evaluated (sum across all environments) num_reset_conditions_evaluated += dones.sum().item() - if env.sim.is_stopped(): + if base_env.sim.is_stopped(): break pbar.close() # Get final statistics - final_successful_reset_conditions = env.recorder_manager.exported_successful_episode_count + final_successful_reset_conditions = base_env.recorder_manager.exported_successful_episode_count print("Reset state recording complete!") print(f"Total reset conditions evaluated: {num_reset_conditions_evaluated}") print(f"Successful reset conditions: {final_successful_reset_conditions}") diff --git a/source/uwlab/uwlab/utils/datasets/torch_dataset_file_handler.py b/source/uwlab/uwlab/utils/datasets/torch_dataset_file_handler.py index f52442d1..e84b2a85 100644 --- a/source/uwlab/uwlab/utils/datasets/torch_dataset_file_handler.py +++ b/source/uwlab/uwlab/utils/datasets/torch_dataset_file_handler.py @@ -97,13 +97,13 @@ def load_episode(self, episode_name: str, device: str = "cpu") -> EpisodeData | raise NotImplementedError("Load episode not supported for preprocessed format") def flush(self): - """Flush any pending data to disk.""" - if self._file_path and self._episode_data: - torch.save(self._episode_data, self._file_path) + # No-op: defer writes until close() to avoid O(N) re-serialization on every episode. + pass def close(self): - """Close the dataset file handler.""" - self.flush() + """Close the dataset file handler — actually writes the file here.""" + if self._file_path and self._episode_data: + torch.save(self._episode_data, self._file_path) self._episode_data = {} self._file_path = None diff --git a/source/uwlab/uwlab/utils/datasets/zarr_dataset_file_handler.py b/source/uwlab/uwlab/utils/datasets/zarr_dataset_file_handler.py index 219f4fd5..53e25ad0 100644 --- a/source/uwlab/uwlab/utils/datasets/zarr_dataset_file_handler.py +++ b/source/uwlab/uwlab/utils/datasets/zarr_dataset_file_handler.py @@ -234,7 +234,11 @@ def _convert_and_save_episode(self, episode: EpisodeData): "obs": processed_obs, "rewards": episode_dict.get("rewards", torch.zeros(num_frames)).cpu().numpy(), "dones": episode_dict.get("dones", torch.cat([torch.zeros(num_frames - 1), torch.ones(1)])).cpu().numpy(), + "expert_mask": episode_dict.get("expert_mask", torch.ones((num_frames, 1))).cpu().numpy(), } + if "expert_actions" in episode_dict.keys(): + episode_data["expert_actions"] = episode_dict["expert_actions"].cpu().numpy() + episode_data["expert_obs"] = episode_dict["expert_obs"].cpu().numpy() # Save episode data to Zarr self._save_episode_to_zarr(episode_data) diff --git a/source/uwlab_assets/uwlab_assets/robots/ur5e_robotiq_gripper/ur5e_robotiq_2f85_gripper.py b/source/uwlab_assets/uwlab_assets/robots/ur5e_robotiq_gripper/ur5e_robotiq_2f85_gripper.py index 5d52f343..f4557664 100644 --- a/source/uwlab_assets/uwlab_assets/robots/ur5e_robotiq_gripper/ur5e_robotiq_2f85_gripper.py +++ b/source/uwlab_assets/uwlab_assets/robots/ur5e_robotiq_gripper/ur5e_robotiq_2f85_gripper.py @@ -59,7 +59,7 @@ UR5E_ARTICULATION = ArticulationCfg( spawn=sim_utils.UsdFileCfg( usd_path=f"{UWLAB_CLOUD_ASSETS_DIR}/Robots/UniversalRobots/Ur5e2f85RobotiqGripperCalibrated/ur5e_robotiq_gripper_d415_mount_safety_calibrated.usd", - activate_contact_sensors=False, + activate_contact_sensors=True, rigid_props=sim_utils.RigidBodyPropertiesCfg( disable_gravity=True, max_depenetration_velocity=5.0, diff --git a/source/uwlab_rl/uwlab_rl/wrappers/diffusion.py b/source/uwlab_rl/uwlab_rl/wrappers/diffusion.py index 0abad3bd..d98eca1e 100644 --- a/source/uwlab_rl/uwlab_rl/wrappers/diffusion.py +++ b/source/uwlab_rl/uwlab_rl/wrappers/diffusion.py @@ -41,7 +41,6 @@ def reset_envs(self, env_indices: list[int]): class LowDimObservationHistory(ObservationHistoryManager): """Manages observation history for low-dimensional policies.""" - def initialize(self, processed_obs: dict[str, torch.Tensor]): """Initialize history as a single tensor.""" obs_shape = processed_obs["obs"].shape @@ -56,12 +55,12 @@ def update(self, processed_obs: dict[str, torch.Tensor]): if self.needs_init: for env_idx in list(self.needs_init): # Fill entire history with the first observation - first_obs = processed_obs["obs"][env_idx : env_idx + 1] # Keep batch dimension + first_obs = processed_obs["obs"][env_idx:env_idx + 1] # Keep batch dimension for step in range(self.n_obs_steps): self.history[env_idx, step] = first_obs[0] self.needs_init.remove(env_idx) # Update history by shifting and adding new observations - self.history[:, :-1] = self.history[:, 1:].clone() + self.history[:, :-1] = self.history[:, 1:] # Add new observation at the end self.history[:, -1] = processed_obs["obs"] @@ -83,13 +82,13 @@ def reset_envs(self, env_indices: list[int]): class ImageObservationHistory(ObservationHistoryManager): """Manages observation history for image-based policies.""" - def __init__(self, num_envs: int, n_obs_steps: int, device: torch.device): + def __init__(self, num_envs: int, obs_keys: list[str], n_obs_steps: int, device: torch.device): super().__init__(num_envs, n_obs_steps, device) - self.obs_keys = None + self.obs_keys = obs_keys def initialize(self, processed_obs: dict[str, torch.Tensor]): """Initialize history as a dictionary of tensors.""" - self.obs_keys = list(processed_obs.keys()) + # self.obs_keys = list(processed_obs.keys()) self.history = {} for key in self.obs_keys: # Shape: (num_envs, n_obs_steps, ...) @@ -107,7 +106,7 @@ def update(self, processed_obs: dict[str, torch.Tensor]): if env_idx < self.num_envs: # Fill entire history with the first observation for each key for key in self.obs_keys: - first_obs = processed_obs[key][env_idx : env_idx + 1] # Keep batch dimension + first_obs = processed_obs[key][env_idx:env_idx + 1] # Keep batch dimension for step in range(self.n_obs_steps): self.history[key][env_idx, step] = first_obs[0] self.needs_init.remove(env_idx) @@ -136,6 +135,77 @@ def reset_envs(self, env_indices: list[int]): self.needs_init.add(i) +class ImageObservationSequence(ImageObservationHistory): + """Stores full trajectory per-environment as lists and returns padded batches with attention masks. + """ + + def initialize(self, processed_obs: dict[str, torch.Tensor]): + """Initialize internal per-key, per-env lists for histories.""" + # Keep the same obs_keys behaviour + # self.obs_keys = list(processed_obs.keys()) + # history will be a dict: key -> list(len=num_envs) of python lists containing tensors + self.history = {} + for key in self.obs_keys: + self.history[key] = [[] for _ in range(self.num_envs)] + + def update(self, processed_obs: dict[str, torch.Tensor], env_indices: list[int]): + """Append new observations to each environment's trajectory list. + + If an environment is marked in `needs_init` it will have its lists cleared + before appending the new (first) observation so that trajectories start after reset. + """ + if self.history is None: + self.initialize(processed_obs) + + # Clear histories for envs that were reset + if self.needs_init and self.obs_keys is not None: + for env_idx in list(self.needs_init): + if env_idx < self.num_envs: + for key in self.obs_keys: + self.history[key][env_idx].clear() + self.needs_init.remove(env_idx) + + # Append current observation for each env and key + if self.obs_keys is not None: + for key in self.obs_keys: + tensor = processed_obs[key] + # Expect tensor shape (num_envs, ...) + for idx, env_idx in enumerate(env_indices): + obs = tensor[idx] + + # store a detached clone to avoid accidental graph retention + self.history[key][env_idx].append(obs.detach().clone().to(self.device)) + + def get_batch(self, env_indices: list[int]) -> dict[str, torch.Tensor]: + """Return a padded batch of observations for the requested envs plus an attention mask. + + Returns a dict mapping each observation key to a tensor of shape + (batch, seq_len, ...). Also includes 'attention_mask' with shape (batch, seq_len), + where 1 indicates a real observation and 0 indicates padding. + """ + if self.history is None or self.obs_keys is None: + return {} + + # Determine per-env sequence lengths (use the first obs_key as canonical) + lengths = [len(self.history[self.obs_keys[0]][env]) for env in env_indices] + max_len = max(lengths) if lengths else 0 + + obs_batch: dict[str, torch.Tensor] = {} + + lengths_tensor = torch.tensor(lengths, device=self.device) + attention_mask = (torch.arange(max_len, device=self.device).unsqueeze(0) < lengths_tensor.unsqueeze(1)).long() + + for key in self.obs_keys: + obs_batch[key] = torch.nn.utils.rnn.pad_sequence( + [torch.stack(self.history[key][env], dim=0) for env in env_indices], + batch_first=True, + padding_value=0.0 + ) + + obs_batch['attention_mask'] = attention_mask + return obs_batch + + class DiffusionPolicyWrapper: """Wraps diffusion policy to handle Isaac Lab environment observations and action execution.""" @@ -147,19 +217,24 @@ def __init__(self, policy, device: torch.device, n_obs_steps: int = 2, num_envs: device: Device to run the policy on. n_obs_steps: Number of observation steps to maintain in history. num_envs: Number of environments to handle. - execute_horizon: Number of actions to execute from each chunk before - replanning. None = execute full chunk (open-loop). 1 = replan - every step (receding horizon). """ self.policy = policy self.device = device self.n_obs_steps = n_obs_steps self.num_envs = num_envs + # Initialize observation history manager based on policy type self.is_image_policy = self._is_image_policy() - if self.is_image_policy: - self.obs_history_manager = ImageObservationHistory(num_envs, n_obs_steps, device) + self.is_transformer = self._is_transformer() + if hasattr(policy.obs_encoder, 'keys'): + obs_keys = policy.obs_encoder.keys + else: + obs_keys = policy.obs_encoder.rgb_keys + policy.obs_encoder.low_dim_keys + if self.is_transformer: + self.obs_history_manager = ImageObservationSequence(num_envs, obs_keys, n_obs_steps, device) + elif self.is_image_policy: + self.obs_history_manager = ImageObservationHistory(num_envs, obs_keys, n_obs_steps, device) else: self.obs_history_manager = LowDimObservationHistory(num_envs, n_obs_steps, device) @@ -169,15 +244,21 @@ def __init__(self, policy, device: torch.device, n_obs_steps: int = 2, num_envs: # Reset the policy to initialize its internal queues self.policy.reset() + def _is_transformer(self) -> bool: + """Detect if the policy is a transformer-based model.""" + policy_class_name = self.policy.__class__.__name__.lower() + transformer_indicators = ['transformer', 'gpt', 'bert', 'dpt'] + return any(indicator in policy_class_name for indicator in transformer_indicators) + def _is_image_policy(self) -> bool: """Detect if this is an image policy based on class name.""" policy_class_name = self.policy.__class__.__name__.lower() - image_policy_indicators = ["image", "hybrid", "video"] + image_policy_indicators = ['image', 'hybrid', 'video'] return any(indicator in policy_class_name for indicator in image_policy_indicators) def reset(self, reset_ids: torch.Tensor): """Reset the policy wrapper and clear observation history and action queue.""" - reset_indices = reset_ids.tolist() if hasattr(reset_ids, "tolist") else reset_ids + reset_indices = reset_ids.tolist() if hasattr(reset_ids, 'tolist') else reset_ids for i in reset_indices: self.action_queue[i].clear() @@ -187,7 +268,7 @@ def reset(self, reset_ids: torch.Tensor): self.obs_history_manager.reset_envs(reset_indices) self.policy.reset() - def predict_action(self, obs_dict: dict[str, Any]) -> torch.Tensor: + def predict_action(self, obs_dict: dict[str, Any], env_indices: list[int] = None) -> torch.Tensor: """Predict action given Isaac Lab environment observations. Args: @@ -200,10 +281,15 @@ def predict_action(self, obs_dict: dict[str, Any]) -> torch.Tensor: processed_obs = self._process_obs(obs_dict) # Update observation history with batched operations - self.obs_history_manager.update(processed_obs) - - # Find environments that need new action chunks - need_new_actions = [i for i in range(self.num_envs) if len(self.action_queue[i]) == 0] + if self.is_transformer: + if env_indices is None: + env_indices = list(range(self.num_envs)) + self.obs_history_manager.update(processed_obs, env_indices) + need_new_actions = [i for i in range(self.num_envs) if len(self.action_queue[i]) == 0 and i in env_indices] + else: + self.obs_history_manager.update(processed_obs) + # Find environments that need new action chunks + need_new_actions = [i for i in range(self.num_envs) if len(self.action_queue[i]) == 0] if need_new_actions: # Get new action chunks for environments that need them @@ -213,10 +299,16 @@ def predict_action(self, obs_dict: dict[str, Any]) -> torch.Tensor: for idx, env_idx in enumerate(need_new_actions): self.action_queue[env_idx].extend(new_actions[idx]) - # Extract next action for each environment - actions = torch.zeros(self.num_envs, self.action_queue[0][0].shape[-1], device=self.device, dtype=torch.float32) - for i in range(self.num_envs): - actions[i] = self.action_queue[i].pop(0) + if self.is_transformer: + actions = torch.zeros(len(env_indices), self.action_queue[env_indices[0]][0].shape[-1], device=self.device, dtype=torch.float32) + for idx, env_idx in enumerate(env_indices): + actions[idx] = self.action_queue[env_idx].pop(0) + return actions + else: + # Extract next action for each environment + actions = torch.zeros(self.num_envs, self.action_queue[0][0].shape[-1], device=self.device, dtype=torch.float32) + for i in range(self.num_envs): + actions[i] = self.action_queue[i].pop(0) return actions @@ -301,9 +393,10 @@ def _get_action_chunks(self, env_indices: list[int]) -> list[torch.Tensor]: obs_batch = self.obs_history_manager.get_batch(env_indices) # Get action chunk from policy - result = self.policy.predict_action(obs_batch) + with torch.no_grad(): + result = self.policy.predict_action(obs_batch) if isinstance(result, dict): - action_chunk = result["action"] + action_chunk = result['action'] else: action_chunk = result @@ -315,8 +408,7 @@ def _get_action_chunks(self, env_indices: list[int]) -> list[torch.Tensor]: env_action_chunk = action_chunk[i] # Shape: (action_chunk_len, action_dim) action_chunks.append(env_action_chunk) else: - # Single action case: (batch_size, action_dim) -> list of (1, action_dim) per env - for i in range(action_chunk.shape[0]): - action_chunks.append(action_chunk[i].unsqueeze(0)) + # Single action case: (batch_size, action_dim) -> (batch_size, 1, action_dim) + action_chunks = action_chunk.unsqueeze(1) return action_chunks diff --git a/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/__init__.py b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/__init__.py new file mode 100644 index 00000000..700ad0c5 --- /dev/null +++ b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/__init__.py @@ -0,0 +1,28 @@ +# Copyright (c) 2024-2026, The UW Lab Project Developers. (https://github.com/uw-lab/UWLab/blob/main/CONTRIBUTORS.md). +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""ASTEROID environments. + +Pick-only variants of the OmniReset UR5e + Robotiq 2F-85 tasks used by ASTEROID +(iterative in-context exploration + distillation, see ``scripts/ASTEROID``). + +Everything here builds on :mod:`uwlab_tasks.manager_based.manipulation.omnireset`: +scene, MDP terms and environment configs are subclassed and only the pick-specific +deltas live in this package (no receptive object, pick-height success, coupled +domain randomization, tactile data-collection observations). +""" + +import os + +ASTEROID_DATASETS_DIR = os.environ.get("ASTEROID_DATASETS_DIR", "Datasets/CubePick") +"""Root of the locally recorded reset-state / grasp datasets used by the ASTEROID configs. + +Layout mirrors the OmniReset asset hub, keyed by the insertive object only:: + + /Resets//resets_.pt + /Grasps//grasps.pt + +Override with the ``ASTEROID_DATASETS_DIR`` environment variable. +""" diff --git a/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/config/__init__.py b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/config/__init__.py new file mode 100644 index 00000000..726b9e49 --- /dev/null +++ b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/config/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) 2024-2026, The UW Lab Project Developers. (https://github.com/uw-lab/UWLab/blob/main/CONTRIBUTORS.md). +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Configurations for ASTEROID environments.""" + +# We leave this file empty since we don't want to expose any configs in this package directly. +# We still need this file to import the "config" module in the parent package. diff --git a/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/config/ur5e_robotiq_2f85/__init__.py b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/config/ur5e_robotiq_2f85/__init__.py new file mode 100644 index 00000000..54a36812 --- /dev/null +++ b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/config/ur5e_robotiq_2f85/__init__.py @@ -0,0 +1,127 @@ +# Copyright (c) 2024-2026, The UW Lab Project Developers. (https://github.com/uw-lab/UWLab/blob/main/CONTRIBUTORS.md). +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import gymnasium as gym + +from . import agents + +## +# Reset-state recording environments +## + +gym.register( + id="Asteroid-UR5eRobotiq2f85-ObjectAnywhereEEAnywhere-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={"env_cfg_entry_point": f"{__name__}.reset_states_cfg:PickObjectAnywhereEEAnywhereResetStatesCfg"}, +) + +gym.register( + id="Asteroid-UR5eRobotiq2f85-ObjectRestingEEGrasped-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={"env_cfg_entry_point": f"{__name__}.reset_states_cfg:PickObjectRestingEEGraspedResetStatesCfg"}, +) + +gym.register( + id="Asteroid-UR5eRobotiq2f85-ObjectAnywhereEEGrasped-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={"env_cfg_entry_point": f"{__name__}.reset_states_cfg:PickObjectAnywhereEEGraspedResetStatesCfg"}, +) + +## +# State-based RL (expert) environments +## + +gym.register( + id="Asteroid-Ur5eRobotiq2f85-RelCartesianOSC-State-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.rl_state_cfg:Ur5eRobotiq2f85PickRelCartesianOSCTrainCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_cfg:AsteroidPPORunnerCfg", + }, +) + +gym.register( + id="Asteroid-Ur5eRobotiq2f85-RelCartesianOSC-State-Finetune-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.rl_state_cfg:Ur5eRobotiq2f85PickRelCartesianOSCFinetuneCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_cfg:AsteroidPPORunnerCfg", + }, +) + +gym.register( + id="Asteroid-Ur5eRobotiq2f85-RelCartesianOSC-State-Play-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.rl_state_cfg:Ur5eRobotiq2f85PickRelCartesianOSCEvalCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_cfg:AsteroidPPORunnerCfg", + }, +) + +gym.register( + id="Asteroid-Ur5eRobotiq2f85-RelCartesianOSC-State-Finetune-Play-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.rl_state_cfg:Ur5eRobotiq2f85PickRelCartesianOSCFinetuneEvalCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_cfg:AsteroidPPORunnerCfg", + }, +) + +## +# Tactile (proprioceptive student) data collection / evaluation environments +## + +gym.register( + id="Asteroid-Ur5eRobotiq2f85-RelCartesianOSC-Tactile-DataCollection-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": ( + f"{__name__}.data_collection_tactile_cfg:Ur5eRobotiq2f85DataCollectionTactileRelCartesianOSCCfg" + ), + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_cfg:AsteroidDAggerRunnerCfg", + }, +) + +gym.register( + id="Asteroid-Ur5eRobotiq2f85-RelCartesianOSC-Tactile-Finetune-DataCollection-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": ( + f"{__name__}.data_collection_tactile_cfg:Ur5eRobotiq2f85DataCollectionFinetuneTactileRelCartesianOSCCfg" + ), + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_cfg:AsteroidDAggerRunnerCfg", + }, +) + +gym.register( + id="Asteroid-Ur5eRobotiq2f85-RelCartesianOSC-Tactile-Play-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": f"{__name__}.data_collection_tactile_cfg:Ur5eRobotiq2f85EvalTactileRelCartesianOSCCfg", + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_cfg:AsteroidDAggerRunnerCfg", + }, +) + +gym.register( + id="Asteroid-Ur5eRobotiq2f85-RelCartesianOSC-Tactile-Finetune-Play-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + disable_env_checker=True, + kwargs={ + "env_cfg_entry_point": ( + f"{__name__}.data_collection_tactile_cfg:Ur5eRobotiq2f85EvalFinetuneTactileRelCartesianOSCCfg" + ), + "rsl_rl_cfg_entry_point": f"{agents.__name__}.rsl_rl_cfg:AsteroidDAggerRunnerCfg", + }, +) diff --git a/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/config/ur5e_robotiq_2f85/actions.py b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/config/ur5e_robotiq_2f85/actions.py new file mode 100644 index 00000000..7f00ba50 --- /dev/null +++ b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/config/ur5e_robotiq_2f85/actions.py @@ -0,0 +1,65 @@ +# Copyright (c) 2024-2026, The UW Lab Project Developers. (https://github.com/uw-lab/UWLab/blob/main/CONTRIBUTORS.md). +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Action configs for the ASTEROID UR5e + Robotiq 2F-85 tasks. + +The default 6-DOF + gripper actions are OmniReset's (re-exported here); this module adds +position-only (3-DOF + gripper) variants. +""" + +from isaaclab.utils import configclass + +from uwlab_assets.robots.ur5e_robotiq_gripper.actions import ROBOTIQ_GRIPPER_BINARY_ACTIONS + +from uwlab_tasks.manager_based.manipulation.omnireset.config.ur5e_robotiq_2f85.actions import ( # noqa: F401 + Ur5eRobotiq2f85RelativeOSCAction, + Ur5eRobotiq2f85RelativeOSCEvalAction, +) + +from ...mdp.actions.actions_cfg import RelCartesianOSCPositionActionCfg + +# Position-only gains (mirrors the pre-train OSC gains; the policy no longer +# commands rotation and the wrist is free to rotate under collisions). +UR5E_ROBOTIQ_2F85_RELATIVE_OSC_POSONLY = RelCartesianOSCPositionActionCfg( + asset_name="robot", + joint_names=["shoulder.*", "elbow.*", "wrist.*"], + body_name="wrist_3_link", + scale_xyz_axisangle=(0.02, 0.02, 0.02, 0.02, 0.02, 0.2), + motion_stiffness=(200.0, 200.0, 200.0, 3.0, 3.0, 3.0), + motion_damping_ratio=(3.0, 3.0, 3.0, 1.0, 1.0, 1.0), + torque_limit=(150.0, 150.0, 150.0, 28.0, 28.0, 28.0), +) + +# Position-only eval / sim2real gains (end-of-curriculum values). +UR5E_ROBOTIQ_2F85_RELATIVE_OSC_EVAL_POSONLY = RelCartesianOSCPositionActionCfg( + asset_name="robot", + joint_names=["shoulder.*", "elbow.*", "wrist.*"], + body_name="wrist_3_link", + scale_xyz_axisangle=(0.01, 0.01, 0.002, 0.02, 0.02, 0.2), + motion_stiffness=(1000.0, 1000.0, 1000.0, 50.0, 50.0, 50.0), + motion_damping_ratio=(1.0, 1.0, 1.0, 1.0, 1.0, 1.0), + torque_limit=(150.0, 150.0, 150.0, 28.0, 28.0, 28.0), +) + + +@configclass +class Ur5eRobotiq2f85RelativeOSCPositionAction: + """Position-only action: 3-DOF Cartesian (x, y, z) arm + binary gripper. + + The policy cannot command wrist rotation -- only translate and open/close the gripper. + Orientation is left uncommanded, so the wrist is free to rotate under collisions. + Total action dim is 4 (3 arm + 1 gripper). + """ + + arm = UR5E_ROBOTIQ_2F85_RELATIVE_OSC_POSONLY + gripper = ROBOTIQ_GRIPPER_BINARY_ACTIONS + + +@configclass +class Ur5eRobotiq2f85RelativeOSCEvalPositionAction: + """Position-only action with high Kp gains (end-of-curriculum values) for eval / data-collection.""" + + arm = UR5E_ROBOTIQ_2F85_RELATIVE_OSC_EVAL_POSONLY + gripper = ROBOTIQ_GRIPPER_BINARY_ACTIONS diff --git a/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/config/ur5e_robotiq_2f85/agents/__init__.py b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/config/ur5e_robotiq_2f85/agents/__init__.py new file mode 100644 index 00000000..202d14d3 --- /dev/null +++ b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/config/ur5e_robotiq_2f85/agents/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) 2024-2026, The UW Lab Project Developers. (https://github.com/uw-lab/UWLab/blob/main/CONTRIBUTORS.md). +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from . import rsl_rl_cfg diff --git a/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/config/ur5e_robotiq_2f85/agents/rsl_rl_cfg.py b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/config/ur5e_robotiq_2f85/agents/rsl_rl_cfg.py new file mode 100644 index 00000000..72153ce7 --- /dev/null +++ b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/config/ur5e_robotiq_2f85/agents/rsl_rl_cfg.py @@ -0,0 +1,58 @@ +# Copyright (c) 2024-2026, The UW Lab Project Developers. (https://github.com/uw-lab/UWLab/blob/main/CONTRIBUTORS.md). +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab.utils import configclass + +from uwlab_rl.rsl_rl.rl_cfg import BehaviorCloningCfg, OffPolicyAlgorithmCfg, RslRlFancyPpoAlgorithmCfg + +from uwlab_tasks.manager_based.manipulation.omnireset.config.ur5e_robotiq_2f85.agents.rsl_rl_cfg import ( + Base_PPORunnerCfg, + my_experts_observation_func, +) + + +@configclass +class AsteroidPPORunnerCfg(Base_PPORunnerCfg): + """PPO runner for the pick-only state expert (same hyperparameters as OmniReset).""" + + experiment_name = "ur5e_robotiq_2f85_asteroid_agent" + + +@configclass +class AsteroidDAggerRunnerCfg(AsteroidPPORunnerCfg): + """PPO + behavior-cloning runner; the expert observes the pick-only policy observation group.""" + + algorithm = RslRlFancyPpoAlgorithmCfg( + value_loss_coef=1.0, + use_clipped_value_loss=True, + normalize_advantage_per_mini_batch=False, + clip_param=0.2, + entropy_coef=0.006, + num_learning_epochs=5, + num_mini_batches=4, + learning_rate=1.0e-4, + schedule="adaptive", + gamma=0.99, + lam=0.95, + desired_kl=0.01, + max_grad_norm=1.0, + offline_algorithm_cfg=OffPolicyAlgorithmCfg( + behavior_cloning_cfg=BehaviorCloningCfg( + experts_path=[""], + experts_loader="torch.jit.load", + experts_observation_group_cfg=( + "uwlab_tasks.manager_based.manipulation.asteroid.config.ur5e_robotiq_2f85.rl_state_cfg" + ":PickObservationsCfg.PolicyCfg" + ), + experts_observation_func=my_experts_observation_func, + experts_action_group_cfg=( + "uwlab_tasks.manager_based.manipulation.asteroid.config.ur5e_robotiq_2f85.actions" + ":Ur5eRobotiq2f85RelativeOSCAction" + ), + cloning_loss_coeff=1.0, + loss_decay=1.0, + ) + ), + ) diff --git a/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/config/ur5e_robotiq_2f85/data_collection_tactile_cfg.py b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/config/ur5e_robotiq_2f85/data_collection_tactile_cfg.py new file mode 100644 index 00000000..b7bec9b6 --- /dev/null +++ b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/config/ur5e_robotiq_2f85/data_collection_tactile_cfg.py @@ -0,0 +1,251 @@ +# Copyright (c) 2024-2026, The UW Lab Project Developers. (https://github.com/uw-lab/UWLab/blob/main/CONTRIBUTORS.md). +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Proprioceptive ("tactile") data-collection and evaluation environments for ASTEROID. + +The student policy sees only proprioception (arm joints, EE pose, last actions and a +normalized gripper-position reading with calibration-drift randomization); the state expert +that supervises it sees the pick-only policy observation group. Used by +``scripts/ASTEROID/collect_demos_asteroid.py`` and ``scripts/ASTEROID/eval_asteroid_policy.py``. +""" + +from __future__ import annotations + +import isaaclab.sim as sim_utils +from isaaclab.managers import EventTermCfg as EventTerm +from isaaclab.managers import ObservationGroupCfg as ObsGroup +from isaaclab.managers import ObservationTermCfg as ObsTerm +from isaaclab.managers import SceneEntityCfg +from isaaclab.managers import TerminationTermCfg as DoneTerm +from isaaclab.sensors import TiledCameraCfg +from isaaclab.utils import configclass + +from ... import mdp as task_mdp +from .actions import Ur5eRobotiq2f85RelativeOSCAction, Ur5eRobotiq2f85RelativeOSCEvalAction +from .rl_state_cfg import PickFinetuneEvalEventCfg, PickObservationsCfg, PickSceneCfg, Ur5eRobotiq2f85PickStateCfg + +## +# Scene +## + + +@configclass +class TactileSceneCfg(PickSceneCfg): + """Pick scene for data collection (no cameras; the student is proprioceptive).""" + + # TODO: add fingertip force/torque sensors once the real-robot counterpart is available. + + +@configclass +class TactileEvalSceneCfg(TactileSceneCfg): + """Adds a high-resolution front camera for evaluation videos.""" + + front_camera = TiledCameraCfg( + prim_path="{ENV_REGEX_NS}/Robot/rgb_front_camera", + update_period=0, + height=1080, + width=1920, + offset=TiledCameraCfg.OffsetCfg( + pos=(1.0770121, -0.21290445, 0.4486344), + rot=(0.70564552, 0.46613815, 0.25072644, 0.47107948), + convention="opengl", + ), + data_types=["rgb"], + spawn=sim_utils.PinholeCameraCfg(focal_length=13.20), + ) + + +## +# Events +## + + +@configclass +class TactileEventCfg(PickFinetuneEvalEventCfg): + """Fixed sysid + OSC gains, 1-path resets, plus gripper-reading calibration randomization.""" + + randomize_gripper_pos_affine = EventTerm( + func=task_mdp.randomize_gripper_pos_affine, + mode="reset", + params={ + "scale_range": (0.9, 1.1), + "offset_range": (-0.03, 0.03), + }, + ) + + +## +# Observations +## + +_GRIPPER_POS_TERM = dict( + func=task_mdp.gripper_pos_normalized, + params={ + "asset_cfg": SceneEntityCfg("robot", joint_names=["left_inner_finger_knuckle_joint"]), + "scale_event_name": "randomize_gripper_pos_affine", + "jitter_std": 0.01, + }, +) + + +@configclass +class TactileObservationsCfg: + @configclass + class TactilePolicyCfg(ObsGroup): + """Student (diffusion) policy obs -- dict-form; keys must match the ``shape_meta`` of the + diffusion-policy task config so the student receives the layout it was trained on.""" + + last_gripper_action = ObsTerm(func=task_mdp.last_action, params={"action_name": "gripper"}) + + last_arm_action = ObsTerm(func=task_mdp.last_action, params={"action_name": "arm"}) + + arm_joint_pos = ObsTerm( + func=task_mdp.joint_pos, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=["shoulder.*", "elbow.*", "wrist.*"])}, + ) + + end_effector_pose = ObsTerm( + func=task_mdp.target_asset_pose_in_root_asset_frame, + params={ + "target_asset_cfg": SceneEntityCfg("robot", body_names="wrist_3_link"), + "root_asset_cfg": SceneEntityCfg("robot"), + "rotation_repr": "axis_angle", + }, + ) + + gripper_pos = ObsTerm(**_GRIPPER_POS_TERM) + + def __post_init__(self): + self.enable_corruption = True + self.concatenate_terms = False + + @configclass + class TactileDataCollectionCfg(ObsGroup): + """Observations recorded to the dataset (policy obs + raw gripper joint angles).""" + + last_gripper_action = ObsTerm(func=task_mdp.last_action, params={"action_name": "gripper"}) + + last_arm_action = ObsTerm(func=task_mdp.last_action, params={"action_name": "arm"}) + + arm_joint_pos = ObsTerm( + func=task_mdp.joint_pos, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=["shoulder.*", "elbow.*", "wrist.*"])}, + ) + + end_effector_pose = ObsTerm( + func=task_mdp.target_asset_pose_in_root_asset_frame, + params={ + "target_asset_cfg": SceneEntityCfg("robot", body_names="wrist_3_link"), + "root_asset_cfg": SceneEntityCfg("robot"), + "rotation_repr": "axis_angle", + }, + ) + + gripper_joint_pos = ObsTerm( + func=task_mdp.joint_pos, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=[".*_inner_finger_knuckle_joint"])}, + ) + + gripper_pos = ObsTerm(**_GRIPPER_POS_TERM) + + def __post_init__(self): + self.enable_corruption = True + self.concatenate_terms = False + + # observation groups + policy: TactilePolicyCfg = TactilePolicyCfg() + data_collection: TactileDataCollectionCfg = TactileDataCollectionCfg() + # Privileged state observations consumed by the JIT-loaded state expert during BC + # supervision -- identical to the group the expert was trained on. Read by + # ``my_experts_observation_func`` via ``env.unwrapped.obs_buf["expert_obs"]``. + expert_obs: PickObservationsCfg.PolicyCfg = PickObservationsCfg.PolicyCfg() + + +## +# Terminations +## + + +@configclass +class TactileTerminationsCfg: + time_out = DoneTerm(func=task_mdp.time_out, time_out=True) + + abnormal_robot = DoneTerm(func=task_mdp.abnormal_robot_state) + + early_success = DoneTerm( + func=task_mdp.early_success_termination, params={"num_consecutive_successes": 10, "min_episode_length": 10} + ) + + success = DoneTerm( + func=task_mdp.consecutive_success_state_with_min_length, + params={"num_consecutive_successes": 10, "min_episode_length": 10}, + ) + + +## +# Environments +## + + +@configclass +class Ur5eRobotiq2f85TactileRelCartesianOSCEvalCfg(Ur5eRobotiq2f85PickStateCfg): + """Tactile base config: fixed sysid + tactile scene / obs / terminations / render.""" + + actions: Ur5eRobotiq2f85RelativeOSCAction = Ur5eRobotiq2f85RelativeOSCAction() + scene: TactileSceneCfg = TactileSceneCfg(num_envs=32, env_spacing=1.5, replicate_physics=False) + observations: TactileObservationsCfg = TactileObservationsCfg() + terminations: TactileTerminationsCfg = TactileTerminationsCfg() + events: TactileEventCfg = TactileEventCfg() + + def __post_init__(self): + super().__post_init__() + + self.episode_length_s = 10.0 + + # speeds up rendering + self.sim.render_interval = self.decimation + + # rerender on reset + self.num_rerenders_on_reset = 1 + + +@configclass +class Ur5eRobotiq2f85DataCollectionTactileRelCartesianOSCCfg(Ur5eRobotiq2f85TactileRelCartesianOSCEvalCfg): + """Data collection with the Stage 1 (soft-gain) expert.""" + + actions: Ur5eRobotiq2f85RelativeOSCAction = Ur5eRobotiq2f85RelativeOSCAction() + + +@configclass +class Ur5eRobotiq2f85DataCollectionFinetuneTactileRelCartesianOSCCfg(Ur5eRobotiq2f85TactileRelCartesianOSCEvalCfg): + """Data collection with the Stage 2 (stiff-gain, finetuned) expert.""" + + actions: Ur5eRobotiq2f85RelativeOSCEvalAction = Ur5eRobotiq2f85RelativeOSCEvalAction() + + +@configclass +class Ur5eRobotiq2f85EvalTactileRelCartesianOSCCfg(Ur5eRobotiq2f85TactileRelCartesianOSCEvalCfg): + """Evaluation of a Stage 1 student, with a front camera for videos.""" + + scene: TactileEvalSceneCfg = TactileEvalSceneCfg(num_envs=32, env_spacing=1.5, replicate_physics=False) + + def __post_init__(self): + super().__post_init__() + self.observations.policy.front_rgb = ObsTerm( + func=task_mdp.process_image, + params={ + "sensor_cfg": SceneEntityCfg("front_camera"), + "data_type": "rgb", + "process_image": True, + "output_size": (1080, 1920), + }, + ) + + +@configclass +class Ur5eRobotiq2f85EvalFinetuneTactileRelCartesianOSCCfg(Ur5eRobotiq2f85EvalTactileRelCartesianOSCCfg): + """Evaluation of a Stage 2 student (stiff gains), with a front camera for videos.""" + + actions: Ur5eRobotiq2f85RelativeOSCEvalAction = Ur5eRobotiq2f85RelativeOSCEvalAction() diff --git a/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/config/ur5e_robotiq_2f85/reset_states_cfg.py b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/config/ur5e_robotiq_2f85/reset_states_cfg.py new file mode 100644 index 00000000..eef3c636 --- /dev/null +++ b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/config/ur5e_robotiq_2f85/reset_states_cfg.py @@ -0,0 +1,251 @@ +# Copyright (c) 2024-2026, The UW Lab Project Developers. (https://github.com/uw-lab/UWLab/blob/main/CONTRIBUTORS.md). +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Reset-state recording environments for the pick-only cube task. + +Subclasses the OmniReset reset-state configs, removes the receptive object and narrows the +object / end-effector sampling regions to a cube resting on the table in front of the robot. +Recorded datasets go to ``/Resets//`` (see +:mod:`uwlab_tasks.manager_based.manipulation.asteroid`). +""" + +from __future__ import annotations + +import numpy as np + +from isaaclab.managers import EventTermCfg as EventTerm +from isaaclab.managers import SceneEntityCfg +from isaaclab.managers import TerminationTermCfg as DoneTerm +from isaaclab.utils import configclass + +from uwlab_assets import UWLAB_CLOUD_ASSETS_DIR + +from uwlab_tasks.manager_based.manipulation.omnireset.config.ur5e_robotiq_2f85.reset_states_cfg import ( + ObjectAnywhereEEAnywhereEventCfg, + ObjectAnywhereEEGraspedEventCfg, + ObjectRestingEEGraspedEventCfg, + ResetStatesSceneCfg, + ResetStatesTerminationCfg, + UR5eRobotiq2f85ResetStatesCfg, + make_insertive_object, +) + +from ... import ASTEROID_DATASETS_DIR +from ... import mdp as task_mdp + +## +# Scene +## + +INSERTIVE_OBJECT_VARIANTS = { + "fbleg": make_insertive_object(f"{UWLAB_CLOUD_ASSETS_DIR}/Props/FurnitureBench/SquareLeg/square_leg.usd"), + "fbdrawerbottom": make_insertive_object( + f"{UWLAB_CLOUD_ASSETS_DIR}/Props/FurnitureBench/DrawerBottom/drawer_bottom.usd" + ), + "peg": make_insertive_object(f"{UWLAB_CLOUD_ASSETS_DIR}/Props/Custom/Peg/peg.usd"), + "cupcake": make_insertive_object(f"{UWLAB_CLOUD_ASSETS_DIR}/Props/Custom/CupCake/cupcake.usd"), + "cube": make_insertive_object(f"{UWLAB_CLOUD_ASSETS_DIR}/Props/Custom/InsertiveCube/insertive_cube.usd"), + "rectangle": make_insertive_object(f"{UWLAB_CLOUD_ASSETS_DIR}/Props/Custom/Rectangle/rectangle.usd"), +} + +variants = {"scene.insertive_object": INSERTIVE_OBJECT_VARIANTS} + + +@configclass +class PickResetStatesSceneCfg(ResetStatesSceneCfg): + """OmniReset reset-state scene without the receptive object; cube by default.""" + + receptive_object = None + insertive_object = make_insertive_object(f"{UWLAB_CLOUD_ASSETS_DIR}/Props/Custom/InsertiveCube/insertive_cube.usd") + + +## +# Events +## + + +@configclass +class PickObjectAnywhereEEAnywhereEventCfg(ObjectAnywhereEEAnywhereEventCfg): + """Cube resting flat on the table in front of the robot; EE hovering above it, pointing down.""" + + receptive_object_material = None + reset_receptive_object_pose = None + + reset_insertive_object_pose = EventTerm( + func=task_mdp.reset_root_states_uniform, + mode="reset", + params={ + "pose_range": { + "x": (0.4, 0.5), + "y": (0.05, 0.15), + "z": (0.01, 0.02), + "roll": (0.0, 0.0), + "pitch": (0.0, 0.0), + "yaw": (-np.pi / 8, np.pi / 8), + }, + "velocity_range": {}, + "asset_cfgs": {"insertive_object": SceneEntityCfg("insertive_object")}, + "offset_asset_cfg": SceneEntityCfg("ur5_metal_support"), + "use_bottom_offset": True, + }, + ) + + reset_end_effector_pose = EventTerm( + func=task_mdp.reset_end_effector_round_fixed_asset, + mode="reset", + params={ + "fixed_asset_cfg": SceneEntityCfg("robot"), + "fixed_asset_offset": None, + "pose_range_b": { + "x": (0.44, 0.46), + "y": (0.09, 0.11), + "z": (0.18, 0.2), + "roll": (0.0, 0.0), + "pitch": (np.pi / 2 - 0.1, np.pi / 2 + 0.1), + "yaw": (np.pi - 0.1, np.pi + 0.1), + }, + "robot_ik_cfg": SceneEntityCfg( + "robot", joint_names=["shoulder.*", "elbow.*", "wrist.*"], body_names="robotiq_base_link" + ), + }, + ) + + +@configclass +class PickObjectRestingEEGraspedEventCfg(ObjectRestingEEGraspedEventCfg): + receptive_object_material = None + reset_receptive_object_pose = None + + reset_insertive_object_pose_from_reset_states = EventTerm( + func=task_mdp.SingleObjectMultiResetManager, + mode="reset", + params={ + "dataset_dir": ASTEROID_DATASETS_DIR, + "reset_types": ["ObjectAnywhereEEAnywhere"], + "probs": [1.0], + }, + ) + + reset_end_effector_pose_from_grasp_dataset = EventTerm( + func=task_mdp.reset_end_effector_from_grasp_dataset, + mode="reset", + params={ + "dataset_dir": ASTEROID_DATASETS_DIR, + "fixed_asset_cfg": SceneEntityCfg("insertive_object"), + "robot_ik_cfg": SceneEntityCfg( + "robot", joint_names=["shoulder.*", "elbow.*", "wrist.*"], body_names="robotiq_base_link" + ), + "gripper_cfg": SceneEntityCfg("robot", joint_names=["finger_joint", ".*right.*", ".*left.*"]), + "pose_range_b": { + "x": (-0.02, 0.02), + "y": (-0.02, 0.02), + "z": (-0.02, 0.02), + "roll": (-np.pi / 16, np.pi / 16), + "pitch": (-np.pi / 16, np.pi / 16), + "yaw": (-np.pi / 16, np.pi / 16), + }, + }, + ) + + +@configclass +class PickObjectAnywhereEEGraspedEventCfg(ObjectAnywhereEEGraspedEventCfg): + receptive_object_material = None + reset_receptive_object_pose = None + + reset_end_effector_pose_from_grasp_dataset = EventTerm( + func=task_mdp.reset_end_effector_from_grasp_dataset, + mode="reset", + params={ + "dataset_dir": ASTEROID_DATASETS_DIR, + "fixed_asset_cfg": SceneEntityCfg("insertive_object"), + "robot_ik_cfg": SceneEntityCfg( + "robot", joint_names=["shoulder.*", "elbow.*", "wrist.*"], body_names="robotiq_base_link" + ), + "gripper_cfg": SceneEntityCfg("robot", joint_names=["finger_joint", ".*right.*", ".*left.*"]), + "pose_range_b": { + "x": (0.0, 0.0), + "y": (0.0, 0.0), + "z": (0.0, 0.0), + "roll": (0.0, 0.0), + "pitch": (0.0, 0.0), + "yaw": (0.0, 0.0), + }, + }, + ) + + +## +# Terminations +## + + +@configclass +class PickResetStatesTerminationCfg(ResetStatesTerminationCfg): + """Reset-state validity check against the insertive object only.""" + + success = DoneTerm( + func=task_mdp.check_reset_state_success, + params={ + "object_cfgs": [SceneEntityCfg("insertive_object")], + "robot_cfg": SceneEntityCfg("robot"), + "ee_body_name": "robotiq_base_link", + "collision_analyzer_cfgs": [ + task_mdp.CollisionAnalyzerCfg( + num_points=1024, + max_dist=0.5, + min_dist=-0.0005, + asset_cfg=SceneEntityCfg("robot"), + obstacle_cfgs=[SceneEntityCfg("insertive_object")], + ), + ], + "max_robot_pos_deviation": 0.1, + "max_object_pos_deviation": np.inf, + "pos_z_threshold": -0.02, + "consecutive_stability_steps": 5, + }, + time_out=True, + ) + + +## +# Environments +## + + +@configclass +class PickResetStatesCfg(UR5eRobotiq2f85ResetStatesCfg): + """Base reset-state environment for the pick-only task.""" + + scene: PickResetStatesSceneCfg = PickResetStatesSceneCfg(num_envs=1, env_spacing=1.5) + terminations: PickResetStatesTerminationCfg = PickResetStatesTerminationCfg() + variants = variants + + +@configclass +class PickObjectAnywhereEEAnywhereResetStatesCfg(PickResetStatesCfg): + events: PickObjectAnywhereEEAnywhereEventCfg = PickObjectAnywhereEEAnywhereEventCfg() + + def __post_init__(self): + super().__post_init__() + self.terminations.success.params["max_object_pos_deviation"] = np.inf + + +@configclass +class PickObjectRestingEEGraspedResetStatesCfg(PickResetStatesCfg): + events: PickObjectRestingEEGraspedEventCfg = PickObjectRestingEEGraspedEventCfg() + + def __post_init__(self): + super().__post_init__() + self.terminations.success.params["max_object_pos_deviation"] = 0.01 + + +@configclass +class PickObjectAnywhereEEGraspedResetStatesCfg(PickResetStatesCfg): + events: PickObjectAnywhereEEGraspedEventCfg = PickObjectAnywhereEEGraspedEventCfg() + + def __post_init__(self): + super().__post_init__() + self.terminations.success.params["max_object_pos_deviation"] = 0.05 diff --git a/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/config/ur5e_robotiq_2f85/rl_state_cfg.py b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/config/ur5e_robotiq_2f85/rl_state_cfg.py new file mode 100644 index 00000000..7dc484c3 --- /dev/null +++ b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/config/ur5e_robotiq_2f85/rl_state_cfg.py @@ -0,0 +1,336 @@ +# Copyright (c) 2024-2026, The UW Lab Project Developers. (https://github.com/uw-lab/UWLab/blob/main/CONTRIBUTORS.md). +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""State-based RL environments for the pick-only cube task (ASTEROID expert training). + +Everything is derived from the OmniReset state environment; this module only expresses the +pick-specific deltas: + +* no receptive object (scene, observations, events), +* pick-height + gripper-down success instead of assembly alignment, +* coupled sysid / OSC-gain / action-scale domain randomization during training, +* reset states loaded from the local ASTEROID dataset directory. +""" + +from __future__ import annotations + +import isaaclab.sim as sim_utils +from isaaclab.assets import RigidObjectCfg +from isaaclab.managers import EventTermCfg as EventTerm +from isaaclab.managers import RewardTermCfg as RewTerm +from isaaclab.managers import SceneEntityCfg +from isaaclab.utils import configclass + +from uwlab_assets import UWLAB_CLOUD_ASSETS_DIR +from uwlab_assets.robots.ur5e_robotiq_gripper import EXPLICIT_UR5E_ROBOTIQ_2F85 + +from uwlab_tasks.manager_based.manipulation.omnireset.config.ur5e_robotiq_2f85.rl_state_cfg import ( + BaseEventCfg, + FinetuneCurriculumsCfg, + FinetuneEvalEventCfg, + FinetuneEventCfg, + ObservationsCfg, + RewardsCfg, + RlStateSceneCfg, + Ur5eRobotiq2f85RlStateCfg, +) + +from ... import ASTEROID_DATASETS_DIR +from ... import mdp as task_mdp +from .actions import Ur5eRobotiq2f85RelativeOSCAction, Ur5eRobotiq2f85RelativeOSCEvalAction + +## +# Scene +## + + +def make_insertive_object(usd_path: str) -> RigidObjectCfg: + """Insertive object with a stiffer solver than OmniReset's (16/2 vs 4/0 iterations).""" + return RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/InsertiveObject", + spawn=sim_utils.UsdFileCfg( + usd_path=usd_path, + scale=(1, 1, 1), + rigid_props=sim_utils.RigidBodyPropertiesCfg( + solver_position_iteration_count=16, + solver_velocity_iteration_count=2, + disable_gravity=False, + kinematic_enabled=False, + ), + mass_props=sim_utils.MassPropertiesCfg(mass=0.02), + ), + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, 0.0), rot=(1.0, 0.0, 0.0, 0.0)), + ) + + +INSERTIVE_OBJECT_VARIANTS = { + "fbleg": make_insertive_object(f"{UWLAB_CLOUD_ASSETS_DIR}/Props/FurnitureBench/SquareLeg/square_leg.usd"), + "fbdrawerbottom": make_insertive_object( + f"{UWLAB_CLOUD_ASSETS_DIR}/Props/FurnitureBench/DrawerBottom/drawer_bottom.usd" + ), + "peg": make_insertive_object(f"{UWLAB_CLOUD_ASSETS_DIR}/Props/Custom/Peg/peg.usd"), + "cupcake": make_insertive_object(f"{UWLAB_CLOUD_ASSETS_DIR}/Props/Custom/CupCake/cupcake.usd"), + "cube": make_insertive_object(f"{UWLAB_CLOUD_ASSETS_DIR}/Props/Custom/InsertiveCube/insertive_cube.usd"), + "rectangle": make_insertive_object(f"{UWLAB_CLOUD_ASSETS_DIR}/Props/Custom/Rectangle/rectangle.usd"), +} + +variants = {"scene.insertive_object": INSERTIVE_OBJECT_VARIANTS} + + +@configclass +class PickSceneCfg(RlStateSceneCfg): + """OmniReset state scene without the receptive object; cube by default.""" + + receptive_object = None + insertive_object = make_insertive_object(f"{UWLAB_CLOUD_ASSETS_DIR}/Props/Custom/InsertiveCube/insertive_cube.usd") + + +## +# Events +## + +RESET_SUCCESS_EXPR = "env.reward_manager.get_term_cfg('progress_context').func.success" + +_TRAIN_RESET_TERM = dict( + func=task_mdp.SingleObjectMultiResetManager, + mode="reset", + params={ + "dataset_dir": ASTEROID_DATASETS_DIR, + "reset_types": ["ObjectAnywhereEEAnywhere", "ObjectRestingEEGrasped", "ObjectAnywhereEEGrasped"], + "probs": [0.34, 0.33, 0.33], + "success": RESET_SUCCESS_EXPR, + }, +) + +_EVAL_RESET_TERM = dict( + func=task_mdp.SingleObjectMultiResetManager, + mode="reset", + params={ + "dataset_dir": ASTEROID_DATASETS_DIR, + "reset_types": ["ObjectAnywhereEEAnywhere"], + "probs": [1.0], + "success": RESET_SUCCESS_EXPR, + }, +) + +_INSERTIVE_OBJECT_MASS_TERM = dict( + func=task_mdp.randomize_rigid_body_mass, + mode="startup", + params={ + "asset_cfg": SceneEntityCfg("insertive_object"), + # cube-sized objects: 20g - 100g + "mass_distribution_params": (0.02, 0.1), + "operation": "abs", + "distribution": "uniform", + "recompute_inertia": True, + }, +) + +_UNIFIED_DR_TERM = dict( + func=task_mdp.randomize_env_cfg_unified, + mode="reset", + params={ + "asset_cfg": SceneEntityCfg("robot"), + "joint_names": [ + "shoulder_pan_joint", + "shoulder_lift_joint", + "elbow_joint", + "wrist_1_joint", + "wrist_2_joint", + "wrist_3_joint", + ], + "actuator_name": "arm", + "action_name": "arm", + "arm_scale_range": (0.8, 1.2), + "delay_range": (0, 1), + "kp_scale_range": (0.8, 1.2), + "terminal_kp": (1000.0, 1000.0, 1000.0, 50.0, 50.0, 50.0), + "terminal_damping_ratio": (1.0, 1.0, 1.0, 1.0, 1.0, 1.0), + "initial_scales": (0.02, 0.02, 0.02, 0.02, 0.02, 0.2), + "target_scales": (0.01, 0.01, 0.002, 0.02, 0.02, 0.2), + "coupled_progress_range": (0.0, 1.5), + "action_scale_progress_range": (0.0, 1.5), + }, +) + + +@configclass +class PickBaseEventCfg(BaseEventCfg): + """OmniReset base events minus the receptive object; lighter insertive object.""" + + receptive_object_material = None + randomize_receptive_object_mass = None + randomize_insertive_object_mass = EventTerm(**_INSERTIVE_OBJECT_MASS_TERM) + + +@configclass +class PickTrainEventCfg(PickBaseEventCfg): + """Training events: 3-path resets + coupled sysid / gain / action-scale randomization.""" + + reset_from_reset_states = EventTerm(**_TRAIN_RESET_TERM) + randomize_env_cfg_unified = EventTerm(**_UNIFIED_DR_TERM) + + +@configclass +class PickTrainEvalEventCfg(PickBaseEventCfg): + """Eval after Stage 1: no sysid / OSC gain randomization, 1-path resets.""" + + reset_from_reset_states = EventTerm(**_EVAL_RESET_TERM) + + +@configclass +class PickFinetuneEventCfg(FinetuneEventCfg): + """Finetune events: OmniReset's curriculum-ramped sysid + OSC gains, 3-path resets, unified DR.""" + + receptive_object_material = None + randomize_receptive_object_mass = None + randomize_insertive_object_mass = EventTerm(**_INSERTIVE_OBJECT_MASS_TERM) + reset_from_reset_states = EventTerm(**_TRAIN_RESET_TERM) + randomize_env_cfg_unified = EventTerm(**_UNIFIED_DR_TERM) + + +@configclass +class PickFinetuneEvalEventCfg(FinetuneEvalEventCfg): + """Eval after Stage 2 / data collection: fixed sysid + OSC gains, 1-path resets.""" + + receptive_object_material = None + randomize_receptive_object_mass = None + randomize_insertive_object_mass = EventTerm(**_INSERTIVE_OBJECT_MASS_TERM) + reset_from_reset_states = EventTerm(**_EVAL_RESET_TERM) + + +## +# Commands / observations / rewards +## + + +@configclass +class PickCommandsCfg: + """Command specifications for the MDP.""" + + task_command = task_mdp.PickTaskCommandCfg( + asset_cfg=SceneEntityCfg("robot", body_names="body"), + resampling_time_range=(1e6, 1e6), + insertive_asset_cfg=SceneEntityCfg("insertive_object"), + ) + + +@configclass +class PickObservationsCfg(ObservationsCfg): + """OmniReset observations minus every receptive-object term.""" + + @configclass + class PolicyCfg(ObservationsCfg.PolicyCfg): + receptive_asset_pose = None + insertive_asset_in_receptive_asset_frame = None + + @configclass + class CriticCfg(ObservationsCfg.CriticCfg): + receptive_asset_pose = None + insertive_asset_in_receptive_asset_frame = None + receptive_object_material_properties = None + receptive_object_mass = None + + policy: PolicyCfg = PolicyCfg() + critic: CriticCfg = CriticCfg() + + +@configclass +class PickRewardsCfg(RewardsCfg): + """Pick-height success in place of assembly alignment; softer action / joint-velocity penalties.""" + + action_rate = RewTerm(func=task_mdp.action_rate_l2_clamped, weight=-1e-4) + + joint_vel = RewTerm( + func=task_mdp.joint_vel_l2_clamped, + weight=-1e-3, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=["shoulder.*", "elbow.*", "wrist.*"])}, + ) + + progress_context = RewTerm( + func=task_mdp.ProgressContextPickOnly, # type: ignore + weight=0.1, + params={ + "insertive_asset_cfg": SceneEntityCfg("insertive_object"), + "robot_asset_cfg": SceneEntityCfg("robot", body_names="robotiq_base_link"), + # success additionally requires the gripper pointing vertically down (approach axis + # aligned with world -z within ~25 deg). Raise toward 1.0 to demand a stricter vertical. + "gripper_down_dot_threshold": 0.9, + }, + ) + + dense_success_reward = RewTerm(func=task_mdp.dense_success_reward_pick_only, weight=0.1, params={"std": 1.0}) + + success_reward = RewTerm(func=task_mdp.success_reward_pick_only, weight=1.0) + + +## +# Environments +## + + +@configclass +class Ur5eRobotiq2f85PickStateCfg(Ur5eRobotiq2f85RlStateCfg): + """Base pick-only state environment (events set by the Train / Finetune / Eval subclasses).""" + + scene: PickSceneCfg = PickSceneCfg(num_envs=32, env_spacing=1.5) + observations: PickObservationsCfg = PickObservationsCfg() + rewards: PickRewardsCfg = PickRewardsCfg() + commands: PickCommandsCfg = PickCommandsCfg() + variants = variants + + def __post_init__(self): + super().__post_init__() + + # Render settings: plain rasterization is enough for state-based training and keeps the + # tactile data-collection cameras cheap. + self.sim.render.enable_dlssg = False + self.sim.render.enable_ambient_occlusion = False + self.sim.render.enable_reflections = False + self.sim.render.enable_dl_denoiser = False + self.sim.render.antialiasing_mode = "DLAA" + + +# Training configuration (Stage 1: implicit actuator, coupled DR, no curriculum) +@configclass +class Ur5eRobotiq2f85PickRelCartesianOSCTrainCfg(Ur5eRobotiq2f85PickStateCfg): + events: PickTrainEventCfg = PickTrainEventCfg() + actions: Ur5eRobotiq2f85RelativeOSCAction = Ur5eRobotiq2f85RelativeOSCAction() + + +# Finetune configuration (Stage 2: explicit actuator, curriculum ramps sysid + gains + scales) +@configclass +class Ur5eRobotiq2f85PickRelCartesianOSCFinetuneCfg(Ur5eRobotiq2f85PickStateCfg): + """Finetune config: loads converged Stage 1 policy, explicit actuator from start, curriculum ramps DR.""" + + events: PickFinetuneEventCfg = PickFinetuneEventCfg() + actions: Ur5eRobotiq2f85RelativeOSCAction = Ur5eRobotiq2f85RelativeOSCAction() + curriculum: FinetuneCurriculumsCfg = FinetuneCurriculumsCfg() + + def __post_init__(self): + super().__post_init__() + self.scene.robot = EXPLICIT_UR5E_ROBOTIQ_2F85.replace(prim_path="{ENV_REGEX_NS}/Robot") + + +# Evaluation configuration (after Stage 1: implicit actuator, soft gains, no sysid DR) +@configclass +class Ur5eRobotiq2f85PickRelCartesianOSCEvalCfg(Ur5eRobotiq2f85PickStateCfg): + """Eval after Stage 1: implicit actuator, soft gains, large action scale, no sysid DR.""" + + events: PickTrainEvalEventCfg = PickTrainEvalEventCfg() + actions: Ur5eRobotiq2f85RelativeOSCAction = Ur5eRobotiq2f85RelativeOSCAction() + + +# Evaluation configuration (after Stage 2: explicit actuator, stiff gains, fixed sysid) +@configclass +class Ur5eRobotiq2f85PickRelCartesianOSCFinetuneEvalCfg(Ur5eRobotiq2f85PickStateCfg): + """Eval after Stage 2: explicit actuator, stiff gains, small action scale, fixed sysid + OSC gains.""" + + events: PickFinetuneEvalEventCfg = PickFinetuneEvalEventCfg() + actions: Ur5eRobotiq2f85RelativeOSCEvalAction = Ur5eRobotiq2f85RelativeOSCEvalAction() + + def __post_init__(self): + super().__post_init__() + self.scene.robot = EXPLICIT_UR5E_ROBOTIQ_2F85.replace(prim_path="{ENV_REGEX_NS}/Robot") diff --git a/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/__init__.py b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/__init__.py new file mode 100644 index 00000000..92e540ca --- /dev/null +++ b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/__init__.py @@ -0,0 +1,17 @@ +# Copyright (c) 2024-2026, The UW Lab Project Developers. (https://github.com/uw-lab/UWLab/blob/main/CONTRIBUTORS.md). +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""MDP terms for ASTEROID environments. + +Re-exports everything from the OmniReset MDP and adds the pick-specific terms. +""" + +from uwlab_tasks.manager_based.manipulation.omnireset.mdp import * # noqa: F401, F403 + +from .commands_cfg import * # noqa: F401, F403 +from .events import * # noqa: F401, F403 +from .observations import * # noqa: F401, F403 +from .recorders import * # noqa: F401, F403 +from .rewards import * # noqa: F401, F403 diff --git a/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/actions/__init__.py b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/actions/__init__.py new file mode 100644 index 00000000..9f7e639a --- /dev/null +++ b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/actions/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2024-2026, The UW Lab Project Developers. (https://github.com/uw-lab/UWLab/blob/main/CONTRIBUTORS.md). +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from .actions_cfg import * # noqa: F401, F403 +from .task_space_actions import * # noqa: F401, F403 diff --git a/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/actions/actions_cfg.py b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/actions/actions_cfg.py new file mode 100644 index 00000000..58820a9d --- /dev/null +++ b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/actions/actions_cfg.py @@ -0,0 +1,28 @@ +# Copyright (c) 2024-2026, The UW Lab Project Developers. (https://github.com/uw-lab/UWLab/blob/main/CONTRIBUTORS.md). +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +from isaaclab.managers.action_manager import ActionTerm +from isaaclab.utils import configclass + +from uwlab_tasks.manager_based.manipulation.omnireset.mdp.actions.actions_cfg import RelCartesianOSCActionCfg + +from . import task_space_actions + + +@configclass +class RelCartesianOSCPositionActionCfg(RelCartesianOSCActionCfg): + """Position-only Relative Cartesian OSC action. + + Identical to :class:`RelCartesianOSCActionCfg` except the policy controls + only the 3-DOF Cartesian position ``[x, y, z]``; the gripper orientation is + left uncommanded and the desired orientation tracks the current orientation, + so collision-induced rotations are accepted rather than resisted. The full + 6-tuple gain / scale / torque fields are retained (the rotation gains only + provide light damping within a policy step). + """ + + class_type: type[ActionTerm] = task_space_actions.RelCartesianOSCPositionAction diff --git a/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/actions/task_space_actions.py b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/actions/task_space_actions.py new file mode 100644 index 00000000..c13c9386 --- /dev/null +++ b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/actions/task_space_actions.py @@ -0,0 +1,63 @@ +# Copyright (c) 2024-2026, The UW Lab Project Developers. (https://github.com/uw-lab/UWLab/blob/main/CONTRIBUTORS.md). +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +from uwlab_tasks.manager_based.manipulation.omnireset.mdp.actions.task_space_actions import RelCartesianOSCAction + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + + from . import actions_cfg + + +class RelCartesianOSCPositionAction(RelCartesianOSCAction): + """Position-only variant of :class:`RelCartesianOSCAction`. + + The policy outputs a 3-DOF Cartesian delta ``[x, y, z]`` only -- it cannot + rotate the gripper. The desired EE orientation simply tracks the *current* + orientation each policy step, so the controller never accumulates a rotation + error to fight: collision-induced rotations are accepted rather than resisted. + + All control machinery (analytical Jacobian, PD torques, clamping) is + inherited unchanged from the parent; only the action interface and the + desired-pose computation differ. + """ + + cfg: actions_cfg.RelCartesianOSCPositionActionCfg + """The configuration of the action term.""" + + def __init__(self, cfg: actions_cfg.RelCartesianOSCPositionActionCfg, env: ManagerBasedEnv): + super().__init__(cfg, env) + # Override the 6-DOF action buffers with 3-DOF (x, y, z) ones. + self._raw_actions = torch.zeros(self.num_envs, 3, device=self.device) + self._processed_actions = torch.zeros(self.num_envs, 3, device=self.device) + + @property + def action_dim(self) -> int: + return 3 + + def process_actions(self, actions: torch.Tensor): + """Scale raw 3-DOF xyz deltas and compute the desired EE position. + + The desired orientation tracks the current EE orientation, so no rotation + error builds up and collision-induced rotations are accepted, not fought. + """ + self._raw_actions[:] = actions + # ``_scale`` is ``(6,)`` by default and ``(num_envs, 6)`` once domain randomization + # (:class:`~..events.randomize_env_cfg_unified`) has written per-env scales into it. + scaled = actions * self._scale[..., :3] + if self._input_clip is not None: + scaled = torch.clamp(scaled, min=self._input_clip[0], max=self._input_clip[1]) + self._processed_actions[:] = scaled + + # Current EE pose in root (base_link) frame. + ee_pos_b, ee_quat_b = self._get_ee_pose_root_frame() + # Position: desired = current + delta. Orientation: track current (no command). + self._ee_pos_des[:] = ee_pos_b + scaled + self._ee_quat_des[:] = ee_quat_b diff --git a/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/commands.py b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/commands.py new file mode 100644 index 00000000..7dad6d0f --- /dev/null +++ b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/commands.py @@ -0,0 +1,83 @@ +# Copyright (c) 2024-2026, The UW Lab Project Developers. (https://github.com/uw-lab/UWLab/blob/main/CONTRIBUTORS.md). +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Command terms for pick-only tasks.""" + +from __future__ import annotations + +import torch +from collections.abc import Sequence +from typing import TYPE_CHECKING + +from isaaclab.assets import Articulation, RigidObject + +from uwlab_tasks.manager_based.manipulation.omnireset.mdp.commands import TaskDependentCommand + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + + from .commands_cfg import PickTaskCommandCfg + + +class PickTaskCommand(TaskDependentCommand): + """Task command for pick-only tasks (a single insertive object, no receptive object). + + Counterpart of :class:`~uwlab_tasks.manager_based.manipulation.omnireset.mdp.commands.TaskCommand` + for scenes without a receptive object. The command itself is a zero vector (the + policy is not goal-conditioned); the term exists to drive the task-dependent reset + events and to log pick metrics. + """ + + cfg: PickTaskCommandCfg + """Configuration for the command generator.""" + + def __init__(self, cfg: PickTaskCommandCfg, env: ManagerBasedEnv): + super().__init__(cfg, env) + + self.insertive_asset: Articulation | RigidObject = env.scene[cfg.insertive_asset_cfg.name] + self._success_expr = cfg.success + + self.metrics["average_object_height"] = torch.zeros(self.num_envs, device=self.device) + self.metrics["end_of_episode_object_height"] = torch.zeros(self.num_envs, device=self.device) + self.metrics["end_of_episode_success_rate"] = torch.zeros(self.num_envs, device=self.device) + + self.object_height = torch.zeros(self.num_envs, device=self.device) + self.success = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + + """ + Properties + """ + + @property + def command(self) -> torch.Tensor: + return torch.zeros(self.num_envs, 3, device=self.device) + + """ + Implementation specific functions. + """ + + def _update_metrics(self): + # logs end of episode data + reset_env = self._env.episode_length_buf == 0 + self.metrics["end_of_episode_object_height"][reset_env] = self.object_height[reset_env] + self.metrics["end_of_episode_success_rate"][reset_env] = self.success[reset_env].float() + + # logs current data + self.object_height[:] = self.insertive_asset.data.root_pos_w[:, 2] - self._env.scene.env_origins[:, 2] + if self._success_expr is not None: + self.success[:] = eval(self._success_expr, {"env": self._env}) + self.metrics["average_object_height"][:] = self.object_height + + def _resample_command(self, env_ids: Sequence[int]): + super()._resample_command(env_ids) + + def _update_command(self): + super()._update_command() + + def _set_debug_vis_impl(self, debug_vis: bool): + pass + + def _debug_vis_callback(self, event): + pass diff --git a/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/commands_cfg.py b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/commands_cfg.py new file mode 100644 index 00000000..0f29a040 --- /dev/null +++ b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/commands_cfg.py @@ -0,0 +1,28 @@ +# Copyright (c) 2024-2026, The UW Lab Project Developers. (https://github.com/uw-lab/UWLab/blob/main/CONTRIBUTORS.md). +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from dataclasses import MISSING + +from isaaclab.managers import SceneEntityCfg +from isaaclab.utils import configclass + +from uwlab_tasks.manager_based.manipulation.omnireset.mdp.commands_cfg import TaskDependentCommandCfg + +from .commands import PickTaskCommand + + +@configclass +class PickTaskCommandCfg(TaskDependentCommandCfg): + """Configuration for :class:`~.commands.PickTaskCommand`.""" + + class_type: type = PickTaskCommand + + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot") + + insertive_asset_cfg: SceneEntityCfg = MISSING + + success: str | None = "env.reward_manager.get_term_cfg('progress_context').func.success" + """Expression (evaluated with ``env`` bound) yielding a per-env success mask used for the + end-of-episode success-rate metric. Set to ``None`` to disable.""" diff --git a/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/events.py b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/events.py new file mode 100644 index 00000000..96fe10e3 --- /dev/null +++ b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/events.py @@ -0,0 +1,348 @@ +# Copyright (c) 2024-2026, The UW Lab Project Developers. (https://github.com/uw-lab/UWLab/blob/main/CONTRIBUTORS.md). +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Event terms for pick-only tasks and ASTEROID domain randomization.""" + +from __future__ import annotations + +import os +import torch +from typing import TYPE_CHECKING + +import isaaclab.utils.math as math_utils +from isaaclab.assets import Articulation, RigidObject +from isaaclab.managers import EventTermCfg, ManagerTermBase, SceneEntityCfg + +from uwlab_tasks.manager_based.manipulation.omnireset.mdp import utils +from uwlab_tasks.manager_based.manipulation.omnireset.mdp.actions.task_space_actions import RelCartesianOSCAction +from uwlab_tasks.manager_based.manipulation.omnireset.mdp.events import MultiResetManager, sample_state_data_set +from uwlab_tasks.manager_based.manipulation.omnireset.mdp.success_monitor_cfg import SuccessMonitorCfg + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + + +class SingleObjectMultiResetManager(MultiResetManager): + """:class:`~uwlab_tasks.manager_based.manipulation.omnireset.mdp.events.MultiResetManager` + for scenes with only an insertive object. + + OmniReset keys reset datasets by the object *pair* directory (``Peg__PegHole``); pick-only + scenes have no receptive object, so datasets are keyed by the insertive object alone + (``/Resets//resets_.pt``). Sampling and state restoration + are inherited unchanged. + """ + + def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): + # NOTE: deliberately skips MultiResetManager.__init__, which requires a receptive object. + ManagerTermBase.__init__(self, cfg, env) + + dataset_dir: str = cfg.params.get("dataset_dir", "") + reset_types: list[str] = cfg.params.get("reset_types", []) + probabilities: list[float] = cfg.params.get("probs", []) + + if not reset_types: + raise ValueError("No reset_types provided") + if len(reset_types) != len(probabilities): + raise ValueError("Number of reset_types must match number of probabilities") + + insertive_usd_path = env.scene["insertive_object"].cfg.spawn.usd_path + pair = utils.object_name_from_usd(insertive_usd_path) + + dataset_files = [f"{dataset_dir}/Resets/{pair}/resets_{rt}.pt" for rt in reset_types] + + self.datasets = [] + num_states = [] + for dataset_file in dataset_files: + local_file_path = utils.safe_retrieve_file_path(dataset_file) + if not os.path.exists(local_file_path): + raise FileNotFoundError(f"Dataset file {dataset_file} could not be accessed or downloaded.") + + dataset = torch.load(local_file_path) + num_states.append(len(dataset["initial_state"]["articulation"]["robot"]["joint_position"])) + init_indices = torch.arange(num_states[-1], device=env.device) + self.datasets.append(sample_state_data_set(dataset, init_indices, env.device)) + + self.probs = torch.tensor(probabilities, device=env.device) / sum(probabilities) + self.num_states = torch.tensor(num_states, device=env.device) + self.num_tasks = len(self.datasets) + + if cfg.params.get("success") is not None: + success_monitor_cfg = SuccessMonitorCfg( + monitored_history_len=100, num_monitored_data=self.num_tasks, device=env.device + ) + self.success_monitor = success_monitor_cfg.class_type(success_monitor_cfg) + + self.task_id = torch.randint(0, self.num_tasks, (self.num_envs,), device=self.device) + + +class reset_root_states_discrete_grid(ManagerTermBase): + """Reset root states by sampling x/y around discrete grid centers. + + X/Y centers are generated from the configured pose range using ``grid_shape``. Each reset + samples one center per env, applies small uniform x/y jitter, then samples the remaining + pose dimensions uniformly from ``pose_range``. + """ + + def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): + super().__init__(cfg, env) + + pose_range_dict = cfg.params.get("pose_range") + velocity_range_dict = cfg.params.get("velocity_range") + + self.pose_range = torch.tensor( + [pose_range_dict.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]], + device=env.device, + ) + self.velocity_range = torch.tensor( + [velocity_range_dict.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]], + device=env.device, + ) + self.asset_cfgs = list(cfg.params.get("asset_cfgs", dict()).values()) + self.offset_asset_cfg = cfg.params.get("offset_asset_cfg") + self.use_bottom_offset = cfg.params.get("use_bottom_offset", False) + + grid_shape = cfg.params.get("grid_shape", (3, 3)) + if len(grid_shape) != 2: + raise ValueError("grid_shape must be a 2-tuple, e.g. (3, 3)") + num_x, num_y = grid_shape + x_centers = torch.linspace(self.pose_range[0, 0], self.pose_range[0, 1], num_x, device=env.device) + y_centers = torch.linspace(self.pose_range[1, 0], self.pose_range[1, 1], num_y, device=env.device) + grid_x, grid_y = torch.meshgrid(x_centers, y_centers, indexing="ij") + self.xy_grid = torch.stack([grid_x.reshape(-1), grid_y.reshape(-1)], dim=-1) + + self.xy_noise_range = torch.tensor(cfg.params.get("xy_noise_range", (-0.01, 0.01)), device=env.device) + + if self.use_bottom_offset: + self.bottom_offset_positions = dict() + for asset_cfg in self.asset_cfgs: + asset: RigidObject | Articulation = env.scene[asset_cfg.name] + metadata = utils.read_metadata_from_usd_directory(asset.cfg.spawn.usd_path) + bottom_offset = metadata.get("bottom_offset") + self.bottom_offset_positions[asset_cfg.name] = ( + torch.tensor(bottom_offset.get("pos"), device=env.device).unsqueeze(0).repeat(env.num_envs, 1) + ) + assert tuple(bottom_offset.get("quat")) == (1.0, 0.0, 0.0, 0.0), ( + "Bottom offset rotation must be (1.0, 0.0, 0.0, 0.0)" + ) + + def __call__( + self, + env: ManagerBasedEnv, + env_ids: torch.Tensor, + pose_range: dict[str, tuple[float, float]], + velocity_range: dict[str, tuple[float, float]], + asset_cfgs: dict[str, SceneEntityCfg] = dict(), + offset_asset_cfg: SceneEntityCfg = None, + use_bottom_offset: bool = False, + grid_shape: tuple[int, int] = (3, 3), + xy_noise_range: tuple[float, float] = (-0.01, 0.01), + ) -> None: + if env_ids is None: + env_ids = torch.arange(env.scene.num_envs, device=env.device) + + num_envs = len(env_ids) + rand_pose_samples = math_utils.sample_uniform( + self.pose_range[:, 0], self.pose_range[:, 1], (num_envs, 6), device=env.device + ) + + grid_ids = torch.randint(0, self.xy_grid.shape[0], (num_envs,), device=env.device) + xy = self.xy_grid[grid_ids] + xy_noise = math_utils.sample_uniform( + self.xy_noise_range[0], self.xy_noise_range[1], (num_envs, 2), device=env.device + ) + rand_pose_samples[:, 0:2] = (xy + xy_noise).clamp(min=self.pose_range[0:2, 0], max=self.pose_range[0:2, 1]) + + orientations_delta = math_utils.quat_from_euler_xyz( + rand_pose_samples[:, 3], rand_pose_samples[:, 4], rand_pose_samples[:, 5] + ) + rand_vel_samples = math_utils.sample_uniform( + self.velocity_range[:, 0], self.velocity_range[:, 1], (num_envs, 6), device=env.device + ) + + for asset_cfg in self.asset_cfgs: + asset: RigidObject | Articulation = env.scene[asset_cfg.name] + root_states = asset.data.default_root_state[env_ids].clone() + positions = root_states[:, 0:3] + env.scene.env_origins[env_ids] + rand_pose_samples[:, 0:3] + + if self.offset_asset_cfg: + offset_asset: RigidObject | Articulation = env.scene[self.offset_asset_cfg.name] + offset_positions = offset_asset.data.default_root_state[env_ids].clone() + positions += offset_positions[:, 0:3] + + if self.use_bottom_offset: + positions -= self.bottom_offset_positions[asset_cfg.name][env_ids, 0:3] + + orientations = math_utils.quat_mul(root_states[:, 3:7], orientations_delta) + velocities = root_states[:, 7:13] + rand_vel_samples + + asset.write_root_pose_to_sim(torch.cat([positions, orientations], dim=-1), env_ids=env_ids) + asset.write_root_velocity_to_sim(velocities, env_ids=env_ids) + + +class randomize_env_cfg_unified(ManagerTermBase): + """Coupled domain randomization over arm joint dynamics, OSC gains and action scaling. + + Samples one ``coupled_progress`` scalar per env and maps it to arm sysid (armature / + friction), actuator delay and :class:`RelCartesianOSCAction` Kp/Kd, so the environment + always lies in a tractable region (high joint friction with soft gains is unsolvable). + Action scaling is unrelated to task feasibility and is randomized independently. + + * arm sysid range: ``[0, full sysid-randomized values]`` + * delay range: ``[delay_min, delay_max]`` + * OSC controller range: ``[action cfg defaults, terminal_kp / terminal_damping_ratio]`` + * action scaling range: ``[initial_scales, target_scales]`` + """ + + def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): + super().__init__(cfg, env) + self.asset_cfg: SceneEntityCfg = cfg.params["asset_cfg"] + self.robot: Articulation = env.scene[self.asset_cfg.name] + self.joint_ids = self.robot.find_joints(cfg.params["joint_names"])[0] + self.actuator_name: str = cfg.params["actuator_name"] + self._action_name: str = cfg.params["action_name"] + self._action_term: RelCartesianOSCAction | None = None + + metadata = utils.read_metadata_from_usd_directory(self.robot.cfg.spawn.usd_path) + sysid = metadata["sysid"] + self.armature = sysid["armature"] + self.static_friction = sysid["static_friction"] + self.dynamic_ratio = sysid["dynamic_ratio"] + self.viscous_friction = sysid["viscous_friction"] + + def _resolve_action_term(self): + if self._action_term is not None: + return + action_term = self._env.action_manager._terms.get(self._action_name) + if action_term is None or not isinstance(action_term, RelCartesianOSCAction): + raise ValueError(f"Action term '{self._action_name}' is not a RelCartesianOSCAction.") + # The action term stores a single (6,) scale; promote it to per-env (num_envs, 6) so we can + # randomize it per env. ``process_actions`` broadcasts either layout. + if action_term._scale.dim() == 1: + action_term._scale = action_term._scale.unsqueeze(0).expand(self.num_envs, -1).clone() + self._action_term = action_term + + def __call__( + self, + env: ManagerBasedEnv, + env_ids: torch.Tensor, + asset_cfg: SceneEntityCfg, + joint_names: list[str], + actuator_name: str, + action_name: str, + arm_scale_range: tuple[float, float] = (0.8, 1.2), + delay_range: tuple[int, int] = (0, 1), + kp_scale_range: tuple[float, float] = (0.8, 1.2), + terminal_kp: tuple[float, ...] = (1000.0, 1000.0, 1000.0, 50.0, 50.0, 50.0), + terminal_damping_ratio: tuple[float, ...] = (1.0, 1.0, 1.0, 1.0, 1.0, 1.0), + initial_scales: tuple[float, ...] = (0.02, 0.02, 0.02, 0.02, 0.02, 0.2), + target_scales: tuple[float, ...] = (0.01, 0.01, 0.002, 0.02, 0.02, 0.2), + coupled_progress_range: tuple[float, float] = (0.0, 1.0), + action_scale_progress_range: tuple[float, float] = (0.0, 1.0), + ) -> None: + self._resolve_action_term() + + if env_ids is None: + env_ids = torch.arange(env.scene.num_envs, device=self.robot.device) + + n = len(env_ids) + n_joints = len(self.joint_ids) + device = self.robot.device + + c_lo, c_hi = coupled_progress_range + coupled_progress = c_lo + torch.rand(n, 1, device=device) * (c_hi - c_lo) + + def _scale_sysid(nominal, scale_range): + lo, hi = scale_range + val = torch.as_tensor(nominal, device=device, dtype=torch.float32) + return val * (lo + torch.rand(n, n_joints, device=device) * (hi - lo)) + + arm_vals = _scale_sysid(self.armature, arm_scale_range) * coupled_progress + sfric_vals = _scale_sysid(self.static_friction, arm_scale_range) * coupled_progress + dratio_vals = _scale_sysid(self.dynamic_ratio, arm_scale_range) * coupled_progress + dfric_vals = torch.minimum(dratio_vals * sfric_vals, sfric_vals) + vfric_vals = _scale_sysid(self.viscous_friction, arm_scale_range) * coupled_progress + + self.robot.write_joint_armature_to_sim(arm_vals, joint_ids=self.joint_ids, env_ids=env_ids) + self.robot.write_joint_friction_coefficient_to_sim( + sfric_vals, + joint_dynamic_friction_coeff=dfric_vals, + joint_viscous_friction_coeff=vfric_vals, + joint_ids=self.joint_ids, + env_ids=env_ids, + ) + + delay_lo, delay_hi = delay_range + if delay_hi > delay_lo: + actuator = self.robot.actuators[self.actuator_name] + if hasattr(actuator, "positions_delay_buffer"): + max_delay = delay_lo + torch.round(coupled_progress.squeeze(-1) * float(delay_hi - delay_lo)).to( + dtype=torch.int32 + ) + min_delay = torch.full_like(max_delay, fill_value=delay_lo) + span = (max_delay - min_delay + 1).clamp(min=1) + # Vectorized integer sampling with per-env bounds in [min_delay, max_delay]. + delays = min_delay + torch.floor(torch.rand(n, device=device) * span.to(torch.float32)).to(torch.int32) + actuator.positions_delay_buffer.set_time_lag(delays, env_ids) + actuator.velocities_delay_buffer.set_time_lag(delays, env_ids) + actuator.efforts_delay_buffer.set_time_lag(delays, env_ids) + + k_lo, k_hi = kp_scale_range + s_xyz = k_lo + torch.rand(n, 1, device=device) * (k_hi - k_lo) + s_rpy = k_lo + torch.rand(n, 1, device=device) * (k_hi - k_lo) + s_dr_xyz = k_lo + torch.rand(n, 1, device=device) * (k_hi - k_lo) + s_dr_rpy = k_lo + torch.rand(n, 1, device=device) * (k_hi - k_lo) + + kp_default = self._action_term._kp_default + dr_default = self._action_term._damping_ratio_default + kp_term = torch.tensor(terminal_kp, device=device, dtype=torch.float32).unsqueeze(0).repeat(n, 1) + dr_term = torch.tensor(terminal_damping_ratio, device=device, dtype=torch.float32).unsqueeze(0).repeat(n, 1) + kp_term[:, :3] *= s_xyz + kp_term[:, 3:] *= s_rpy + dr_term[:, :3] *= s_dr_xyz + dr_term[:, 3:] *= s_dr_rpy + + new_kp = kp_default.unsqueeze(0) + coupled_progress * (kp_term - kp_default.unsqueeze(0)) + new_dr = dr_default.unsqueeze(0) + coupled_progress * (dr_term - dr_default.unsqueeze(0)) + self._action_term._kp[env_ids] = new_kp + self._action_term._kd[env_ids] = 2.0 * torch.sqrt(new_kp) * new_dr + + a_lo, a_hi = action_scale_progress_range + action_progress = a_lo + torch.rand(n, 1, device=device) * (a_hi - a_lo) + initial = torch.tensor(initial_scales, device=device, dtype=torch.float32).unsqueeze(0) + target = torch.tensor(target_scales, device=device, dtype=torch.float32).unsqueeze(0) + self._action_term._scale[env_ids] = initial + action_progress * (target - initial) + + +class randomize_gripper_pos_affine(ManagerTermBase): + """Per-env affine noise on the ``gripper_pos`` observation: ``pos -> pos * scale + offset``. + + Stores ``(num_envs,)`` ``scale`` and ``offset`` tensors that + :func:`~.observations.gripper_pos_normalized` reads via + ``env.event_manager.get_term_cfg().func``. Both are resampled per env at each + reset and fixed within an episode. + + Makes the policy robust to calibration / encoder drift on the real Robotiq's POS register, + so it cannot latch onto an absolute threshold like ``pos > 0.93`` to detect a grasp. + """ + + def __init__(self, cfg: EventTermCfg, env: ManagerBasedEnv): + super().__init__(cfg, env) + self.scale = torch.ones(env.num_envs, device=env.device) + self.offset = torch.zeros(env.num_envs, device=env.device) + + def __call__( + self, + env: ManagerBasedEnv, + env_ids, + scale_range: tuple[float, float] = (0.75, 1.25), + offset_range: tuple[float, float] = (-0.1, 0.1), + ) -> None: + if env_ids is None: + env_ids = torch.arange(env.num_envs, device=env.device) + n = env_ids.shape[0] + s_lo, s_hi = scale_range + o_lo, o_hi = offset_range + self.scale[env_ids] = s_lo + torch.rand(n, device=env.device) * (s_hi - s_lo) + self.offset[env_ids] = o_lo + torch.rand(n, device=env.device) * (o_hi - o_lo) diff --git a/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/observations.py b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/observations.py new file mode 100644 index 00000000..ba73329b --- /dev/null +++ b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/observations.py @@ -0,0 +1,108 @@ +# Copyright (c) 2024-2026, The UW Lab Project Developers. (https://github.com/uw-lab/UWLab/blob/main/CONTRIBUTORS.md). +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Observation terms for the ASTEROID tactile / proprioceptive student policies.""" + +from __future__ import annotations + +import math +import torch +from typing import TYPE_CHECKING + +import isaaclab.utils.math as math_utils +from isaaclab.assets import Articulation +from isaaclab.managers import SceneEntityCfg + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv, ManagerBasedRLEnv + + +def gripper_pos_normalized( + env: ManagerBasedEnv, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot", joint_names=["left_inner_finger_knuckle_joint"]), + full_close_angle: float = -math.pi / 4, + scale_event_name: str | None = None, + jitter_std: float = 0.0, +) -> torch.Tensor: + """Robotiq 2F-85 gripper position as a ``[0, 1]`` scalar (0 = open, 1 = closed). + + Mirrors the real-world Robotiq POS register (``robotiq_gripper.get_current_position``), + normalized to ``[0, 1]`` instead of the firmware's 0-255. + + Inverts the sim-side mapping ``inner_finger_knuckle_joint_angle = full_close_angle * pos``, + where ``full_close_angle`` is the joint angle at fully-closed (negative on the UWLab URDF: + ~-pi/4). For multi-joint ``asset_cfg`` the mean angle is used; the 2F-85's left/right + knuckles are mimic-coupled so they track each other. + + Args: + env: The environment. + asset_cfg: Scene entity + joint(s) to read. Defaults to the left inner_finger_knuckle joint. + full_close_angle: Joint angle at fully closed (radians, signed). Default ``-pi/4`` matches + UWLab's 2F-85 URDF convention; pass ``+pi/4`` (or use a different joint) if your robot's + joint axis is the opposite sign. + scale_event_name: Name of a :class:`~.events.randomize_gripper_pos_affine` event term whose + per-env ``scale`` / ``offset`` are applied to the reading (calibration-drift DR). + jitter_std: Std of per-step Gaussian noise added on top (freshly sampled each call). + + Returns: + Tensor of shape ``(num_envs, 1)`` in ``[0.0, 1.0]`` (before scale / offset / jitter). + """ + robot: Articulation = env.scene[asset_cfg.name] + joint_ids = asset_cfg.joint_ids if asset_cfg.joint_ids is not None else slice(None) + angle = robot.data.joint_pos[:, joint_ids] + if angle.dim() > 1 and angle.shape[-1] > 1: + angle = angle.mean(dim=-1, keepdim=True) + elif angle.dim() == 1: + angle = angle.unsqueeze(-1) + pos = (angle / full_close_angle).clamp(0.0, 1.0) + if scale_event_name is not None: + # For class-based terms (ManagerTermBase subclasses) cfg.func is the instantiated object. + try: + scale_term = env.event_manager.get_term_cfg(scale_event_name).func + except ValueError as e: + raise RuntimeError( + f"gripper_pos_normalized: event term '{scale_event_name}' not registered on event_manager." + ) from e + if not hasattr(scale_term, "scale"): + raise RuntimeError( + f"gripper_pos_normalized: event term '{scale_event_name}' has no .scale attr " + f"(got {type(scale_term).__name__})." + ) + # scale / offset: (num_envs,) -> (num_envs, 1) for broadcast. + pos = pos * scale_term.scale.unsqueeze(-1) + if hasattr(scale_term, "offset"): + pos = pos + scale_term.offset.unsqueeze(-1) + if jitter_std > 0.0: + pos = pos + torch.randn_like(pos) * jitter_std + return pos.to(torch.float32) + + +def fingertip_contact_force_b( + env: ManagerBasedRLEnv, + contact_sensor_name: str, + root_asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + root_body_name: str = "robotiq_base_link", +) -> torch.Tensor: + """Contact force from a single fingertip contact sensor, expressed in a body frame. + + Args: + env: The environment to extract contact forces from. + contact_sensor_name: Name of the contact sensor to read from. + root_asset_cfg: Asset whose body frame the force is expressed in. + root_body_name: Body of ``root_asset_cfg`` to use as reference frame. + + Returns: + Contact force in body frame. Shape: ``(num_envs, 3)``. + """ + root_asset: Articulation = env.scene[root_asset_cfg.name] + root_body_idx = root_asset.body_names.index(root_body_name) + root_quat_w = root_asset.data.body_link_quat_w[:, root_body_idx].view(-1, 4) + + contact_sensor = env.scene.sensors[contact_sensor_name] + # force_matrix_w is flattened, so we reshape to (num_envs, 3) + force_w = contact_sensor.data.force_matrix_w.view(env.num_envs, 3) + + # Rotation only: forces are free vectors. + return math_utils.quat_apply_inverse(root_quat_w, force_w) diff --git a/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/recorders/__init__.py b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/recorders/__init__.py new file mode 100644 index 00000000..0a39ea2c --- /dev/null +++ b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/recorders/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) 2024-2026, The UW Lab Project Developers. (https://github.com/uw-lab/UWLab/blob/main/CONTRIBUTORS.md). +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Recorder terms for ASTEROID data collection.""" + +from .recorders import * # noqa: F401, F403 +from .recorders_cfg import * # noqa: F401, F403 diff --git a/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/recorders/recorders.py b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/recorders/recorders.py new file mode 100644 index 00000000..005facc6 --- /dev/null +++ b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/recorders/recorders.py @@ -0,0 +1,35 @@ +# Copyright (c) 2024-2026, The UW Lab Project Developers. (https://github.com/uw-lab/UWLab/blob/main/CONTRIBUTORS.md). +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch + +from isaaclab.managers.recorder_manager import RecorderTerm + + +class PreStepExpertMaskRecorder(RecorderTerm): + """Records, per step, whether the expert (1) or the exploration policy (0) produced the action. + + The mask is pushed in from the data-collection script (see + ``scripts/ASTEROID/collect_demos_asteroid.py``) via :meth:`set_mask` before each step. + """ + + def __init__(self, cfg, env): + super().__init__(cfg, env) + self._expert_mask = torch.ones((env.num_envs, 1), device=env.device) + self._exploration_horizon = None + + def set_mask(self, expert_mask: torch.Tensor): + """Set the expert mask data externally.""" + self._expert_mask = expert_mask + + def set_exploration_horizon(self, exploration_horizon: torch.Tensor): + """Set the exploration horizon data externally.""" + self._exploration_horizon = exploration_horizon + + def record_pre_step(self) -> tuple[str, torch.Tensor]: + """Record the expert mask before each step.""" + return "expert_mask", self._expert_mask.clone() diff --git a/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/recorders/recorders_cfg.py b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/recorders/recorders_cfg.py new file mode 100644 index 00000000..d429274b --- /dev/null +++ b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/recorders/recorders_cfg.py @@ -0,0 +1,27 @@ +# Copyright (c) 2024-2026, The UW Lab Project Developers. (https://github.com/uw-lab/UWLab/blob/main/CONTRIBUTORS.md). +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab.managers.recorder_manager import RecorderTerm, RecorderTermCfg +from isaaclab.utils import configclass + +from uwlab_tasks.manager_based.manipulation.omnireset.mdp.recorders.recorders_cfg import ( + ActionStateRecorderManagerCfg, +) + +from . import recorders + + +@configclass +class PreStepExpertMaskRecorderCfg(RecorderTermCfg): + """Configuration for the expert action mask recorder term (for DAgger-style data).""" + + class_type: type[RecorderTerm] = recorders.PreStepExpertMaskRecorder + + +@configclass +class AsteroidActionStateRecorderManagerCfg(ActionStateRecorderManagerCfg): + """OmniReset's action/state recorder plus the per-step expert mask.""" + + record_pre_step_expert_mask = PreStepExpertMaskRecorderCfg() diff --git a/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/rewards.py b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/rewards.py new file mode 100644 index 00000000..d69bbb74 --- /dev/null +++ b/source/uwlab_tasks/uwlab_tasks/manager_based/manipulation/asteroid/mdp/rewards.py @@ -0,0 +1,112 @@ +# Copyright (c) 2024-2026, The UW Lab Project Developers. (https://github.com/uw-lab/UWLab/blob/main/CONTRIBUTORS.md). +# All Rights Reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Reward terms for pick-only tasks.""" + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +import isaaclab.utils.math as math_utils +from isaaclab.assets import Articulation, RigidObject +from isaaclab.managers import ManagerTermBase, RewardTermCfg, SceneEntityCfg + +from uwlab_tasks.manager_based.manipulation.omnireset.assembly_keypoints import Offset +from uwlab_tasks.manager_based.manipulation.omnireset.mdp import utils +from uwlab_tasks.manager_based.manipulation.omnireset.mdp.success_monitor_cfg import SuccessMonitorCfg + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + + +class ProgressContextPickOnly(ManagerTermBase): + """Pick-only success context (no receptive object). + + Success is ``insertive object lifted above pick_height_threshold`` AND ``gripper pointing + vertically down``. Other reward / termination terms read :attr:`success`, + :attr:`insertive_asset_z` and :attr:`continuous_success_counter` from this term. + """ + + def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + self.insertive_asset: Articulation | RigidObject = env.scene[cfg.params.get("insertive_asset_cfg").name] # type: ignore + + insertive_meta = utils.read_metadata_from_usd_directory(self.insertive_asset.cfg.spawn.usd_path) + self.insertive_asset_offset = Offset( + pos=tuple(insertive_meta.get("assembled_offset").get("pos")), + quat=tuple(insertive_meta.get("assembled_offset").get("quat")), + ) + + # Gripper orientation tracking: success additionally requires the gripper to point + # vertically down. The gripper's approach axis is the local axis given by + # ``gripper_approach_direction`` in the robot metadata (local +x for the Robotiq 2f85); + # "pointing down" means that axis, expressed in world frame, aligns with world -z. + self.robot: Articulation = env.scene[cfg.params.get("robot_asset_cfg").name] # type: ignore + robot_meta = utils.read_metadata_from_usd_directory(self.robot.cfg.spawn.usd_path) + approach_dir = robot_meta.get("gripper_approach_direction", [1.0, 0.0, 0.0]) + self.gripper_approach_dir = torch.tensor(approach_dir, dtype=torch.float32, device=env.device).view(1, 3) + self.gripper_body_id = self.robot.find_bodies(cfg.params.get("robot_asset_cfg").body_names)[0][0] + self.gripper_pointing_down = torch.zeros((env.num_envs), dtype=torch.bool, device=env.device) + + self.insertive_asset_z = torch.zeros((env.num_envs), device=env.device) + self.success = torch.zeros((self._env.num_envs), dtype=torch.bool, device=self._env.device) + self.continuous_success_counter = torch.zeros((self._env.num_envs), dtype=torch.int32, device=self._env.device) + + success_monitor_cfg = SuccessMonitorCfg(monitored_history_len=100, num_monitored_data=1, device=env.device) + self.success_monitor = success_monitor_cfg.class_type(success_monitor_cfg) + + def reset(self, env_ids: torch.Tensor | None = None) -> None: + super().reset(env_ids) + self.continuous_success_counter[:] = 0 + + def __call__( + self, + env: ManagerBasedRLEnv, + insertive_asset_cfg: SceneEntityCfg, + robot_asset_cfg: SceneEntityCfg, + command_context: str = "task_command", + pick_height_threshold: float = 0.02, + gripper_down_dot_threshold: float = 0.9, + ) -> torch.Tensor: + # Object lifted above threshold? + insertive_z_pos = self.insertive_asset.data.root_pos_w[:, 2] + self.insertive_asset_z[:] = insertive_z_pos + + # Gripper pointing vertically down? Rotate the local approach axis into world frame and + # require its z-component to be close to -1. A threshold of -1.0 disables the constraint. + if gripper_down_dot_threshold <= -1.0: + self.gripper_pointing_down[:] = True + else: + gripper_quat_w = self.robot.data.body_link_quat_w[:, self.gripper_body_id] + approach_w = math_utils.quat_apply(gripper_quat_w, self.gripper_approach_dir.expand(env.num_envs, 3)) + self.gripper_pointing_down[:] = (-approach_w[:, 2]) >= gripper_down_dot_threshold + + self.success[:] = (insertive_z_pos > pick_height_threshold) & self.gripper_pointing_down + + self.continuous_success_counter[:] = torch.where( + self.success, self.continuous_success_counter + 1, torch.zeros_like(self.continuous_success_counter) + ) + self.success_monitor.success_update( + torch.zeros(env.num_envs, dtype=torch.int32, device=env.device), self.success + ) + + return torch.zeros(env.num_envs, device=env.device) + + +def dense_success_reward_pick_only( + env: ManagerBasedRLEnv, std: float, context: str = "progress_context" +) -> torch.Tensor: + """Dense shaping toward lifting the object: ``exp(-max(0, 0.4 - z) / std)``.""" + context_term: ManagerTermBase = env.reward_manager.get_term_cfg(context).func # type: ignore + insertive_asset_z: torch.Tensor = getattr(context_term, "insertive_asset_z") + return torch.exp(-torch.clamp(0.4 - insertive_asset_z, min=0) / std) + + +def success_reward_pick_only(env: ManagerBasedRLEnv, context: str = "progress_context") -> torch.Tensor: + """Sparse reward: 1 while :class:`ProgressContextPickOnly` reports success.""" + context_term: ManagerTermBase = env.reward_manager.get_term_cfg(context).func # type: ignore + success: torch.Tensor = getattr(context_term, "success") + return torch.where(success, 1.0, 0.0)