From d6893f0f1822cdc81933b9e6fee8561d32a590ea Mon Sep 17 00:00:00 2001 From: HuiDong Chen Date: Wed, 15 Jul 2026 12:54:18 +0800 Subject: [PATCH 01/14] franka pour from Maximilian --- .../teleop_franka_pour_spacemouse.py | 223 ++ .../generate_franka_pour_reset_dataset.py | 77 + .../validate_franka_pour_reset_dataset.py | 310 ++ .../visualize_franka_pour_reset_dataset.py | 335 +++ .../max-franka-pour-reset-dataset.rst | 5 + .../changelog.d/feat-proxycoupledsolver.rst | 10 + .../max-franka-pour-reset-dataset.rst | 5 + .../changelog.d/mpm-task-support.minor.rst | 19 + .../isaaclab_newton/cloner/__init__.pyi | 2 + .../isaaclab_newton/cloner/replicate.py | 42 + .../isaaclab_newton/physics/mpm_manager.py | 13 +- .../physics/mpm_manager_cfg.py | 27 +- .../isaaclab_newton/physics/newton_manager.py | 16 +- .../physics/test_newton_fabric_body_sync.py | 18 + .../test_newton_manager_abstraction.py | 405 ++- .../changelog.d/pin-newton-c7ae7c7.rst | 6 + .../isaaclab_rl/entrypoints/common.py | 14 +- .../test/test_entrypoints_common.py | 71 + .../changelog.d/franka-pour.minor.rst | 21 + .../contrib/franka_pour/README.md | 61 + .../contrib/franka_pour/__init__.py | 12 + .../franka_pour/_reset_collision_screen.py | 333 +++ .../contrib/franka_pour/config/__init__.py | 4 + .../franka_pour/config/franka/__init__.py | 44 + .../config/franka/agents/__init__.py | 4 + .../config/franka/agents/rsl_rl_ppo_cfg.py | 125 + .../contrib/franka_pour/cube_bowl_mesh.py | 175 ++ .../contrib/franka_pour/cube_bowl_spawner.py | 166 ++ .../franka_pour/cube_bowl_spawner_cfg.py | 52 + .../contrib/franka_pour/cup_media.py | 125 + .../contrib/franka_pour/mdp/__init__.py | 108 + .../contrib/franka_pour/mdp/actions.py | 906 ++++++ .../contrib/franka_pour/mdp/curriculums.py | 181 ++ .../contrib/franka_pour/mdp/events.py | 22 + .../contrib/franka_pour/mdp/observations.py | 206 ++ .../contrib/franka_pour/mdp/reset_dataset.py | 183 ++ .../contrib/franka_pour/mdp/reset_mixture.py | 17 + .../contrib/franka_pour/mdp/rewards.py | 1031 +++++++ .../contrib/franka_pour/mdp/terminations.py | 313 ++ .../contrib/franka_pour/media_fill.py | 103 + .../contrib/franka_pour/pour_env.py | 2610 +++++++++++++++++ .../contrib/franka_pour/pour_env_cfg.py | 2139 ++++++++++++++ .../franka_pour/reset_dataset_generator.py | 1898 ++++++++++++ .../contrib/franka_pour/reset_utils.py | 532 ++++ .../isaaclab_tasks/utils/__init__.pyi | 17 +- .../utils/adaptive_reset_sampler.py | 464 +++ .../isaaclab_tasks/utils/reset_dataset.py | 295 ++ .../test_franka_pour_cube_bowl_mesh.py | 82 + .../test_franka_pour_cube_bowl_spawner.py | 213 ++ .../contrib/test_franka_pour_curriculum.py | 862 ++++++ .../test/contrib/test_franka_pour_env_cfg.py | 2095 +++++++++++++ .../test/contrib/test_franka_pour_mdp.py | 1754 +++++++++++ .../contrib/test_franka_pour_media_fill.py | 93 + .../contrib/test_franka_pour_reset_bridge.py | 151 + .../contrib/test_franka_pour_reset_dataset.py | 175 ++ ...est_franka_pour_reset_dataset_generator.py | 534 ++++ .../contrib/test_franka_pour_reset_utils.py | 409 +++ .../test/contrib/test_franka_pour_runtime.py | 1283 ++++++++ .../test_franka_pour_spacemouse_teleop.py | 128 + .../contrib/test_franka_pour_visualization.py | 417 +++ .../test/core/test_adaptive_reset_sampler.py | 158 + .../test/utils/test_reset_dataset.py | 214 ++ .../changelog.d/pin-newton-c7ae7c7.rst | 6 + 63 files changed, 22276 insertions(+), 43 deletions(-) create mode 100644 scripts/environments/teleoperation/teleop_franka_pour_spacemouse.py create mode 100644 scripts/tools/generate_franka_pour_reset_dataset.py create mode 100644 scripts/tools/validate_franka_pour_reset_dataset.py create mode 100644 scripts/tools/visualize_franka_pour_reset_dataset.py create mode 100644 source/isaaclab/changelog.d/max-franka-pour-reset-dataset.rst create mode 100644 source/isaaclab_newton/changelog.d/feat-proxycoupledsolver.rst create mode 100644 source/isaaclab_newton/changelog.d/max-franka-pour-reset-dataset.rst create mode 100644 source/isaaclab_newton/changelog.d/mpm-task-support.minor.rst create mode 100644 source/isaaclab_physx/changelog.d/pin-newton-c7ae7c7.rst create mode 100644 source/isaaclab_tasks/changelog.d/franka-pour.minor.rst create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/README.md create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/__init__.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/_reset_collision_screen.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/config/__init__.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/config/franka/__init__.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/config/franka/agents/__init__.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/config/franka/agents/rsl_rl_ppo_cfg.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/cube_bowl_mesh.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/cube_bowl_spawner.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/cube_bowl_spawner_cfg.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/cup_media.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/__init__.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/actions.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/curriculums.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/events.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/observations.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/reset_dataset.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/reset_mixture.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/rewards.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/terminations.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/media_fill.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/pour_env.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/pour_env_cfg.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/reset_dataset_generator.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/reset_utils.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/utils/adaptive_reset_sampler.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/utils/reset_dataset.py create mode 100644 source/isaaclab_tasks/test/contrib/test_franka_pour_cube_bowl_mesh.py create mode 100644 source/isaaclab_tasks/test/contrib/test_franka_pour_cube_bowl_spawner.py create mode 100644 source/isaaclab_tasks/test/contrib/test_franka_pour_curriculum.py create mode 100644 source/isaaclab_tasks/test/contrib/test_franka_pour_env_cfg.py create mode 100644 source/isaaclab_tasks/test/contrib/test_franka_pour_mdp.py create mode 100644 source/isaaclab_tasks/test/contrib/test_franka_pour_media_fill.py create mode 100644 source/isaaclab_tasks/test/contrib/test_franka_pour_reset_bridge.py create mode 100644 source/isaaclab_tasks/test/contrib/test_franka_pour_reset_dataset.py create mode 100644 source/isaaclab_tasks/test/contrib/test_franka_pour_reset_dataset_generator.py create mode 100644 source/isaaclab_tasks/test/contrib/test_franka_pour_reset_utils.py create mode 100644 source/isaaclab_tasks/test/contrib/test_franka_pour_runtime.py create mode 100644 source/isaaclab_tasks/test/contrib/test_franka_pour_spacemouse_teleop.py create mode 100644 source/isaaclab_tasks/test/contrib/test_franka_pour_visualization.py create mode 100644 source/isaaclab_tasks/test/core/test_adaptive_reset_sampler.py create mode 100644 source/isaaclab_tasks/test/utils/test_reset_dataset.py create mode 100644 source/isaaclab_visualizers/changelog.d/pin-newton-c7ae7c7.rst diff --git a/scripts/environments/teleoperation/teleop_franka_pour_spacemouse.py b/scripts/environments/teleoperation/teleop_franka_pour_spacemouse.py new file mode 100644 index 000000000000..382723568e93 --- /dev/null +++ b/scripts/environments/teleoperation/teleop_franka_pour_spacemouse.py @@ -0,0 +1,223 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""SpaceMouse teleoperation for the joint-position Franka Pour task.""" + +from __future__ import annotations + +import argparse +import sys + +import torch + +DEFAULT_TASK = "Isaac-Pour-Franka-Teleop-v0" + + +def joint_targets_to_actions( + *, + joint_targets: torch.Tensor, + action_offset: torch.Tensor, + action_scale: float | torch.Tensor, + lower_limits: torch.Tensor, + upper_limits: torch.Tensor, +) -> torch.Tensor: + """Encode bounded joint targets in the task's normalized joint-action coordinates.""" + joint_targets = torch.clamp(joint_targets, min=lower_limits, max=upper_limits) + scale = torch.as_tensor(action_scale, dtype=joint_targets.dtype, device=joint_targets.device) + if torch.any(scale == 0.0): + raise ValueError("Joint-position action scale must be nonzero.") + return (joint_targets - action_offset) / scale + + +def compose_env_action(arm_action: torch.Tensor, gripper_command: torch.Tensor) -> torch.Tensor: + """Append the normalized symmetric-gripper command to seven arm joint-position actions.""" + if gripper_command.ndim == 1: + gripper_command = gripper_command.unsqueeze(-1) + return torch.cat((arm_action, gripper_command), dim=-1) + + +def apply_tcp_offset_to_jacobian( + jacobian: torch.Tensor, + body_quat: torch.Tensor, + offset_pos: torch.Tensor, +) -> torch.Tensor: + """Move a root-frame geometric Jacobian to a hand-local tool-centre point.""" + from isaaclab.utils import math as math_utils + + result = jacobian.clone() + offset_pos_root = math_utils.quat_apply(body_quat, offset_pos) + result[:, :3, :] += torch.bmm(-math_utils.skew_symmetric_matrix(offset_pos_root), result[:, 3:, :]) + return result + + +def _base_frame_tcp_state_and_jacobian( + robot, + *, + body_idx: int, + joint_ids: list[int], + offset_pos: torch.Tensor, + offset_rot: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return TCP pose and Jacobian in the robot root frame.""" + from isaaclab.utils import math as math_utils + + body_pose_w = robot.data.body_link_pose_w.torch[:, body_idx] + root_pose_w = robot.data.root_link_pose_w.torch + body_pos_b, body_quat_b = math_utils.subtract_frame_transforms( + root_pose_w[:, :3], + root_pose_w[:, 3:7], + body_pose_w[:, :3], + body_pose_w[:, 3:7], + ) + tcp_pos_b, tcp_quat_b = math_utils.combine_frame_transforms( + body_pos_b, + body_quat_b, + offset_pos, + offset_rot, + ) + + jacobian_idx = body_idx - 1 if robot.is_fixed_base else body_idx + jacobian_joint_ids = [joint_id + robot.num_base_dofs for joint_id in joint_ids] + jacobian_w = robot.data.body_link_jacobian_w.torch[:, jacobian_idx, :, jacobian_joint_ids] + root_rot_w = math_utils.matrix_from_quat(math_utils.quat_inv(root_pose_w[:, 3:7])) + jacobian_b = jacobian_w.clone() + jacobian_b[:, :3, :] = torch.bmm(root_rot_w, jacobian_b[:, :3, :]) + jacobian_b[:, 3:, :] = torch.bmm(root_rot_w, jacobian_b[:, 3:, :]) + jacobian_b = apply_tcp_offset_to_jacobian(jacobian_b, body_quat_b, offset_pos) + return tcp_pos_b, tcp_quat_b, jacobian_b + + +def _build_arg_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="SpaceMouse teleoperation for Franka Pour.") + parser.add_argument("--task", default=DEFAULT_TASK, help="Registered Franka Pour task name.") + parser.add_argument("--num_envs", type=int, default=1, help="Number of environments driven by one device.") + parser.add_argument("--max_steps", type=int, default=-1, help="Stop after N steps; negative runs until close.") + parser.add_argument("--mock", action="store_true", help="Use zero device commands for a hardware-free smoke test.") + parser.add_argument("--pos_sensitivity", type=float, default=0.05, help="Translation sensitivity [m].") + parser.add_argument("--rot_sensitivity", type=float, default=0.05, help="Rotation sensitivity [rad].") + parser.add_argument("--ik_damping", type=float, default=0.05, help="Damped-least-squares regularization.") + return parser + + +def main() -> None: + """Launch the task and convert SpaceMouse Cartesian deltas to joint-position actions.""" + import gymnasium as gym + + from isaaclab.app import add_launcher_args, launch_simulation + from isaaclab.controllers import DifferentialIKController, DifferentialIKControllerCfg + from isaaclab.devices import Se3SpaceMouse, Se3SpaceMouseCfg + + import isaaclab_tasks # noqa: F401 + from isaaclab_tasks.utils import resolve_task_config, setup_preset_cli + + parser = _build_arg_parser() + add_launcher_args(parser) + parser.set_defaults(visualizer=["kit"]) + args_cli, hydra_args = setup_preset_cli(parser) + sys.argv = [sys.argv[0], *hydra_args] + env_cfg, _ = resolve_task_config(args_cli.task, "") + + with launch_simulation(env_cfg, args_cli): + env_cfg.scene.num_envs = args_cli.num_envs + if args_cli.device is not None: + env_cfg.sim.device = args_cli.device + env = gym.make(args_cli.task, cfg=env_cfg).unwrapped + + try: + robot = env.scene["robot"] + arm_cfg = env.cfg.actions.arm_action + arm_term = env.action_manager.get_term("arm_action") + joint_ids, joint_names = robot.find_joints( + arm_cfg.joint_names, + preserve_order=arm_cfg.preserve_order, + ) + if len(joint_ids) != 7: + raise RuntimeError(f"Expected seven Franka arm joints, found {joint_names}.") + body_ids, body_names = robot.find_bodies(env.cfg.tcp_body_name) + if len(body_ids) != 1: + raise RuntimeError(f"Expected one TCP body, found {body_names}.") + ik = DifferentialIKController( + DifferentialIKControllerCfg( + command_type="pose", + use_relative_mode=True, + ik_method="dls", + ik_params={"lambda_val": args_cli.ik_damping}, + ), + num_envs=env.num_envs, + device=env.device, + ) + spacemouse = None + if not args_cli.mock: + spacemouse = Se3SpaceMouse( + Se3SpaceMouseCfg( + pos_sensitivity=args_cli.pos_sensitivity, + rot_sensitivity=args_cli.rot_sensitivity, + gripper_term=True, + sim_device=env.device, + ) + ) + reset_requested = False + + def request_reset() -> None: + nonlocal reset_requested + reset_requested = True + + if spacemouse is not None: + spacemouse.add_callback("R", request_reset) + env.reset() + ik.reset() + if spacemouse is not None: + spacemouse.reset() + print(spacemouse) + print("SpaceMouse drives the TCP through joint-position actions; R resets the environment.") + + offset_pos = torch.tensor(env.cfg.tcp_offset_pos, device=env.device).repeat(env.num_envs, 1) + offset_rot = torch.tensor(env.cfg.tcp_offset_rot, device=env.device).repeat(env.num_envs, 1) + + step = 0 + with torch.inference_mode(): + while args_cli.max_steps < 0 or step < args_cli.max_steps: + if args_cli.max_steps < 0 and env.sim.visualizers: + if not any(v.is_running() and not v.is_closed for v in env.sim.visualizers): + break + if reset_requested: + env.reset() + if spacemouse is not None: + spacemouse.reset() + ik.reset() + reset_requested = False + + command = ( + spacemouse.advance() + if spacemouse is not None + else torch.zeros(7, device=env.device, dtype=torch.float32) + ) + tcp_pos_b, tcp_quat_b, jacobian_b = _base_frame_tcp_state_and_jacobian( + robot, + body_idx=body_ids[0], + joint_ids=joint_ids, + offset_pos=offset_pos, + offset_rot=offset_rot, + ) + ik.set_command(command[:6].repeat(env.num_envs, 1), tcp_pos_b, tcp_quat_b) + joint_pos = robot.data.joint_pos.torch[:, joint_ids] + joint_targets = ik.compute(tcp_pos_b, tcp_quat_b, jacobian_b, joint_pos) + joint_limits = robot.data.soft_joint_pos_limits.torch[:, joint_ids] + arm_action = joint_targets_to_actions( + joint_targets=joint_targets, + action_offset=arm_term.action_offset, + action_scale=arm_term.action_scale, + lower_limits=joint_limits[..., 0], + upper_limits=joint_limits[..., 1], + ) + gripper_command = command[6].repeat(env.num_envs) + env.step(compose_env_action(arm_action, gripper_command)) + step += 1 + finally: + env.close() + + +if __name__ == "__main__": + main() diff --git a/scripts/tools/generate_franka_pour_reset_dataset.py b/scripts/tools/generate_franka_pour_reset_dataset.py new file mode 100644 index 000000000000..559733a9c38b --- /dev/null +++ b/scripts/tools/generate_franka_pour_reset_dataset.py @@ -0,0 +1,77 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Generate an oversampled candidate pool for the Franka Pour reset dataset.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +from isaaclab.app import add_launcher_args, launch_simulation + +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.contrib.franka_pour.pour_env import FrankaPourResetSamplerEnv +from isaaclab_tasks.contrib.franka_pour.reset_dataset_generator import ( + FRANKA_POUR_RESET_DATASET_TASK_ID, + FrankaPourResetDatasetGenerator, + FrankaPourResetDatasetGeneratorCfg, + save_reset_dataset, +) +from isaaclab_tasks.utils.parse_cfg import parse_env_cfg + +_REPO_ROOT = Path(__file__).resolve().parents[2] +_DEFAULT_OUTPUT = _REPO_ROOT / "datasets/franka_pour/reset_dataset_candidates.pt" +_CANDIDATE_GRASPING_COUNT = 14_000 +_CANDIDATE_NON_GRASPING_COUNT = 12_000 +_CANDIDATE_NEAR_POUR_COUNT = 2_000 + + +def main() -> None: + """Launch one task world, sample exact category quotas, and atomically save them.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--task", default=FRANKA_POUR_RESET_DATASET_TASK_ID) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--batch_size", type=int, default=256) + parser.add_argument("--max_attempt_multiplier", type=int, default=100) + parser.add_argument("--grasping_count", type=int, default=_CANDIDATE_GRASPING_COUNT) + parser.add_argument("--non_grasping_count", type=int, default=_CANDIDATE_NON_GRASPING_COUNT) + parser.add_argument("--near_pour_grasp_count", type=int, default=_CANDIDATE_NEAR_POUR_COUNT) + parser.add_argument( + "--output", + type=Path, + default=_DEFAULT_OUTPUT, + help="Candidate dataset written before dynamic validation.", + ) + add_launcher_args(parser) + args = parser.parse_args() + + # One scene world is sufficient: IK and collision candidates are replicated in sampler-owned + # Newton models. This avoids allocating an MPM training batch merely to generate reset data. + env_cfg = parse_env_cfg(args.task, device=args.device, num_envs=1) + env_cfg.seed = args.seed + with launch_simulation(env_cfg, args): + env = FrankaPourResetSamplerEnv(env_cfg) + try: + sampler_cfg = FrankaPourResetDatasetGeneratorCfg( + grasping_count=args.grasping_count, + non_grasping_count=args.non_grasping_count, + near_pour_grasp_count=args.near_pour_grasp_count, + batch_size=args.batch_size, + seed=args.seed, + max_attempt_multiplier=args.max_attempt_multiplier, + ) + payload = FrankaPourResetDatasetGenerator(env, sampler_cfg).generate() + save_reset_dataset(payload, args.output) + print(f"[INFO] Wrote {payload['metadata']['state_count']} candidate states to {args.output.resolve()}.") + print(f"[INFO] Contract SHA-256: {payload['contract_sha256']}") + print(f"[INFO] Content SHA-256: {payload['content_sha256']}") + print("[INFO] Next: validate_franka_pour_reset_dataset.py") + finally: + env.close() + + +if __name__ == "__main__": + main() diff --git a/scripts/tools/validate_franka_pour_reset_dataset.py b/scripts/tools/validate_franka_pour_reset_dataset.py new file mode 100644 index 000000000000..f1219e1fb159 --- /dev/null +++ b/scripts/tools/validate_franka_pour_reset_dataset.py @@ -0,0 +1,310 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Validate Franka Pour reset candidates and write the production reset dataset. + +Every candidate is restored into the real task and simulated under a neutral arm command. Grasping +states receive a close-gripper command. The validator rejects non-finite, out-of-bounds, or +particle-workspace failures and requires grasping states to retain bilateral contact after settling. +The output remains balanced across grasp side and broad/near-pour strata. +""" + +from __future__ import annotations + +import argparse +import json +from dataclasses import replace +from pathlib import Path + +import torch + +from isaaclab.app import add_launcher_args, launch_simulation + +import isaaclab_tasks # noqa: F401 +from isaaclab_tasks.contrib.franka_pour.mdp.terminations import source_grasp_milestones +from isaaclab_tasks.contrib.franka_pour.pour_env import FrankaPourResetDatasetValidationEnv +from isaaclab_tasks.contrib.franka_pour.reset_dataset_generator import ( + FRANKA_POUR_RESET_DATASET_TASK_ID, + GRASPING_CATEGORY, + NON_GRASPING_CATEGORY, + RESET_DATASET_GRASPING_COUNT, + RESET_DATASET_NEAR_POUR_COUNT, + RESET_DATASET_NON_GRASPING_COUNT, + FrankaPourResetDatasetGeneratorCfg, + build_reset_dataset_payload, + normalize_grasp_objectives, + save_reset_dataset, + select_production_reset_rows, + validate_reset_dataset, +) +from isaaclab_tasks.utils.parse_cfg import parse_env_cfg + +_REPO_ROOT = Path(__file__).resolve().parents[2] +_DEFAULT_INPUT = _REPO_ROOT / "datasets/franka_pour/reset_dataset_candidates.pt" +_DEFAULT_OUTPUT = _REPO_ROOT / "datasets/franka_pour/reset_dataset.pt" +_DEFAULT_REPORT = _REPO_ROOT / "datasets/franka_pour/reset_dataset_validation.json" +_FAILURE_NAMES = ( + "nonfinite", + "extreme_rigid_state", + "particle_out_of_bounds", + "grasp_lost", + "near_pour_missed_target", + "near_pour_excessive_spill", +) + + +def _never_terminate( + env, + dwell_time_s: float | None = None, + min_lift_height: float | None = None, + max_tcp_distance: float | None = None, + max_gripper_width_error: float | None = None, + max_gripper_command: float | None = None, + terminate: bool | None = None, +) -> torch.Tensor: + """Keep validation rows alive for the complete settling window.""" + del ( + dwell_time_s, + min_lift_height, + max_tcp_distance, + max_gripper_width_error, + max_gripper_command, + terminate, + ) + return torch.zeros(env.num_envs, device=env.device, dtype=torch.bool) + + +def _parse_args() -> argparse.Namespace: + """Parse and validate command-line arguments.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--task", default=FRANKA_POUR_RESET_DATASET_TASK_ID) + parser.add_argument("--input", type=Path, default=_DEFAULT_INPUT, help="Candidate reset dataset.") + parser.add_argument("--output", type=Path, default=_DEFAULT_OUTPUT, help="Validated production dataset.") + parser.add_argument("--report", type=Path, default=_DEFAULT_REPORT, help="Compact JSON validation report.") + parser.add_argument("--num_envs", type=int, default=512) + parser.add_argument( + "--steps", + type=int, + default=60, + help="Policy steps simulated per state; 60 is two seconds at 30 Hz.", + ) + parser.add_argument( + "--settle_steps", + type=int, + default=8, + help="Initial steps excluded from grasp-retention checks.", + ) + parser.add_argument( + "--failure_dwell_steps", + type=int, + default=2, + help="Consecutive failed contact checks required to reject a grasp.", + ) + add_launcher_args(parser) + args = parser.parse_args() + if args.num_envs <= 0 or args.steps <= 0: + parser.error("num_envs and steps must be positive.") + if args.settle_steps < 0 or args.settle_steps >= args.steps: + parser.error("settle_steps must lie in [0, steps).") + if args.failure_dwell_steps <= 0: + parser.error("failure_dwell_steps must be positive.") + return args + + +def _physical_failure_bits(env) -> torch.Tensor: + """Return bit-packed physical failures for every environment.""" + bits = (~env.state_finite()).to(torch.uint8) + bits |= (~env.rigid_state_in_bounds()).to(torch.uint8) << 1 + bits |= (~env.particles_in_workspace()).to(torch.uint8) << 2 + return bits + + +def _build_validated_payload( + payload: dict, + keep: torch.Tensor, + *, + steps: int, + settle_steps: int, + failure_dwell_steps: int, + failure_counts: dict[str, int], + balance_trimmed: int, +) -> dict: + """Build a self-consistent dataset from the dynamically valid rows.""" + states = {name: value[keep].clone() for name, value in payload["states"].items()} + grasping = states["category"] == GRASPING_CATEGORY + states["objective"][grasping] = normalize_grasp_objectives(states["objective_raw"][grasping]) + + sampler_cfg = replace( + FrankaPourResetDatasetGeneratorCfg(**payload["metadata"]["sampler_cfg"]), + grasping_count=RESET_DATASET_GRASPING_COUNT, + non_grasping_count=RESET_DATASET_NON_GRASPING_COUNT, + near_pour_grasp_count=RESET_DATASET_NEAR_POUR_COUNT, + ) + metadata = dict(payload["metadata"]) + metadata["objective_raw_min_max"] = torch.stack( + (states["objective_raw"][grasping].min(), states["objective_raw"][grasping].max()) + ) + metadata["dynamic_validation"] = { + "source_content_sha256": payload["content_sha256"], + "steps": steps, + "settle_steps": settle_steps, + "failure_dwell_steps": failure_dwell_steps, + "failure_counts": failure_counts, + "balance_trimmed": balance_trimmed, + } + return build_reset_dataset_payload( + states, + payload["particle_layouts"]["local_position"], + metadata, + sampler_cfg, + ) + + +def main() -> None: + """Validate every candidate and write the filtered production dataset.""" + args = _parse_args() + input_path = args.input.expanduser().resolve() + payload = torch.load(input_path, map_location="cpu", weights_only=True) + validate_reset_dataset(payload) + row_count = int(payload["metadata"]["state_count"]) + env_count = min(args.num_envs, row_count) + + env_cfg = parse_env_cfg(args.task, device=args.device, num_envs=env_count) + env_cfg.seed = int(payload["metadata"]["seed"]) + env_cfg.reset_dataset_path = str(input_path) + env_cfg.reset_dataset_content_sha256 = payload["content_sha256"] + env_cfg.curriculum_freeze = True + for term_name in ( + "failure", + "extreme_rigid_state", + "lost_grasp", + "spill", + "particle_out_of_bounds", + "success", + "time_out", + ): + getattr(env_cfg.terminations, term_name).func = _never_terminate + + failures = torch.zeros(row_count, dtype=torch.uint8) + with launch_simulation(env_cfg, args): + env = FrankaPourResetDatasetValidationEnv(env_cfg) + try: + task = env + task.sim._app_control_on_stop_handle = None + actions = torch.zeros((env_count, task.action_manager.total_action_dim), device=task.device) + gripper_term_index = task.action_manager.active_terms.index("gripper_action") + if task.action_manager.action_term_dim[gripper_term_index] != 1: + raise RuntimeError("Reset validation requires a one-dimensional gripper action.") + gripper_action_index = sum(task.action_manager.action_term_dim[:gripper_term_index]) + category = payload["states"]["category"].to(task.device) + grasp_region = payload["states"]["grasp_region"].to(task.device) + + for first_row in range(0, row_count, env_count): + active_count = min(env_count, row_count - first_row) + rows = torch.arange(first_row, first_row + active_count, device=task.device) + padded_rows = rows + if active_count < env_count: + padded_rows = torch.cat((rows, rows[:1].expand(env_count - active_count))) + task._forced_reset_dataset_row.copy_(padded_rows) + env.reset() + if not bool(torch.equal(task.reset_dataset_row_id, padded_rows)): + raise RuntimeError("Forced reset rows did not restore the expected dataset entries.") + + actions.zero_() + grasping = category[padded_rows] == GRASPING_CATEGORY + near_pour = grasping & (grasp_region[padded_rows] == 1) + actions[:, gripper_action_index] = torch.where( + grasping, + actions.new_tensor(-1.0), + actions.new_tensor(1.0), + ) + batch_failures = _physical_failure_bits(task) + loss_streak = torch.zeros(env_count, device=task.device, dtype=torch.long) + near_pour_succeeded = ~near_pour + near_pour_spilled = torch.zeros(env_count, device=task.device, dtype=torch.bool) + for step in range(args.steps): + env.step(actions) + batch_failures |= _physical_failure_bits(task) + target_fraction = task.count_in_target() / max(task.num_particles, 1) + spill_fraction = task.spilled_fraction() + within_spill_limit = spill_fraction <= float(task.cfg.max_spill_fraction) + near_pour_succeeded |= ( + near_pour & (target_fraction >= float(task.cfg.pour_target_frac)) & within_spill_limit + ) + near_pour_spilled |= near_pour & ~within_spill_limit + if step < args.settle_steps: + continue + preloaded = source_grasp_milestones( + task, + min_lift_height=task.cfg.success_min_lift_height, + max_tcp_distance=task.cfg.success_max_tcp_distance, + max_gripper_width_error=task.cfg.success_max_gripper_width_error, + max_gripper_command=task.cfg._resolved_success_max_gripper_command(), + )[1] + loss_streak = torch.where( + grasping & preloaded, + torch.zeros_like(loss_streak), + torch.where(grasping, loss_streak + 1, torch.zeros_like(loss_streak)), + ) + batch_failures |= (loss_streak >= args.failure_dwell_steps).to(torch.uint8) << 3 + + final_preloaded = source_grasp_milestones( + task, + min_lift_height=task.cfg.success_min_lift_height, + max_tcp_distance=task.cfg.success_max_tcp_distance, + max_gripper_width_error=task.cfg.success_max_gripper_width_error, + max_gripper_command=task.cfg._resolved_success_max_gripper_command(), + )[1] + batch_failures |= (grasping & ~final_preloaded).to(torch.uint8) << 3 + batch_failures |= (~near_pour_succeeded).to(torch.uint8) << 4 + batch_failures |= near_pour_spilled.to(torch.uint8) << 5 + failures[first_row : first_row + active_count] = batch_failures[:active_count].cpu() + passed = int((batch_failures[:active_count] == 0).sum()) + print(f"[RESET VALIDATION] rows {first_row}:{first_row + active_count} passed {passed}/{active_count}") + finally: + env.close() + + valid = failures == 0 + keep, balance_trimmed = select_production_reset_rows(payload["states"], valid) + failure_counts = {name: int(((failures & (1 << bit)) != 0).sum()) for bit, name in enumerate(_FAILURE_NAMES)} + validated_payload = _build_validated_payload( + payload, + keep, + steps=args.steps, + settle_steps=args.settle_steps, + failure_dwell_steps=args.failure_dwell_steps, + failure_counts=failure_counts, + balance_trimmed=int(balance_trimmed.sum()), + ) + output_path = args.output.expanduser().resolve() + save_reset_dataset(validated_payload, output_path) + + category = payload["states"]["category"] + report = { + "input": str(input_path), + "output": str(output_path), + "source_content_sha256": payload["content_sha256"], + "content_sha256": validated_payload["content_sha256"], + "candidate_count": row_count, + "dynamically_valid_count": int(valid.sum()), + "retained_count": int(keep.sum()), + "balance_trimmed_count": int(balance_trimmed.sum()), + "failure_counts": failure_counts, + "category_counts": { + "non_grasping": int((keep & (category == NON_GRASPING_CATEGORY)).sum()), + "grasping": int((keep & (category == GRASPING_CATEGORY)).sum()), + }, + } + report_path = args.report.expanduser().resolve() + report_path.parent.mkdir(parents=True, exist_ok=True) + report_path.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8") + print( + f"[RESET VALIDATION] retained {report['retained_count']}/{row_count}; " + f"failures={failure_counts}; dataset={output_path}; report={report_path}" + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/tools/visualize_franka_pour_reset_dataset.py b/scripts/tools/visualize_franka_pour_reset_dataset.py new file mode 100644 index 000000000000..12d4fe123339 --- /dev/null +++ b/scripts/tools/visualize_franka_pour_reset_dataset.py @@ -0,0 +1,335 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Render every Franka Pour reset state into four overlaid PNGs. + +This is an offline dataset diagnostic: it uses Pinocchio and Isaac Lab's Franka +URDF to reconstruct the robot link poses, so it does not launch Kit, Newton, or +a GPU simulation. Every state is rendered; the script never subsamples. +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import matplotlib +import numpy as np +import pinocchio as pin +import torch + +from isaaclab_tasks.contrib.franka_pour.reset_dataset_generator import validate_reset_dataset + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 +from matplotlib.lines import Line2D # noqa: E402 +from mpl_toolkits.mplot3d.art3d import Line3DCollection # noqa: E402 + +_REPO_ROOT = Path(__file__).resolve().parents[2] +_DEFAULT_DATASET = _REPO_ROOT / "datasets/franka_pour/reset_dataset.pt" +_DEFAULT_OUTPUT_DIR = _REPO_ROOT / "datasets/franka_pour/reset_dataset_visualization" +_DEFAULT_URDF = _REPO_ROOT / "source/isaaclab/isaaclab/controllers/config/data/lula_franka_gen.urdf" + +_ROBOT_FRAMES = tuple(f"panda_link{index}" for index in range(8)) + ( + "panda_hand", + "panda_leftfinger", + "panda_rightfinger", +) +_ROBOT_EDGES = np.asarray(tuple((index, index + 1) for index in range(8)) + ((8, 9), (8, 10))) +_BOX_EDGES = np.asarray( + ( + (0, 1), + (1, 3), + (3, 2), + (2, 0), + (4, 5), + (5, 7), + (7, 6), + (6, 4), + (0, 4), + (1, 5), + (2, 6), + (3, 7), + ) +) + + +def _objective_buckets(objective: torch.Tensor) -> tuple[tuple[str, str, torch.Tensor, str], ...]: + """Return the requested four disjoint objective buckets.""" + buckets = ( + ("minus_1", "Non-grasping states: J = -1", objective == -1.0, "#58677c"), + ("0_to_0p5", "Grasping states: 0 ≤ J < 0.5", (objective >= 0.0) & (objective < 0.5), "#276fbf"), + ( + "0p5_to_0p9", + "Grasping states: 0.5 ≤ J < 0.9", + (objective >= 0.5) & (objective < 0.9), + "#d98e04", + ), + ( + "0p9_to_1", + "Grasping states: 0.9 ≤ J ≤ 1", + (objective >= 0.9) & (objective <= 1.0), + "#c43832", + ), + ) + assignment_count = torch.stack([mask for _, _, mask, _ in buckets]).sum(dim=0) + if not bool((assignment_count == 1).all()): + invalid = objective[assignment_count != 1] + raise ValueError( + f"Every state must belong to exactly one objective bucket; found {invalid.numel()} invalid values." + ) + return buckets + + +def _load_dataset(path: Path) -> dict: + """Load the dataset and validate the fields needed by this diagnostic.""" + payload = torch.load(path, map_location="cpu", weights_only=True) + validate_reset_dataset(payload) + try: + states = payload["states"] + contract = payload["metadata"]["task_contract"] + required_states = { + "arm_joint_position": (7,), + "finger_joint_position": (2,), + "source_root_pose": (7,), + "target_root_pose": (7,), + "objective": (), + } + state_count = int(states["objective"].shape[0]) + for name, trailing_shape in required_states.items(): + value = states[name] + if not isinstance(value, torch.Tensor) or tuple(value.shape) != (state_count, *trailing_shape): + raise ValueError(f"Invalid {name!r} tensor shape: expected {(state_count, *trailing_shape)}.") + for name in ("source_box_half", "target_box_half"): + if len(contract[name]) != 3: + raise ValueError(f"Invalid task-contract field {name!r}.") + except KeyError as error: + raise ValueError(f"Dataset is missing required field {error.args[0]!r}.") from error + + if state_count == 0 or not bool(torch.isfinite(states["objective"]).all()): + raise ValueError("The dataset must contain finite objective values and at least one state.") + _objective_buckets(states["objective"]) + return payload + + +def _robot_frame_positions( + urdf_path: Path, + arm_joint_position: torch.Tensor, + finger_joint_position: torch.Tensor, +) -> np.ndarray: + """Evaluate all dataset Franka frame positions with Pinocchio FK.""" + model = pin.buildModelFromUrdf(str(urdf_path)) + data = model.createData() + q = pin.neutral(model) + joint_names = tuple(f"panda_joint{index}" for index in range(1, 8)) + ( + "panda_finger_joint1", + "panda_finger_joint2", + ) + joint_indices = np.asarray([model.joints[model.getJointId(name)].idx_q for name in joint_names]) + frame_ids = tuple(model.getFrameId(name) for name in _ROBOT_FRAMES) + if any(frame_id >= len(model.frames) for frame_id in frame_ids): + raise RuntimeError("The Franka URDF does not contain every expected robot frame.") + + joint_position = torch.cat((arm_joint_position, finger_joint_position), dim=1).numpy() + positions = np.empty((joint_position.shape[0], len(frame_ids), 3), dtype=np.float32) + for row, cached_q in enumerate(joint_position): + q[joint_indices] = cached_q + pin.framesForwardKinematics(model, data, q) + positions[row] = np.asarray([data.oMf[frame_id].translation for frame_id in frame_ids]) + return positions + + +def _quat_xyzw_to_matrix(quaternion: np.ndarray) -> np.ndarray: + """Convert a batch of normalized XYZW quaternions to rotation matrices.""" + quaternion = quaternion / np.linalg.norm(quaternion, axis=-1, keepdims=True).clip(min=1.0e-12) + x, y, z, w = np.moveaxis(quaternion, -1, 0) + return np.stack( + ( + 1.0 - 2.0 * (y * y + z * z), + 2.0 * (x * y - z * w), + 2.0 * (x * z + y * w), + 2.0 * (x * y + z * w), + 1.0 - 2.0 * (x * x + z * z), + 2.0 * (y * z - x * w), + 2.0 * (x * z - y * w), + 2.0 * (y * z + x * w), + 1.0 - 2.0 * (x * x + y * y), + ), + axis=-1, + ).reshape(-1, 3, 3) + + +def _box_segments(poses: torch.Tensor, half_extents: tuple[float, float, float]) -> np.ndarray: + """Return world-space edges of bottom-origin oriented collision-proxy boxes.""" + half_x, half_y, half_z = half_extents + corners = np.asarray( + tuple((x, y, z) for z in (0.0, 2.0 * half_z) for y in (-half_y, half_y) for x in (-half_x, half_x)), + dtype=np.float32, + ) + poses_np = poses.numpy() + rotation = _quat_xyzw_to_matrix(poses_np[:, 3:7]) + world_corners = np.einsum("nij,kj->nki", rotation, corners) + poses_np[:, None, :3] + return world_corners[:, _BOX_EDGES] + + +def _table_grid(lower_xy: np.ndarray, upper_xy: np.ndarray, cells: int = 10) -> np.ndarray: + """Create a z=0 grid spanning the actual tabletop support footprint.""" + x_values = np.linspace(lower_xy[0], upper_xy[0], cells + 1) + y_values = np.linspace(lower_xy[1], upper_xy[1], cells + 1) + segments = [((x, lower_xy[1], 0.0), (x, upper_xy[1], 0.0)) for x in x_values] + segments += [((lower_xy[0], y, 0.0), (upper_xy[0], y, 0.0)) for y in y_values] + return np.asarray(segments, dtype=np.float32) + + +def _common_bounds(*point_sets: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Compute fixed plot bounds shared by all four output images.""" + points = np.concatenate([value.reshape(-1, 3) for value in point_sets], axis=0) + lower = points.min(axis=0) + upper = points.max(axis=0) + lower = np.minimum(lower, np.asarray((-0.15, -0.70, 0.0), dtype=np.float32)) + upper = np.maximum(upper, np.asarray((0.90, 0.70, 0.85), dtype=np.float32)) + padding = np.maximum(0.035 * (upper - lower), np.asarray((0.02, 0.02, 0.015))) + return lower - padding, upper + padding + + +def _render_bucket( + output_path: Path, + title: str, + mask: np.ndarray, + robot_color: str, + robot_segments: np.ndarray, + source_segments: np.ndarray, + target_segments: np.ndarray, + table_segments: np.ndarray, + bounds: tuple[np.ndarray, np.ndarray], + *, + dpi: int, + azimuth: float, + elevation: float, + content_sha256: str, +) -> None: + """Render one objective bucket with every matching state overlaid.""" + count = int(mask.sum()) + alpha = min(0.72, max(0.006, 12.0 / max(count, 1))) + figure = plt.figure(figsize=(12, 10), constrained_layout=True) + axis = figure.add_subplot(111, projection="3d", computed_zorder=False) + axis.add_collection3d(Line3DCollection(table_segments, colors="#9da3aa", linewidths=0.7, alpha=0.65)) + for segments, color, width, alpha_scale, zorder in ( + (robot_segments, robot_color, 0.55, 1.0, 3), + (source_segments, "#e56b2f", 0.40, 0.85, 4), + (target_segments, "#159a9c", 0.40, 0.85, 2), + ): + selected = segments[mask].reshape(-1, 2, 3) + axis.add_collection3d( + Line3DCollection( + selected, + colors=color, + linewidths=width, + alpha=min(1.0, alpha * alpha_scale), + rasterized=True, + zorder=zorder, + ) + ) + + lower, upper = bounds + axis.set_xlim(lower[0], upper[0]) + axis.set_ylim(lower[1], upper[1]) + axis.set_zlim(lower[2], upper[2]) + axis.set_box_aspect(upper - lower) + axis.view_init(elev=elevation, azim=azimuth) + axis.set_proj_type("ortho") + axis.set_xlabel("x [m]") + axis.set_ylabel("y [m]") + axis.set_zlabel("z [m]") + axis.grid(False) + axis.xaxis.pane.set_alpha(0.0) + axis.yaxis.pane.set_alpha(0.0) + axis.zaxis.pane.set_alpha(0.0) + axis.set_title(f"{title}\n{count:,} states — all overlaid, no subsampling", fontsize=16, fontweight="semibold") + axis.legend( + handles=( + Line2D((0,), (0,), color=robot_color, linewidth=2.0, label="Franka link skeleton"), + Line2D((0,), (0,), color="#e56b2f", linewidth=2.0, label="Source cup collision proxy"), + Line2D((0,), (0,), color="#159a9c", linewidth=2.0, label="Target bowl collision proxy"), + Line2D((0,), (0,), color="#9da3aa", linewidth=2.0, label="Table plane"), + ), + loc="upper left", + frameon=False, + ) + output_path.parent.mkdir(parents=True, exist_ok=True) + figure.savefig( + output_path, + dpi=dpi, + facecolor="white", + metadata={ + "Title": title, + "Description": f"All {count} matching states; dataset content SHA-256: {content_sha256}", + }, + ) + plt.close(figure) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--input", type=Path, default=_DEFAULT_DATASET, help="Validated reset dataset.") + parser.add_argument("--output_dir", type=Path, default=_DEFAULT_OUTPUT_DIR, help="Directory for the four PNGs.") + parser.add_argument("--urdf", type=Path, default=_DEFAULT_URDF, help="Franka URDF used for forward kinematics.") + parser.add_argument("--dpi", type=int, default=180) + parser.add_argument("--camera_azimuth", type=float, default=-55.0) + parser.add_argument("--camera_elevation", type=float, default=26.0) + args = parser.parse_args() + if args.dpi <= 0: + raise ValueError("--dpi must be positive.") + + dataset_path = args.input.expanduser().resolve() + output_dir = args.output_dir.expanduser().resolve() + urdf_path = args.urdf.expanduser().resolve() + payload = _load_dataset(dataset_path) + states = payload["states"] + contract = payload["metadata"]["task_contract"] + robot_positions = _robot_frame_positions( + urdf_path, + states["arm_joint_position"], + states["finger_joint_position"], + ) + robot_segments = robot_positions[:, _ROBOT_EDGES] + source_segments = _box_segments(states["source_root_pose"], tuple(contract["source_box_half"])) + target_segments = _box_segments(states["target_root_pose"], tuple(contract["target_box_half"])) + bounds = _common_bounds(robot_segments, source_segments, target_segments) + table_segments = _table_grid( + np.asarray(contract["tabletop_support_lower_xy"], dtype=np.float32), + np.asarray(contract["tabletop_support_upper_xy"], dtype=np.float32), + ) + + outputs = [] + buckets = _objective_buckets(states["objective"]) + for slug, title, mask, robot_color in buckets: + output_path = output_dir / f"reset_states_{slug}.png" + _render_bucket( + output_path, + title, + mask.numpy(), + robot_color, + robot_segments, + source_segments, + target_segments, + table_segments, + bounds, + dpi=args.dpi, + azimuth=args.camera_azimuth, + elevation=args.camera_elevation, + content_sha256=payload["content_sha256"], + ) + outputs.append(output_path) + + counts = [int(mask.sum()) for _, _, mask, _ in buckets] + print(f"[INFO] Rendered all {sum(counts):,} reset states without subsampling.") + print(f"[INFO] Bucket counts (-1, [0,.5), [.5,.9), [.9,1]): {counts}.") + for output_path in outputs: + print(f"[INFO] Wrote {output_path}") + + +if __name__ == "__main__": + main() diff --git a/source/isaaclab/changelog.d/max-franka-pour-reset-dataset.rst b/source/isaaclab/changelog.d/max-franka-pour-reset-dataset.rst new file mode 100644 index 000000000000..8bd33aba3c20 --- /dev/null +++ b/source/isaaclab/changelog.d/max-franka-pour-reset-dataset.rst @@ -0,0 +1,5 @@ +Changed +^^^^^^^ + +* Changed training video capture to cap visualization at one environment and to + run only on rank zero during distributed training. diff --git a/source/isaaclab_newton/changelog.d/feat-proxycoupledsolver.rst b/source/isaaclab_newton/changelog.d/feat-proxycoupledsolver.rst new file mode 100644 index 000000000000..585f0864c5f4 --- /dev/null +++ b/source/isaaclab_newton/changelog.d/feat-proxycoupledsolver.rst @@ -0,0 +1,10 @@ +Changed +^^^^^^^ + +* Extracted the repeated solver-kwargs filtering pattern from + :class:`~isaaclab_newton.physics.NewtonFeatherstoneManager`, + :class:`~isaaclab_newton.physics.NewtonMJWarpManager`, and + :class:`~isaaclab_newton.physics.NewtonXPBDManager` into a shared + :meth:`~isaaclab_newton.physics.NewtonManager._filter_solver_kwargs` helper, + so :class:`NewtonManager` subclasses can reuse it when forwarding + ``solver_cfg`` fields to a Newton solver constructor. diff --git a/source/isaaclab_newton/changelog.d/max-franka-pour-reset-dataset.rst b/source/isaaclab_newton/changelog.d/max-franka-pour-reset-dataset.rst new file mode 100644 index 000000000000..4a6d88f31d5b --- /dev/null +++ b/source/isaaclab_newton/changelog.d/max-franka-pour-reset-dataset.rst @@ -0,0 +1,5 @@ +Fixed +^^^^^ + +* Fixed standalone Newton GL training videos to render only environment zero by + default, independent of the simulation batch size. diff --git a/source/isaaclab_newton/changelog.d/mpm-task-support.minor.rst b/source/isaaclab_newton/changelog.d/mpm-task-support.minor.rst new file mode 100644 index 000000000000..e0f554aebeee --- /dev/null +++ b/source/isaaclab_newton/changelog.d/mpm-task-support.minor.rst @@ -0,0 +1,19 @@ +Added +^^^^^ + +* Added scoped Newton builder-world hooks and independent clone-source builder + copies for tasks that extend replicated Newton worlds. +* Added isolated-world and bounded sparse-grid capacity options to + :class:`~isaaclab_newton.physics.MPMSolverCfg`. +* Added :meth:`~isaaclab_newton.physics.NewtonManager.reset_solver_state` for + clearing solver-owned history after selective simulation-state rewrites. + +Fixed +^^^^^ + +* Fixed Newton articulation poses written during environment reset not reaching + Fabric until a later physics step in declarative MPM scenes. +* Fixed graph-capable Newton solvers being captured before the environment's + initial reset and added solver preparation and status checks around replay. +* Fixed empty reset masks unnecessarily clearing solver-owned state every + physics step. diff --git a/source/isaaclab_newton/isaaclab_newton/cloner/__init__.pyi b/source/isaaclab_newton/isaaclab_newton/cloner/__init__.pyi index a2af3ccf629e..57095f08c06d 100644 --- a/source/isaaclab_newton/isaaclab_newton/cloner/__init__.pyi +++ b/source/isaaclab_newton/isaaclab_newton/cloner/__init__.pyi @@ -6,6 +6,7 @@ __all__ = [ "NewtonReplicateContext", "PHYSICS_CONTEXT", + "copy_newton_source_builder", "newton_builder_world_hook", "newton_physics_replicate", ] @@ -13,6 +14,7 @@ __all__ = [ from .replicate import ( NewtonReplicateContext, PHYSICS_CONTEXT, + copy_newton_source_builder, newton_builder_world_hook, newton_physics_replicate, ) diff --git a/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py b/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py index 0b72d79dd113..d7c2dd0a26ce 100644 --- a/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py +++ b/source/isaaclab_newton/isaaclab_newton/cloner/replicate.py @@ -6,10 +6,12 @@ from __future__ import annotations import contextlib +import copy import re from collections.abc import Callable, Iterator, Sequence from typing import TYPE_CHECKING, TypeAlias +import numpy as np import torch from newton import ModelBuilder from newton._src.usd.schemas import SchemaResolverNewton, SchemaResolverPhysx @@ -65,6 +67,46 @@ def newton_builder_world_hook( hooks.remove(hook) +def copy_newton_source_builder(source_path: str) -> ModelBuilder: + """Return an independent mutable copy of a retained clone-source builder. + + Args: + source_path: Clone-source prim path retained by the active replication + plan. + + Returns: + Builder copy that can be finalized or modified independently. + + Raises: + RuntimeError: If the source path is not part of the active clone plan. + """ + prototype = NewtonManager._cl_protos.get(source_path) + if prototype is None: + available = ", ".join(sorted(NewtonManager._cl_protos)) + raise RuntimeError(f"No Newton clone source for {source_path!r}. Available: {available}") + + def copy_mutable(value): + if isinstance(value, list): + return [copy_mutable(item) for item in value] + if isinstance(value, dict): + return {key: copy_mutable(item) for key, item in value.items()} + if isinstance(value, set): + return set(value) + if isinstance(value, np.ndarray): + return value.copy() + return value + + builder = copy.copy(prototype) + for name, value in vars(prototype).items(): + if isinstance(value, (list, dict, set, np.ndarray)): + setattr(builder, name, copy_mutable(value)) + builder.shape_source = [ + source.copy() if callable(getattr(source, "copy", None)) else copy.copy(source) + for source in prototype.shape_source + ] + return builder + + def _build_newton_builder_from_mapping( stage: Usd.Stage, sources: Sequence[str], diff --git a/source/isaaclab_newton/isaaclab_newton/physics/mpm_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/mpm_manager.py index 4cc6f6798270..20b6099a0bd5 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/mpm_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/mpm_manager.py @@ -28,6 +28,10 @@ def _make_solver_config(solver_cfg: MPMSolverCfg) -> SolverImplicitMPM.Config: grid_type=solver_cfg.grid_type, grid_padding=solver_cfg.grid_padding, max_active_cell_count=solver_cfg.max_active_cell_count, + max_leaf_node_count=solver_cfg.max_leaf_node_count, + max_lower_node_count=solver_cfg.max_lower_node_count, + max_upper_node_count=solver_cfg.max_upper_node_count, + separate_worlds=solver_cfg.separate_worlds, transfer_scheme=solver_cfg.transfer_scheme, integration_scheme=solver_cfg.integration_scheme, critical_fraction=solver_cfg.critical_fraction, @@ -120,15 +124,6 @@ def _build_solver(cls, model: Model, solver_cfg: MPMSolverCfg) -> None: NewtonManager._supports_rigid_body_force_input = False cls._project_outside_colliders = solver_cfg.project_outside_colliders - @classmethod - def _supports_cuda_graph_capture(cls) -> bool: - """Return ``True`` only for fixed-grid MPM. - - Sparse and dense grids reallocate as particles move, which is not - capturable in a CUDA graph; the fixed grid keeps a static topology. - """ - return cls._solver.grid_type == "fixed" - @classmethod def _step_solver( cls, state_0: State, state_1: State, control: Control, contacts: Contacts | None, substep_dt: float diff --git a/source/isaaclab_newton/isaaclab_newton/physics/mpm_manager_cfg.py b/source/isaaclab_newton/isaaclab_newton/physics/mpm_manager_cfg.py index 9b8c49489b49..44b31375afa5 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/mpm_manager_cfg.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/mpm_manager_cfg.py @@ -64,7 +64,32 @@ class MPMSolverCfg(NewtonSolverCfg): """Number of empty cells to add around particles when allocating the grid.""" max_active_cell_count: int = -1 - """Maximum active cell count for dense-grid active subsets. ``-1`` means unlimited.""" + """Maximum active cell count shared by all worlds. + + A positive value also enables capacity-bounded sparse-grid rebuilding. + ``-1`` retains allocating sparse behavior and leaves other grids unbounded. + """ + + max_leaf_node_count: int = -1 + """Maximum sparse-grid leaf-node count shared by all worlds. + + ``-1`` derives the capacity from :attr:`max_active_cell_count`. + """ + + max_lower_node_count: int = -1 + """Maximum sparse-grid lower internal-node count shared by all worlds. + + ``-1`` derives the capacity from the initial topology. + """ + + max_upper_node_count: int = -1 + """Maximum sparse-grid upper internal-node count shared by all worlds. + + ``-1`` derives the capacity from the initial topology. + """ + + separate_worlds: bool = False + """Whether each Newton world uses an independent local MPM grid environment.""" transfer_scheme: Literal["apic", "pic"] = "apic" """Particle-grid transfer scheme.""" diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py index fdcaf956a403..a18cea3961d4 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py @@ -1588,22 +1588,30 @@ def start_simulation(cls) -> None: @staticmethod def _initialize_fabric_body_prims(stage, fabric_hierarchy, usdrt, body_bindings: Sequence[tuple[str, int]]) -> None: """Initialize Fabric body prims used by Newton transform sync.""" + complete_hierarchy = bool(NewtonManager._mpm_object_registry) for prim_path, body_index in body_bindings: prim = stage.GetPrimAtPath(prim_path) if prim.IsValid(): xformable_prim = usdrt.Rt.Xformable(prim) xformable_prim.SetWorldXformFromUsd() else: + if complete_hierarchy: + # Solver-only bodies without USD prims are skipped for MPM hierarchy sync. + continue prim = stage.DefinePrim(prim_path, "Xform") xformable_prim = usdrt.Rt.Xformable(prim) xformable_prim.CreateFabricHierarchyWorldMatrixAttr() prim.CreateAttribute(NewtonManager._newton_index_attr, usdrt.Sdf.ValueTypeNames.UInt, custom=True) prim.GetAttribute(NewtonManager._newton_index_attr).Set(body_index) - # Tag with PhysicsRigidBodyAPI so FabricHierarchyGpuUpdateOptions.RIGID_BODY - # applies Inverse propagation (preserves Newton's world transforms and derives - # local) instead of Forward. - prim.AddAppliedSchema("PhysicsRigidBodyAPI") + if complete_hierarchy: + # MPM scenes publish absolute body poses through Fabric reset-stack locals. + fabric_hierarchy.set_reset_xform_stack(usdrt.Sdf.Path(prim_path), True) + else: + # Tag with PhysicsRigidBodyAPI so FabricHierarchyGpuUpdateOptions.RIGID_BODY + # applies Inverse propagation (preserves Newton's world transforms and derives + # local) instead of Forward. + prim.AddAppliedSchema("PhysicsRigidBodyAPI") fabric_hierarchy.update_world_xforms() diff --git a/source/isaaclab_newton/test/physics/test_newton_fabric_body_sync.py b/source/isaaclab_newton/test/physics/test_newton_fabric_body_sync.py index 865e9e725c47..54105b816c02 100644 --- a/source/isaaclab_newton/test/physics/test_newton_fabric_body_sync.py +++ b/source/isaaclab_newton/test/physics/test_newton_fabric_body_sync.py @@ -179,10 +179,14 @@ def CreateFabricHierarchyWorldMatrixAttr(self): class _FakeFabricHierarchy: def __init__(self): self.update_world_xforms_count = 0 + self.reset_xform_stacks = [] def update_world_xforms(self): self.update_world_xforms_count += 1 + def set_reset_xform_stack(self, path, enabled): + self.reset_xform_stacks.append((path, enabled)) + class _FakeRt: Xformable = _FakeXformable @@ -194,6 +198,7 @@ class _FakeValueTypeNames: class _FakeSdf: ValueTypeNames = _FakeValueTypeNames + Path = str class _FakeUsdrt: @@ -327,6 +332,19 @@ def test_root_pose_write_is_visible_on_next_render_without_step(): sim.register_interactive_scene(None) +def test_initialize_mpm_body_prims_reset_fabric_stack(monkeypatch): + prim = _FakePrim() + prim_path = "/World/envs/env_0/Robot/base" + stage = _FakeStage({prim_path: prim}) + fabric_hierarchy = _FakeFabricHierarchy() + monkeypatch.setattr(NewtonManager, "_mpm_object_registry", [object()], raising=False) + + NewtonManager._initialize_fabric_body_prims(stage, fabric_hierarchy, _FakeUsdrt, [(prim_path, 3)]) + + assert prim.applied_schemas == [] + assert fabric_hierarchy.reset_xform_stacks == [(prim_path, True)] + + @pytest.mark.isaacsim_ci @pytest.mark.skipif(not wp.get_cuda_device_count(), reason="CUDA is unavailable") def test_periodic_cable_is_skipped_by_fabric_sync(): diff --git a/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py b/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py index d4b69cca87ad..ecd7525f06cd 100644 --- a/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py +++ b/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py @@ -31,6 +31,7 @@ import numpy as np import pytest import warp as wp +from isaaclab_newton.cloner import copy_newton_source_builder, newton_builder_world_hook from isaaclab_newton.physics import ( FeatherstoneSolverCfg, KaminoSolverCfg, @@ -373,6 +374,9 @@ def test_mpm_solver_cfg_maps_only_newton_solver_fields(): ("grid_type", "dense"), ("grid_padding", 4), ("max_active_cell_count", 1024), + ("max_leaf_node_count", 512), + ("max_lower_node_count", 128), + ("max_upper_node_count", 32), ("transfer_scheme", "pic"), ("integration_scheme", "gimp"), ("critical_fraction", 0.25), @@ -381,6 +385,7 @@ def test_mpm_solver_cfg_maps_only_newton_solver_fields(): ("collider_basis", "Q1"), ("strain_basis", "P1d"), ("velocity_basis", "B2"), + ("separate_worlds", True), ] @@ -595,24 +600,39 @@ def counting_project(*args, **kwargs): assert calls["n"] == 0 +def test_mpm_solver_cfg_preserves_shared_world_default(): + """World-isolated MPM remains opt-in for backward compatibility.""" + + assert MPMSolverCfg().separate_worlds is False + + @pytest.mark.parametrize( - "grid_type, expected", + "grid_type, advertised, expected", [ - ("fixed", True), - ("sparse", False), - ("dense", False), + ("fixed", True, True), + ("sparse", True, True), + ("dense", False, False), ], ) -def test_mpm_cuda_graph_capture_supports_only_fixed_grid(monkeypatch, grid_type, expected): - """Newton implicit MPM is CUDA-graph capturable only with a fixed grid.""" +def test_mpm_cuda_graph_capture_uses_solver_capability(monkeypatch, grid_type, advertised, expected): + """The manager delegates graph safety to Newton's resolved solver configuration.""" - monkeypatch.setattr(NewtonManager, "_solver", SimpleNamespace(grid_type=grid_type), raising=False) + solver = SimpleNamespace(grid_type=grid_type, supports_graph_capture=advertised) + monkeypatch.setattr(NewtonManager, "_solver", solver, raising=False) assert NewtonMPMManager._supports_cuda_graph_capture() is expected -def test_mpm_unsupported_cuda_graph_capture_uses_eager_execution(monkeypatch): - """Sparse/dense MPM should not enter a CUDA graph capture window.""" +def test_cuda_graph_capture_keeps_legacy_solver_support(monkeypatch): + """Solvers without the optional capability property retain their existing capture path.""" + + monkeypatch.setattr(NewtonManager, "_solver", SimpleNamespace(), raising=False) + + assert NewtonManager._supports_cuda_graph_capture() is True + + +def test_solver_advertised_unsupported_cuda_graph_capture_uses_eager_execution(monkeypatch): + """A solver capability rejection prevents the manager from entering capture.""" from isaaclab.physics import PhysicsManager monkeypatch.setattr( @@ -622,46 +642,379 @@ def test_mpm_unsupported_cuda_graph_capture_uses_eager_execution(monkeypatch): raising=False, ) monkeypatch.setattr(PhysicsManager, "_device", "cuda:0", raising=False) - monkeypatch.setattr(NewtonManager, "_solver", SimpleNamespace(grid_type="sparse"), raising=False) + monkeypatch.setattr(NewtonManager, "_solver", SimpleNamespace(supports_graph_capture=False), raising=False) monkeypatch.setattr(NewtonManager, "_graph", object(), raising=False) - monkeypatch.setattr(NewtonManager, "_graph_capture_pending", True, raising=False) + monkeypatch.setattr(NewtonManager, "_graph_capture_pending", "standard", raising=False) - NewtonMPMManager._capture_or_defer_graph() + NewtonManager._capture_or_defer_graph() assert NewtonManager._graph is None - assert NewtonManager._graph_capture_pending is False + assert NewtonManager._graph_capture_pending is None -def test_cuda_graph_capture_uses_simulation_device(monkeypatch): - """CUDA graph capture should use the simulation device instead of Warp's default device.""" +@pytest.mark.parametrize("usdrt_stage, expected_mode", [(None, "standard"), (object(), "relaxed")]) +def test_cuda_graph_capture_is_deferred_with_explicit_mode(monkeypatch, usdrt_stage, expected_mode): + """Headless and RTX runs schedule their respective capture modes until the first step.""" from isaaclab.physics import PhysicsManager - captured_devices = [] + monkeypatch.setattr(PhysicsManager, "_cfg", SimpleNamespace(use_cuda_graph=True), raising=False) + monkeypatch.setattr(PhysicsManager, "_device", "cuda:1", raising=False) + monkeypatch.setattr(NewtonManager, "_usdrt_stage", usdrt_stage, raising=False) + monkeypatch.setattr(NewtonManager, "_solver", SimpleNamespace(supports_graph_capture=True), raising=False) + monkeypatch.setattr(NewtonManager, "_graph", object(), raising=False) + monkeypatch.setattr(NewtonManager, "_graph_capture_pending", None, raising=False) + + NewtonManager._capture_or_defer_graph() + + assert NewtonManager._graph is None + assert NewtonManager._graph_capture_pending == expected_mode + + +def test_standard_cuda_graph_capture_prepares_solver_before_recording(monkeypatch): + """Solver-owned persistent resources are prepared before the standard capture window.""" + events = [] + contacts = object() captured_graph = object() class FakeScopedCapture: - def __init__(self, device=None): - captured_devices.append(device) + def __init__(self, device=None, **_kwargs): + assert device == "cuda:0" self.graph = captured_graph def __enter__(self): + events.append(("capture", None)) return self def __exit__(self, exc_type, exc_value, traceback): return False - monkeypatch.setattr(PhysicsManager, "_cfg", SimpleNamespace(use_cuda_graph=True), raising=False) - monkeypatch.setattr(PhysicsManager, "_device", "cuda:1", raising=False) - monkeypatch.setattr(NewtonManager, "_usdrt_stage", None, raising=False) - monkeypatch.setattr(NewtonManager, "_solver", None, raising=False) + solver = SimpleNamespace( + supports_graph_capture=True, + prepare_graph_capture=lambda received: events.append(("prepare", received)), + ) + monkeypatch.setattr(NewtonManager, "_solver", solver, raising=False) + monkeypatch.setattr(NewtonManager, "_contacts", contacts, raising=False) monkeypatch.setattr(NewtonManager, "_is_all_graphable", classmethod(lambda cls: False)) - monkeypatch.setattr(NewtonManager, "_simulate_physics_only", classmethod(lambda cls: None)) + monkeypatch.setattr( + NewtonManager, + "_simulate_physics_only", + classmethod(lambda cls: events.append(("simulate", None))), + ) monkeypatch.setattr(wp, "ScopedCapture", FakeScopedCapture) - NewtonManager._capture_or_defer_graph() + graph = NewtonManager._capture_standard_graph("cuda:0") + + assert events == [("prepare", contacts), ("capture", None), ("simulate", None)] + assert graph is captured_graph + + +def test_relaxed_cuda_graph_capture_prepares_solver_before_warmup(monkeypatch): + """The RTX-compatible path prepares solver resources before its eager allocation warmup.""" + import isaaclab_newton.physics.newton_manager as newton_manager_module + + events = [] + contacts = object() + solver = SimpleNamespace(prepare_graph_capture=lambda received: events.append(("prepare", received))) + fake_cudart = SimpleNamespace(cudaStreamCreateWithFlags=lambda *_args: 1) + + monkeypatch.setattr(newton_manager_module, "_cudart", fake_cudart) + monkeypatch.setattr(NewtonManager, "_solver", solver, raising=False) + monkeypatch.setattr(NewtonManager, "_contacts", contacts, raising=False) + monkeypatch.setattr(NewtonManager, "_is_all_graphable", classmethod(lambda cls: False)) + monkeypatch.setattr( + NewtonManager, + "_simulate_physics_only", + classmethod(lambda cls: events.append(("warmup", None))), + ) + monkeypatch.setattr(wp, "get_stream", lambda *_args, **_kwargs: object()) + monkeypatch.setattr(wp, "synchronize_stream", lambda *_args, **_kwargs: None) + + assert NewtonManager._capture_relaxed_graph("cpu") == (None, True) + assert events == [("prepare", contacts), ("warmup", None)] + + +@pytest.mark.parametrize("all_graphable", [True, False]) +def test_manager_checks_solver_status_after_graph_replay(monkeypatch, all_graphable): + """Asynchronous solver failures are inspected after either manager graph path replays.""" + from isaaclab.physics import PhysicsManager + + events = [] + mask = SimpleNamespace(zero_=lambda: None) + solver = SimpleNamespace(check_status=lambda: events.append("status")) + monkeypatch.setattr(PhysicsManager, "_sim", SimpleNamespace(is_playing=lambda: True), raising=False) + monkeypatch.setattr(PhysicsManager, "_cfg", SimpleNamespace(use_cuda_graph=True), raising=False) + monkeypatch.setattr(PhysicsManager, "_device", "cuda:0", raising=False) + monkeypatch.setattr(PhysicsManager, "_sim_time", 0.0, raising=False) + monkeypatch.setattr(NewtonManager, "_model_changes", set(), raising=False) + monkeypatch.setattr(NewtonManager, "_solver_reset_pending", False, raising=False) + monkeypatch.setattr(NewtonManager, "_graph_capture_pending", None, raising=False) + monkeypatch.setattr(NewtonManager, "_graph", object(), raising=False) + monkeypatch.setattr(NewtonManager, "_solver", solver, raising=False) + monkeypatch.setattr(NewtonManager, "_world_reset_mask", mask, raising=False) + monkeypatch.setattr(NewtonManager, "_fk_reset_mask", mask, raising=False) + monkeypatch.setattr(NewtonManager, "_needs_collision_pipeline", False, raising=False) + monkeypatch.setattr(NewtonManager, "_needs_fk_before_step", False, raising=False) + monkeypatch.setattr(NewtonManager, "_solver_dt", 1.0 / 120.0, raising=False) + monkeypatch.setattr(NewtonManager, "_num_substeps", 1, raising=False) + monkeypatch.setattr(NewtonManager, "_decimation", 1, raising=False) + monkeypatch.setattr(NewtonManager, "_adapter", None, raising=False) + monkeypatch.setattr(NewtonManager, "_post_actuator_callbacks", [], raising=False) + monkeypatch.setattr(NewtonManager, "_usdrt_stage", None, raising=False) + monkeypatch.setattr(NewtonManager, "_particle_visual_prims", {}, raising=False) + monkeypatch.setattr(NewtonManager, "_is_all_graphable", classmethod(lambda cls: all_graphable)) + monkeypatch.setattr(NewtonManager, "_reset_solver_internals", classmethod(lambda cls, world_mask: None)) + monkeypatch.setattr(wp, "capture_launch", lambda graph: events.append("replay")) + monkeypatch.setattr(NewtonManager, "_log_solver_debug", classmethod(lambda cls: events.append("debug"))) + + NewtonManager.step() + + assert events == ["replay", "status", "debug"] + + +@pytest.mark.parametrize("mode", ["standard", "relaxed"]) +@pytest.mark.parametrize("all_graphable", [True, False]) +def test_deferred_graph_capture_runs_after_reset_setup_then_replays_once(monkeypatch, mode, all_graphable): + """The first reset is consumed before capture, then the new graph advances exactly one step.""" + from isaaclab.physics import PhysicsManager + + events = [] + graph = object() + + class Mask: + def __init__(self, name): + self.name = name + + def zero_(self): + events.append(f"zero_{self.name}") + + adapter = SimpleNamespace(step=lambda *_args: events.append("actuator")) if not all_graphable else None + solver = SimpleNamespace(check_status=lambda: events.append("status")) + monkeypatch.setattr(PhysicsManager, "_sim", SimpleNamespace(is_playing=lambda: True), raising=False) + monkeypatch.setattr(PhysicsManager, "_cfg", SimpleNamespace(use_cuda_graph=True), raising=False) + monkeypatch.setattr(PhysicsManager, "_device", "cuda:0", raising=False) + monkeypatch.setattr(PhysicsManager, "_sim_time", 0.0, raising=False) + monkeypatch.setattr(NewtonManager, "_model_changes", set(), raising=False) + monkeypatch.setattr(NewtonManager, "_solver_reset_pending", True, raising=False) + monkeypatch.setattr(NewtonManager, "_graph_capture_pending", mode, raising=False) + monkeypatch.setattr(NewtonManager, "_graph", None, raising=False) + monkeypatch.setattr(NewtonManager, "_solver", solver, raising=False) + monkeypatch.setattr(NewtonManager, "_world_reset_mask", Mask("world"), raising=False) + monkeypatch.setattr(NewtonManager, "_fk_reset_mask", Mask("fk"), raising=False) + monkeypatch.setattr(NewtonManager, "_needs_collision_pipeline", True, raising=False) + monkeypatch.setattr(NewtonManager, "_needs_fk_before_step", False, raising=False) + monkeypatch.setattr(NewtonManager, "_solver_dt", 1.0 / 120.0, raising=False) + monkeypatch.setattr(NewtonManager, "_num_substeps", 1, raising=False) + monkeypatch.setattr(NewtonManager, "_decimation", 1, raising=False) + monkeypatch.setattr(NewtonManager, "_adapter", adapter, raising=False) + monkeypatch.setattr(NewtonManager, "_post_actuator_callbacks", [], raising=False) + monkeypatch.setattr(NewtonManager, "_usdrt_stage", None, raising=False) + monkeypatch.setattr(NewtonManager, "_particle_visual_prims", {}, raising=False) + monkeypatch.setattr(NewtonManager, "_is_all_graphable", classmethod(lambda cls: all_graphable)) + monkeypatch.setattr( + NewtonManager, + "_reset_solver_internals", + classmethod(lambda cls, world_mask: events.append("reset")), + ) + monkeypatch.setattr( + NewtonManager, + "_eval_fk", + classmethod(lambda cls, world_mask, fk_mask: events.append("fk")), + ) + monkeypatch.setattr( + NewtonManager, + "_capture_standard_graph", + classmethod(lambda cls, device: events.append("capture_standard") or graph), + ) + monkeypatch.setattr( + NewtonManager, + "_capture_relaxed_graph", + classmethod(lambda cls, device: events.append("capture_relaxed") or (graph, True)), + ) + monkeypatch.setattr(wp, "capture_launch", lambda captured: events.append("replay")) + monkeypatch.setattr(NewtonManager, "_log_solver_debug", classmethod(lambda cls: events.append("debug"))) + monkeypatch.setattr( + NewtonManager, + "_simulate_full", + classmethod(lambda cls: pytest.fail("captured step must not also execute eagerly")), + ) + monkeypatch.setattr( + NewtonManager, + "_simulate_physics_only", + classmethod(lambda cls: pytest.fail("captured step must not also execute eagerly")), + ) + + NewtonManager.step() + + expected = ["reset", "fk", "zero_world", "zero_fk"] + if not all_graphable: + expected.append("actuator") + expected.append(f"capture_{mode}") + if mode == "standard": + expected.extend(["replay", "status"]) + expected.append("debug") + assert events == expected + assert NewtonManager._graph_capture_pending is None + + +def test_reset_solver_state_resets_distinct_buffers_and_deduplicates_aliases(monkeypatch): + """Selective resets cannot revive stale history after a state-buffer swap.""" + calls = [] + state_0 = object() + state_1 = object() + world_mask = SimpleNamespace(zero_=lambda: None) + solver = SimpleNamespace(reset=lambda state, *, world_mask, flags: calls.append((state, world_mask, flags))) + monkeypatch.setattr(NewtonManager, "_solver", solver, raising=False) + monkeypatch.setattr(NewtonManager, "_model", SimpleNamespace(world_count=2), raising=False) + monkeypatch.setattr(NewtonManager, "_state_0", state_0, raising=False) + monkeypatch.setattr(NewtonManager, "_state_1", state_1, raising=False) + + NewtonManager.reset_solver_state(world_mask=world_mask, flags=17) + assert calls == [(state_1, world_mask, 17), (state_0, world_mask, 17)] + + calls.clear() + monkeypatch.setattr(NewtonManager, "_state_1", state_0, raising=False) + NewtonManager.reset_solver_state() + assert calls == [(state_0, None, None)] + + +def test_solver_internal_reset_is_event_gated_and_single_world_uses_full_reset(monkeypatch): + """Absent invalidation does no work; a dirty single world uses a full reset.""" + calls = [] + state = object() + state_1 = object() + world_mask = SimpleNamespace(zero_=lambda: None) + solver = SimpleNamespace(reset=lambda *args, **kwargs: calls.append((args, kwargs))) + monkeypatch.setattr(NewtonManager, "_solver", solver, raising=False) + monkeypatch.setattr(NewtonManager, "_model", SimpleNamespace(world_count=1), raising=False) + monkeypatch.setattr(NewtonManager, "_state_0", state, raising=False) + monkeypatch.setattr(NewtonManager, "_state_1", state_1, raising=False) + monkeypatch.setattr(NewtonManager, "_solver_reset_pending", False, raising=False) + monkeypatch.setattr(NewtonManager, "_world_reset_mask", world_mask, raising=False) + monkeypatch.setattr(NewtonManager, "_fk_reset_mask", SimpleNamespace(zero_=lambda: None), raising=False) + monkeypatch.setattr(NewtonManager, "_eval_fk", classmethod(lambda cls, *_args: None)) + monkeypatch.setattr(NewtonManager, "_usdrt_stage", None, raising=False) + + NewtonManager.forward() + assert calls == [] + + NewtonManager._solver_reset_pending = True + NewtonManager.forward() + assert calls == [((state,), {"world_mask": None, "flags": 0})] + assert NewtonManager._solver_reset_pending is False + + calls.clear() + NewtonManager.reset_solver_state(world_mask=world_mask, flags=17) + assert calls == [ + ((state_1,), {"world_mask": world_mask, "flags": 17}), + ((state,), {"world_mask": world_mask, "flags": 17}), + ] + + +def test_forward_publishes_reset_fk_to_fabric(monkeypatch): + """A public forward call must publish reset joint poses before returning.""" + events = [] + world_mask = SimpleNamespace(zero_=lambda: events.append("clear_world")) + fk_mask = SimpleNamespace(zero_=lambda: events.append("clear_fk")) + + monkeypatch.setattr(NewtonManager, "_world_reset_mask", world_mask, raising=False) + monkeypatch.setattr(NewtonManager, "_fk_reset_mask", fk_mask, raising=False) + monkeypatch.setattr(NewtonManager, "_solver_reset_pending", True, raising=False) + monkeypatch.setattr(NewtonManager, "_usdrt_stage", object(), raising=False) + monkeypatch.setattr( + NewtonManager, + "_reset_solver_internals", + classmethod(lambda cls, mask: events.append(("reset", mask))), + ) + monkeypatch.setattr( + NewtonManager, + "_eval_fk", + classmethod(lambda cls, worlds, articulations: events.append(("fk", worlds, articulations))), + ) + monkeypatch.setattr( + NewtonManager, + "_mark_transforms_dirty", + classmethod(lambda cls: events.append("mark_transforms")), + ) + monkeypatch.setattr( + NewtonManager, + "sync_transforms_to_usd", + classmethod(lambda cls: events.append("sync_transforms")), + ) + + NewtonManager.forward() + + assert events == [ + ("reset", world_mask), + ("fk", world_mask, fk_mask), + "mark_transforms", + "sync_transforms", + "clear_fk", + "clear_world", + ] + + +def test_newton_builder_world_hook_is_scoped_and_preserves_existing_registration(monkeypatch): + def existing(*args): + pass + + def added(*args): + pass + + hooks = [existing] + monkeypatch.setattr(NewtonManager, "_per_world_builder_hooks", hooks, raising=False) + + with newton_builder_world_hook(added): + assert hooks == [existing, added] + assert hooks == [existing] + + with pytest.raises(RuntimeError, match="already registered"): + with newton_builder_world_hook(existing): + pass + assert hooks == [existing] + + +def test_newton_builder_world_hook_cleans_up_after_error(monkeypatch): + def hook(*args): + pass + + hooks = [] + monkeypatch.setattr(NewtonManager, "_per_world_builder_hooks", hooks, raising=False) + + with pytest.raises(RuntimeError, match="stop"): + with newton_builder_world_hook(hook): + raise RuntimeError("stop") + + assert hooks == [] + + +def test_copy_newton_source_builder_detaches_mutable_state_and_geometry(monkeypatch): + geometry = SimpleNamespace(name="mesh", copy=lambda: SimpleNamespace(name="mesh-copy")) + prototype = SimpleNamespace( + values=[1], + mapping={"rows": [2]}, + array=np.asarray([3.0], dtype=np.float32), + shape_source=[geometry, None], + ) + monkeypatch.setattr(NewtonManager, "_cl_protos", {"/World/envs/env_0": prototype}, raising=False) + + builder = copy_newton_source_builder("/World/envs/env_0") + builder.values.append(4) + builder.mapping["rows"].append(5) + builder.array[0] = 6.0 + + assert builder is not prototype + assert builder.shape_source[0].name == "mesh-copy" + assert builder.shape_source[0] is not prototype.shape_source[0] + assert prototype.values == [1] + assert prototype.mapping == {"rows": [2]} + assert prototype.array == pytest.approx([3.0]) + + +def test_copy_newton_source_builder_rejects_unknown_source(monkeypatch): + monkeypatch.setattr(NewtonManager, "_cl_protos", {"/World/known": object()}, raising=False) - assert captured_devices == ["cuda:1"] - assert NewtonManager._graph is captured_graph + with pytest.raises(RuntimeError, match="/World/missing.*Available: /World/known"): + copy_newton_source_builder("/World/missing") # --------------------------------------------------------------------------- diff --git a/source/isaaclab_physx/changelog.d/pin-newton-c7ae7c7.rst b/source/isaaclab_physx/changelog.d/pin-newton-c7ae7c7.rst new file mode 100644 index 000000000000..457c02822bf3 --- /dev/null +++ b/source/isaaclab_physx/changelog.d/pin-newton-c7ae7c7.rst @@ -0,0 +1,6 @@ +Changed +^^^^^^^ + +* Changed the ``newton[sim]`` dependency pin of the ``newton`` extra to Newton + commit ``b691d94db1a03c514de03bfdaf27cb9136fc766f`` and required + ``newton-usd-schemas>=0.4.0`` for Newton's USD parsing. diff --git a/source/isaaclab_rl/isaaclab_rl/entrypoints/common.py b/source/isaaclab_rl/isaaclab_rl/entrypoints/common.py index c2872b2aa1f0..e130ab470540 100644 --- a/source/isaaclab_rl/isaaclab_rl/entrypoints/common.py +++ b/source/isaaclab_rl/isaaclab_rl/entrypoints/common.py @@ -431,12 +431,22 @@ def add_common_train_args( def enable_cameras_for_video(args_cli: argparse.Namespace) -> None: - """Enable camera rendering when video recording or sensor capture is requested. + """Configure rank-local rendering for video and sensor capture. Args: args_cli: Parsed command-line arguments. """ - if getattr(args_cli, "video", False) or getattr(args_cli, "capture_env_sensors", 0) > 0: + video_enabled = getattr(args_cli, "video", False) + if video_enabled and getattr(args_cli, "distributed", False) and int(os.getenv("RANK", "0")) != 0: + args_cli.video = False + video_enabled = False + + if video_enabled: + # Training videos intentionally show one world. Besides avoiding an unreadable overlay of + # thousands of environments, this keeps rendering overhead independent of the training batch. + args_cli.max_visible_envs = 1 + + if video_enabled or getattr(args_cli, "capture_env_sensors", 0) > 0: args_cli.enable_cameras = True diff --git a/source/isaaclab_rl/test/test_entrypoints_common.py b/source/isaaclab_rl/test/test_entrypoints_common.py index 8e1f467e4f08..3b0f484617e6 100644 --- a/source/isaaclab_rl/test/test_entrypoints_common.py +++ b/source/isaaclab_rl/test/test_entrypoints_common.py @@ -197,6 +197,77 @@ def test_common_train_args_include_sensor_capture_options() -> None: assert args_cli.capture_env_sensors_format == "file" +def test_common_train_args_video_flag_uses_capture_defaults() -> None: + """The video flag is sufficient while length and interval remain optional overrides.""" + parser = argparse.ArgumentParser() + add_common_train_args(parser, agent_default=None, agent_help="", include_agent=False) + + default_args = parser.parse_args(["--video"]) + override_args = parser.parse_args(["--video", "--video_length", "17", "--video_interval", "29"]) + + assert default_args.video + assert default_args.video_length == 200 + assert default_args.video_interval == 2000 + assert override_args.video_length == 17 + assert override_args.video_interval == 29 + + +def test_enable_cameras_for_video_limits_primary_rank_to_one_environment(monkeypatch: pytest.MonkeyPatch) -> None: + """Primary distributed rank records one environment when video is requested.""" + monkeypatch.setenv("RANK", "0") + args_cli = argparse.Namespace( + video=True, + distributed=True, + capture_env_sensors=0, + enable_cameras=False, + max_visible_envs=None, + ) + + enable_cameras_for_video(args_cli) + + assert args_cli.video + assert args_cli.enable_cameras + assert args_cli.max_visible_envs == 1 + + +def test_enable_cameras_for_video_disables_non_primary_distributed_rank( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Non-primary distributed ranks neither record nor render training video.""" + monkeypatch.setenv("RANK", "3") + args_cli = argparse.Namespace( + video=True, + distributed=True, + capture_env_sensors=0, + enable_cameras=False, + max_visible_envs=None, + ) + + enable_cameras_for_video(args_cli) + + assert not args_cli.video + assert not args_cli.enable_cameras + assert args_cli.max_visible_envs is None + + +def test_enable_cameras_for_video_ignores_rank_for_single_process(monkeypatch: pytest.MonkeyPatch) -> None: + """A stale rank environment variable does not disable single-process video.""" + monkeypatch.setenv("RANK", "3") + args_cli = argparse.Namespace( + video=True, + distributed=False, + capture_env_sensors=0, + enable_cameras=False, + max_visible_envs=None, + ) + + enable_cameras_for_video(args_cli) + + assert args_cli.video + assert args_cli.enable_cameras + assert args_cli.max_visible_envs == 1 + + def test_enable_cameras_for_video_enables_cameras_for_sensor_capture() -> None: """Sensor capture requires camera rendering even when normal video capture is disabled.""" args_cli = argparse.Namespace(video=False, capture_env_sensors=1, enable_cameras=False) diff --git a/source/isaaclab_tasks/changelog.d/franka-pour.minor.rst b/source/isaaclab_tasks/changelog.d/franka-pour.minor.rst new file mode 100644 index 000000000000..6127d771ea71 --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/franka-pour.minor.rst @@ -0,0 +1,21 @@ +Added +^^^^^ + +* Added the ``Isaac-Pour-Franka-v0`` contributed task, where a Franka pours + granular MPM media between scene-owned cups using proxy-coupled MJWarp and + implicit-MPM solvers. +* Added both staged procedural resets and a reset-dataset training preset. The + latter uses offline rejection sampling, Newton IK and collision validation, + adaptive competence-weighted replay, general fixed-weight rewards, and + particle-based success. +* Added CUDA-graph-captured sparse-grid training with isolated MPM worlds, + sparse-grid playback, visible MPM particles, video-friendly camera framing, + and SpaceMouse teleoperation presets for Franka Pour. + +Deprecated +^^^^^^^^^^ + +* Deprecated the experimental ``Reset-Mixture`` task, configuration, runner, + and curriculum names in favor of their ``Reset-Dataset`` counterparts. The + compatibility names do not make older Cartesian-IK policy checkpoints + compatible with the new relative-joint policy. diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/README.md b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/README.md new file mode 100644 index 000000000000..4fec5d34f679 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/README.md @@ -0,0 +1,61 @@ + + +# Franka Pour reset datasets + +The reset-dataset task uses a two-step offline pipeline. Generation proposes and statically screens +an oversampled pool of 14,000 grasping states (including 2,000 near-pour states) and 12,000 +non-grasping states. Validation restores every candidate in the real task, simulates it, requires +near-pour states to deliver at least 30% without excessive spill, and selects exactly 10,000 valid +states from each category. The normal training environment rejects the intermediate candidate file. +Generation and validation default to the same reset-dataset task configuration; custom ``--task`` +overrides must match in both commands because the cache records a strict physics contract. + +From the repository root, run: + +```bash +./isaaclab.sh -p scripts/tools/generate_franka_pour_reset_dataset.py \ + --device cuda:0 --viz none +./isaaclab.sh -p scripts/tools/validate_franka_pour_reset_dataset.py \ + --device cuda:0 --viz none +``` + +If validation reports an insufficient quota, rerun generation with larger +`--grasping_count`, `--non_grasping_count`, or `--near_pour_grasp_count` values. + +This produces: + +- `datasets/franka_pour/reset_dataset_candidates.pt`: oversampled intermediate proposals; never + use for training. +- `datasets/franka_pour/reset_dataset.pt`: dynamically validated production data. +- `datasets/franka_pour/reset_dataset_validation.json`: compact validation report. + +`datasets/` is intentionally ignored by Git. Treat datasets as external build artifacts, preserve +the content SHA-256 printed by the validator with experiment metadata, and distribute the `.pt` +file separately from source changes. + +Train the reset-dataset task and pin the exact artifact with: + +```bash +./isaaclab.sh train --rl_library rsl_rl \ + --task Isaac-Pour-Franka-Reset-Dataset-v0 \ + --num_envs 2048 --logger wandb --video \ + env.reset_dataset_content_sha256= +``` + +`--video` records one environment on rank zero. The traditional procedural task remains available +as `Isaac-Pour-Franka-v0` and does not require a dataset. + +The former `Reset-Mixture` task and Python names remain as deprecated lookup aliases for one +release. They select this reset-dataset implementation; they do not make its 8-action +relative-joint policy compatible with older 7-action Cartesian-IK checkpoints. + +The adaptive reset sampler is intentionally process-local during distributed training: each rank +learns from its own completed episodes and rank zero supplies the logged curriculum metrics. RSL-RL +checkpoints currently restore the policy and optimizer, not this transient sampling history, so a +resumed run starts the sampler again from its configured initial frontier. The reusable sampler +still exposes `state_dict()` for a future generic environment-state checkpoint integration. diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/__init__.py new file mode 100644 index 000000000000..48e6bd7f47a8 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/__init__.py @@ -0,0 +1,12 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Franka pour task: grasp a bowl full of granular MPM media and pour it into a second bowl. + +The submodules are intentionally import-light at the package level so pure-geometry helpers +(:mod:`cube_bowl_mesh`, :mod:`media_fill`) can be imported and unit-tested without launching the +simulator. The environment classes live in :mod:`pour_env` / :mod:`pour_env_cfg` and the gym +registration in :mod:`config.franka`. +""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/_reset_collision_screen.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/_reset_collision_screen.py new file mode 100644 index 000000000000..7f563d13af1c --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/_reset_collision_screen.py @@ -0,0 +1,333 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Exact endpoint collision screening for independently mixed Franka Pour resets.""" + +from __future__ import annotations + +import newton +import torch +import warp as wp + +from isaaclab.utils import math as math_utils + +_DISTAL_TABLE_BODIES = { + "panda_link3", + "panda_link4", + "panda_link5", + "panda_link6", + "panda_link7", + "panda_hand", + "panda_leftfinger", + "panda_rightfinger", +} +_COMPLETE_TABLE_BODIES = _DISTAL_TABLE_BODIES | {"panda_link1", "panda_link2"} + + +@wp.kernel(enable_backward=False) +def _mark_colliding_reset_worlds( + contact_count: wp.array(dtype=wp.int32), + contact_max: int, + contact_shape0: wp.array(dtype=wp.int32), + contact_shape1: wp.array(dtype=wp.int32), + contact_point0: wp.array(dtype=wp.vec3), + contact_point1: wp.array(dtype=wp.vec3), + contact_normal: wp.array(dtype=wp.vec3), + contact_margin0: wp.array(dtype=wp.float32), + contact_margin1: wp.array(dtype=wp.float32), + shape_body: wp.array(dtype=wp.int32), + shape_world: wp.array(dtype=wp.int32), + shape_obstacle: wp.array(dtype=wp.int32), + shape_table: wp.array(dtype=wp.int32), + body_world: wp.array(dtype=wp.int32), + body_q: wp.array(dtype=wp.transform), + body_robot: wp.array(dtype=wp.int32), + body_tests_table: wp.array(dtype=wp.int32), + body_allows_source_contact: wp.array(dtype=wp.int32), + shape_source: wp.array(dtype=wp.int32), + shape_margin: wp.array(dtype=wp.float32), + check_self_collision: int, + penetration_tolerance: float, + source_finger_penetration_tolerance: float, + colliding_worlds: wp.array(dtype=wp.int32), +): + contact_index = wp.tid() + if contact_index >= contact_max or contact_index >= contact_count[0]: + return + + shape0 = contact_shape0[contact_index] + shape1 = contact_shape1[contact_index] + obstacle0 = shape_obstacle[shape0] != 0 + obstacle1 = shape_obstacle[shape1] != 0 + body0 = shape_body[shape0] + body1 = shape_body[shape1] + if not obstacle0 and not obstacle1: + if check_self_collision == 0 or body0 < 0 or body1 < 0 or body0 == body1: + return + if body_robot[body0] == 0 or body_robot[body1] == 0: + return + world0 = body_world[body0] + world1 = body_world[body1] + if world0 < 0 or world0 != world1: + return + point0_w = wp.transform_point(body_q[body0], contact_point0[contact_index]) + point1_w = wp.transform_point(body_q[body1], contact_point1[contact_index]) + separation = wp.dot(contact_normal[contact_index], point1_w - point0_w) + separation = separation - contact_margin0[contact_index] - contact_margin1[contact_index] + if separation < -penetration_tolerance: + wp.atomic_max(colliding_worlds, world0, 1) + return + if obstacle0 and obstacle1: + return + + obstacle_shape = shape0 + robot_body = body1 + if not obstacle0: + obstacle_shape = shape1 + robot_body = body0 + if robot_body < 0 or body_robot[robot_body] == 0: + return + if shape_table[obstacle_shape] != 0 and body_tests_table[robot_body] == 0: + return + + world = body_world[robot_body] + if world < 0 or shape_world[obstacle_shape] != world: + return + transform0 = wp.transform_identity() + transform1 = wp.transform_identity() + if body0 >= 0: + transform0 = body_q[body0] + if body1 >= 0: + transform1 = body_q[body1] + point0_w = wp.transform_point(transform0, contact_point0[contact_index]) + point1_w = wp.transform_point(transform1, contact_point1[contact_index]) + separation = wp.dot(contact_normal[contact_index], point1_w - point0_w) + separation = separation - contact_margin0[contact_index] - contact_margin1[contact_index] + allowed_penetration = penetration_tolerance + if shape_source[obstacle_shape] != 0 and body_allows_source_contact[robot_body] != 0: + # Newton's contact margins include both geometric effective radii and the shapes' + # speculative collision margins. Keep the effective radii in the distance, but do not + # treat an intentionally configured speculative margin as physical finger-cup overlap. + separation = separation + shape_margin[shape0] + shape_margin[shape1] + allowed_penetration = source_finger_penetration_tolerance + if separation < -allowed_penetration: + wp.atomic_max(colliding_worlds, world, 1) + + +def collision_free_reset_candidates( + prototype_builder: newton.ModelBuilder, + robot_q: torch.Tensor, + source_positions: torch.Tensor, + source_quaternions: torch.Tensor, + target_positions: torch.Tensor, + *, + source_box_half: tuple[float, float, float], + target_vertices, + target_indices, + collider_margin: float, + device: str, + penetration_tolerance: float = 1.0e-4, + allow_source_finger_contact: bool = False, + source_finger_penetration_tolerance: float = 0.0, + check_self_collision: bool = False, + check_complete_robot_table: bool = False, +) -> torch.Tensor: + """Return exact collision validity for explicit ``(robot, source, target)`` reset triples.""" + candidate_count = robot_q.shape[0] + coordinate_count = prototype_builder.joint_coord_count + if robot_q.shape != (candidate_count, coordinate_count): + raise ValueError(f"robot_q must have shape (N, {coordinate_count}), got {tuple(robot_q.shape)}.") + if source_positions.shape != (candidate_count, 3) or source_quaternions.shape != (candidate_count, 4): + raise ValueError("Source reset poses must have shapes (N, 3) and (N, 4).") + if target_positions.shape != (candidate_count, 3): + raise ValueError("Target reset positions must have shape (N, 3).") + if candidate_count == 0: + return torch.empty(0, device=device, dtype=torch.bool) + if penetration_tolerance < 0.0: + raise ValueError("penetration_tolerance must be nonnegative.") + if source_finger_penetration_tolerance < 0.0: + raise ValueError("source_finger_penetration_tolerance must be nonnegative.") + + half_x, half_y, half_z = (float(value) for value in source_box_half) + center_offset = torch.zeros_like(source_positions) + center_offset[:, 2] = half_z + source_centers = source_positions + math_utils.quat_apply(source_quaternions, center_offset) + source_centers_cpu = source_centers.detach().cpu().tolist() + source_quaternions_cpu = source_quaternions.detach().cpu().tolist() + target_positions_cpu = target_positions.detach().cpu().tolist() + target_mesh = newton.Mesh(target_vertices, target_indices, compute_inertia=False, is_solid=False) + shape_cfg = newton.ModelBuilder.ShapeConfig( + density=0.0, + margin=float(collider_margin), + has_shape_collision=True, + has_particle_collision=False, + is_visible=False, + ) + + builder = newton.ModelBuilder(up_axis=prototype_builder.up_axis) + for candidate, (source_center, source_quaternion, target_position) in enumerate( + zip(source_centers_cpu, source_quaternions_cpu, target_positions_cpu, strict=True) + ): + builder.begin_world(label=f"reset_candidate_{candidate}") + builder.add_builder(prototype_builder) + builder.add_shape_box( + -1, + xform=wp.transform(wp.vec3(*source_center), wp.quat(*source_quaternion)), + hx=half_x, + hy=half_y, + hz=half_z, + cfg=shape_cfg, + label="ResetCandidate/Source", + ) + builder.add_shape_mesh( + -1, + xform=wp.transform(wp.vec3(*target_position), wp.quat_identity()), + mesh=target_mesh, + cfg=shape_cfg, + label="ResetCandidate/Target", + ) + builder.add_shape_plane( + -1, + xform=wp.transform_identity(), + width=0.0, + length=0.0, + cfg=shape_cfg, + label="ResetCandidate/Table", + ) + builder.end_world() + + model = builder.finalize(device=device) + if model.world_count != candidate_count: + raise RuntimeError(f"Expected {candidate_count} reset worlds, got {model.world_count}.") + model_coordinate_count = model.joint_coord_count // candidate_count + if model_coordinate_count != coordinate_count: + raise RuntimeError( + f"Expected {coordinate_count} robot coordinates per reset world, got {model_coordinate_count}." + ) + + body_names = [str(label).rsplit("/", 1)[-1] for label in model.body_label] + body_robot = torch.as_tensor( + ["/Robot/" in str(label) for label in model.body_label], device=device, dtype=torch.int32 + ) + table_body_names = _COMPLETE_TABLE_BODIES if check_complete_robot_table else _DISTAL_TABLE_BODIES + body_tests_table = torch.as_tensor([name in table_body_names for name in body_names], device=device, dtype=torch.int32) + body_allows_source_contact = torch.as_tensor( + [allow_source_finger_contact and name in {"panda_leftfinger", "panda_rightfinger"} for name in body_names], + device=device, + dtype=torch.int32, + ) + prototype_robot_count = sum("/Robot/" in str(label) for label in prototype_builder.body_label) + if prototype_robot_count <= 0 or int(body_robot.sum()) != candidate_count * prototype_robot_count: + raise RuntimeError("Reset validation did not import the expected explicit /Robot/ bodies.") + prototype_table_body_count = sum( + str(label).rsplit("/", 1)[-1] in table_body_names for label in prototype_builder.body_label + ) + if prototype_table_body_count != len(table_body_names) or int(body_tests_table.sum()) != ( + candidate_count * prototype_table_body_count + ): + raise RuntimeError("Reset validation did not import the expected table-tested robot bodies.") + if allow_source_finger_contact: + prototype_finger_count = sum( + str(label).rsplit("/", 1)[-1] in {"panda_leftfinger", "panda_rightfinger"} + for label in prototype_builder.body_label + ) + if prototype_finger_count != 2 or int(body_allows_source_contact.sum()) != 2 * candidate_count: + raise RuntimeError("Reset validation did not import exactly two source-contact finger bodies per world.") + + shape_labels = [str(label) for label in model.shape_label] + explicit_obstacle = torch.as_tensor( + [ + label.endswith(("ResetCandidate/Source", "ResetCandidate/Target", "ResetCandidate/Table")) + for label in shape_labels + ], + device=device, + dtype=torch.int32, + ) + prototype_table = torch.as_tensor( + [ + check_complete_robot_table and ("/Table/" in label or label.endswith("/Table")) + for label in shape_labels + ], + device=device, + dtype=torch.int32, + ) + shape_obstacle = torch.maximum(explicit_obstacle, prototype_table) + shape_table = torch.as_tensor( + [label.endswith("ResetCandidate/Table") for label in shape_labels], + device=device, + dtype=torch.int32, + ) + shape_table = torch.maximum(shape_table, prototype_table) + shape_source = torch.as_tensor( + [label.endswith("ResetCandidate/Source") for label in shape_labels], + device=device, + dtype=torch.int32, + ) + if ( + int(explicit_obstacle.sum()) != 3 * candidate_count + or int(shape_source.sum()) != candidate_count + ): + raise RuntimeError("Reset validation did not build exactly one source, target, and table obstacle per world.") + if check_complete_robot_table and int(prototype_table.sum()) == 0: + raise RuntimeError("Complete reset validation requires the SeattleLab table collision shapes in the prototype.") + + pipeline = newton.CollisionPipeline( + model, + broad_phase="explicit", + include_static_kinematic_pairs=False, + soft_contact_max=0, + verify_buffers=True, + ) + contacts = pipeline.contacts() + if contacts.rigid_contact_max < 3 * candidate_count: + raise RuntimeError( + f"Reset validation contact capacity {contacts.rigid_contact_max} is below {3 * candidate_count}." + ) + state = model.state() + wp.to_torch(model.joint_q).reshape(candidate_count, coordinate_count).copy_(robot_q) + newton.eval_fk(model, model.joint_q, model.joint_qd, state) + pipeline.collide(state, contacts) + wp.synchronize_device(model.device) + generated_contact_count = int(wp.to_torch(contacts.rigid_contact_count)[0]) + if generated_contact_count > contacts.rigid_contact_max: + raise RuntimeError( + f"Reset validation generated {generated_contact_count} contacts for capacity " + f"{contacts.rigid_contact_max}." + ) + colliding_worlds = wp.zeros(candidate_count, dtype=wp.int32, device=model.device) + wp.launch( + _mark_colliding_reset_worlds, + dim=contacts.rigid_contact_max, + inputs=[ + contacts.rigid_contact_count, + contacts.rigid_contact_max, + contacts.rigid_contact_shape0, + contacts.rigid_contact_shape1, + contacts.rigid_contact_point0, + contacts.rigid_contact_point1, + contacts.rigid_contact_normal, + contacts.rigid_contact_margin0, + contacts.rigid_contact_margin1, + model.shape_body, + model.shape_world, + wp.from_torch(shape_obstacle, dtype=wp.int32), + wp.from_torch(shape_table, dtype=wp.int32), + model.body_world, + state.body_q, + wp.from_torch(body_robot, dtype=wp.int32), + wp.from_torch(body_tests_table, dtype=wp.int32), + wp.from_torch(body_allows_source_contact, dtype=wp.int32), + wp.from_torch(shape_source, dtype=wp.int32), + model.shape_margin, + int(check_self_collision), + float(penetration_tolerance), + float(source_finger_penetration_tolerance), + ], + outputs=[colliding_worlds], + device=model.device, + ) + wp.synchronize_device(model.device) + result = (wp.to_torch(colliding_worlds) == 0).clone() + return result diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/config/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/config/__init__.py new file mode 100644 index 000000000000..460a30569089 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/config/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/config/franka/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/config/franka/__init__.py new file mode 100644 index 000000000000..10996c46019f --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/config/franka/__init__.py @@ -0,0 +1,44 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Gym registration for the Franka two-bowl pour MPM task.""" + +import gymnasium as gym + +from . import agents + +_ENV = "isaaclab_tasks.contrib.franka_pour.pour_env:FrankaPourEnv" +_CFG = "isaaclab_tasks.contrib.franka_pour.pour_env_cfg" +_AGENT = f"{agents.__name__}.rsl_rl_ppo_cfg:FrankaPourPPORunnerCfg" +_RESET_DATASET_AGENT = f"{agents.__name__}.rsl_rl_ppo_cfg:FrankaPourResetDatasetPPORunnerCfg" +_RESET_MIXTURE_AGENT = f"{agents.__name__}.rsl_rl_ppo_cfg:FrankaPourResetMixturePPORunnerCfg" + + +def _register(task_id: str, env_cfg: str, agent_cfg: str | None = None) -> None: + """Register one Franka Pour task variant.""" + kwargs = {"env_cfg_entry_point": f"{_CFG}:{env_cfg}"} + if agent_cfg is not None: + kwargs["rsl_rl_cfg_entry_point"] = agent_cfg + gym.register(id=task_id, entry_point=_ENV, disable_env_checker=True, kwargs=kwargs) + +_register("Isaac-Pour-Franka-v0", "FrankaPourEnvCfg", _AGENT) +_register("Isaac-Pour-Franka-Play-v0", "FrankaPourEnvCfg_PLAY", _AGENT) +_register("Isaac-Pour-Franka-Teleop-v0", "FrankaPourEnvCfg_TELEOP") + +for suffix, cfg_name in ( + ("", "FrankaPourEnvCfg_RESET_DATASET"), + ("-Eval", "FrankaPourEnvCfg_RESET_DATASET_EVAL"), + ("-Play", "FrankaPourEnvCfg_RESET_DATASET_PLAY"), +): + _register(f"Isaac-Pour-Franka-Reset-Dataset{suffix}-v0", cfg_name, _RESET_DATASET_AGENT) + +# Deprecated task IDs retained for one release. They select the reset-dataset implementation; +# Cartesian-IK checkpoints from the experimental task remain incompatible with its joint policy. +for suffix, cfg_name in ( + ("", "FrankaPourEnvCfg_RESET_MIXTURE"), + ("-Eval", "FrankaPourEnvCfg_RESET_MIXTURE_EVAL"), + ("-Play", "FrankaPourEnvCfg_RESET_MIXTURE_PLAY"), +): + _register(f"Isaac-Pour-Franka-Reset-Mixture{suffix}-v0", cfg_name, _RESET_MIXTURE_AGENT) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/config/franka/agents/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/config/franka/agents/__init__.py new file mode 100644 index 000000000000..460a30569089 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/config/franka/agents/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/config/franka/agents/rsl_rl_ppo_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/config/franka/agents/rsl_rl_ppo_cfg.py new file mode 100644 index 000000000000..63dd83e4881f --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/config/franka/agents/rsl_rl_ppo_cfg.py @@ -0,0 +1,125 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from isaaclab.utils.configclass import configclass + +from isaaclab_rl.rsl_rl import RslRlMLPModelCfg, RslRlOnPolicyRunnerCfg, RslRlPpoAlgorithmCfg + + +@configclass +class FrankaPourPPORunnerCfg(RslRlOnPolicyRunnerCfg): + num_steps_per_env = 32 + max_iterations = 3000 + save_interval = 50 + experiment_name = "franka_pour" + run_name = "pour" + clip_actions = 1.0 + # W&B uses the active CLI login (or WANDB_API_KEY); set WANDB_MODE=offline + # to keep run data local without changing the task config. + logger = "wandb" + wandb_project = "franka-pour-mpm" + obs_groups = {"actor": ["policy"], "critic": ["policy", "privileged"]} + actor = RslRlMLPModelCfg( + hidden_dims=[256, 128, 64], + activation="elu", + # Curriculum stages intentionally change reset offsets and geometry. Running empirical + # normalization makes the first unseen stage arbitrarily out-of-distribution because + # stage-constant features have zero variance. Observations are physically scaled in the + # environment instead, matching the standard Franka Lift policy setup. + obs_normalization=False, + # Relative joint increments and the continuous gripper command already map normalized + # actions into bounded physical ranges. Use the standard state-independent Gaussian with + # modest initial exploration around the contact-sensitive backward resets. + distribution_cfg=RslRlMLPModelCfg.GaussianDistributionCfg(init_std=0.1, std_type="log"), + ) + critic = RslRlMLPModelCfg( + hidden_dims=[256, 128, 64], + activation="elu", + obs_normalization=False, + ) + algorithm = RslRlPpoAlgorithmCfg( + value_loss_coef=1.0, + use_clipped_value_loss=True, + clip_param=0.2, + # Retain modest exploration when the curriculum first exposes unseen reset geometry. + entropy_coef=1.0e-3, + # Match Isaac Lab's Franka manipulation update depth. The earlier one-fifth-rate, + # two-epoch update left the actor statistically unchanged after millions of transitions. + num_learning_epochs=5, + num_mini_batches=4, + # The contact-sensitive, low-noise policy produces very small KL early in the backward + # curriculum. Keep the standard Franka base rate fixed so an adaptive schedule cannot + # amplify it during those initially short episodes. + learning_rate=1.0e-4, + schedule="fixed", + gamma=0.99, + lam=0.95, + desired_kl=0.01, + max_grad_norm=1.0, + ) + + +@configclass +class FrankaPourResetDatasetPPORunnerCfg(FrankaPourPPORunnerCfg): + """PPO runner calibrated for adaptive reset-dataset training.""" + + @configclass + class ExplorationDistributionCfg(RslRlMLPModelCfg.HeteroscedasticGaussianDistributionCfg): + """State-dependent exploration bounded for independently sampled action noise.""" + + std_range: tuple[float, float] = (0.05, 0.75) + + # The environment runs at 30 Hz, so 32 transitions span a 1.067-second rollout. + num_steps_per_env = 32 + save_interval = 25 + # Keep the 8-action relative-joint policy separate from incompatible Cartesian-IK checkpoints. + experiment_name = "franka_pour_reset_dataset_joint_rel" + run_name = "reset_dataset_joint_rel" + actor = RslRlMLPModelCfg( + hidden_dims=[512, 256, 128, 64], + activation="elu", + # The adaptive cache deliberately changes its reset support as competence grows. Running + # statistics would therefore move the policy's coordinate system at every frontier + # expansion and can destroy behavior on previously mastered rows. The observations are + # physically scaled by the environment, so keep their representation stationary. + obs_normalization=False, + # Mainline RSL-RL has no temporally correlated gSDE. Retain state-dependent exploration, + # but use the standard Isaac Lab entropy scale below so independently sampled noise does + # not grow merely because the environment clips actions. + distribution_cfg=ExplorationDistributionCfg( + # Start well inside the learned range. Initializing exactly at the former upper clamp + # pinned both standard deviation and entropy in the failed joint-relative runs. + init_std=0.25, + std_type="log", + ), + ) + critic = RslRlMLPModelCfg( + hidden_dims=[512, 256, 128, 64], + activation="elu", + obs_normalization=False, + ) + algorithm = RslRlPpoAlgorithmCfg( + value_loss_coef=1.0, + use_clipped_value_loss=True, + clip_param=0.2, + entropy_coef=1.0e-3, + num_learning_epochs=5, + num_mini_batches=4, + learning_rate=1.0e-4, + # The stopped run's small early KL let the adaptive schedule raise this to 1.14e-3, + # coincident with the first curriculum expansion. Keep update size stationary while the + # reset distribution itself changes. + schedule="fixed", + # Preserve the 10 Hz configuration's physical discount and GAE time constants after + # increasing the policy rate by three. + gamma=0.99 ** (1.0 / 3.0), + lam=0.95 ** (1.0 / 3.0), + desired_kl=0.01, + max_grad_norm=1.0, + ) + + +# Deprecated compatibility name; use ``FrankaPourResetDatasetPPORunnerCfg``. +FrankaPourResetMixturePPORunnerCfg = FrankaPourResetDatasetPPORunnerCfg diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/cube_bowl_mesh.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/cube_bowl_mesh.py new file mode 100644 index 000000000000..4c55fad5d1a4 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/cube_bowl_mesh.py @@ -0,0 +1,175 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Watertight hollow-cube bowl collision mesh for the Franka pour task. + +The bowl is an open-top rectangular box (a "hollow cube") generated as a single closed, +consistently outward-wound triangle mesh, so an MPM collider built from it has an unambiguous +inside/outside. The walls and floor must be **at least ~1.5 grid voxels thick**: +MPM resolves mesh colliders at grid-node resolution, so a sub-voxel wall has +no solid interior on the grid and particles tunnel through it. + +Local frame: ``z=0`` is the outer base (table-facing); the cavity floor is at ``z=bottom_thickness``; +the rim is at ``z = bottom_thickness + cavity_depth``. The bowl is centred on the z axis in x/y. +""" + +from __future__ import annotations + +import numpy as np + +# Corner ordering for an axis-aligned ring, CCW viewed from +z: (-,-), (+,-), (+,+), (-,+). +_RING_SIGNS = ((-1.0, -1.0), (1.0, -1.0), (1.0, 1.0), (-1.0, 1.0)) +# Outward (away-from-solid) horizontal normal of outer wall ``k`` (the edge from ring corner k to k+1). +_OUTER_WALL_NORMALS = ((0.0, -1.0, 0.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0), (-1.0, 0.0, 0.0)) + + +def _validate_closed_oriented_mesh(indices: np.ndarray, name: str) -> None: + """Validate that every triangle edge has exactly one oppositely directed partner.""" + triangles = indices.reshape(-1, 3) + directed_edges: dict[tuple[int, int], int] = {} + for triangle in triangles: + edges = ( + (int(triangle[0]), int(triangle[1])), + (int(triangle[1]), int(triangle[2])), + (int(triangle[2]), int(triangle[0])), + ) + for edge in edges: + directed_edges[edge] = directed_edges.get(edge, 0) + 1 + + duplicate_edges = [edge for edge, count in directed_edges.items() if count != 1] + if duplicate_edges: + raise RuntimeError(f"{name} mesh winding is inconsistent: {len(duplicate_edges)} directed edges not unique.") + boundary_edges = [(start, end) for start, end in directed_edges if (end, start) not in directed_edges] + if boundary_edges: + raise RuntimeError(f"{name} mesh is not watertight: {len(boundary_edges)} boundary edges.") + + +def make_cube_bowl_mesh( + *, + inner_width: float, + inner_depth: float, + cavity_depth: float, + wall_thickness: float, + bottom_thickness: float, + validate: bool = True, +) -> tuple[np.ndarray, np.ndarray]: + """Build a watertight, outward-wound open-top hollow-cube bowl collision mesh. + + Args: + inner_width: Inner cavity size along x [m]. + inner_depth: Inner cavity size along y [m]. + cavity_depth: Inner cavity height from the cavity floor to the rim [m]. + wall_thickness: Side-wall thickness [m]; keep >= ~1.5 * MPM voxel_size to avoid tunnelling. + bottom_thickness: Base thickness below the cavity floor [m]; keep >= ~1.5 * voxel_size. + validate: If True, assert the mesh is a closed, consistently-oriented manifold. + + Returns: + ``(vertices, indices)`` with vertices ``(V, 3)`` float32 and triangle indices ``(3F,)`` int32. + """ + ihx, ihy = 0.5 * float(inner_width), 0.5 * float(inner_depth) + ohx, ohy = ihx + float(wall_thickness), ihy + float(wall_thickness) + bt = float(bottom_thickness) + total_h = bt + float(cavity_depth) + + def ring(hx: float, hy: float, z: float) -> np.ndarray: + return np.array([[sx * hx, sy * hy, z] for sx, sy in _RING_SIGNS], dtype=np.float64) + + outer_base = ring(ohx, ohy, 0.0) # [0:4) + outer_top = ring(ohx, ohy, total_h) # [4:8) + inner_top = ring(ihx, ihy, total_h) # [8:12) + inner_floor = ring(ihx, ihy, bt) # [12:16) + vertices = np.vstack([outer_base, outer_top, inner_top, inner_floor]).astype(np.float32) + outer_base_start, outer_top_start, inner_top_start, inner_floor_start = 0, 4, 8, 12 + + faces: list[int] = [] + + def quad(i0: int, i1: int, i2: int, i3: int, normal: tuple[float, float, float]) -> None: + """Emit two triangles for a cyclic quad, wound CCW about ``normal`` (outward from the solid).""" + p0, p1, p2 = vertices[i0], vertices[i1], vertices[i2] + geo_n = np.cross(p1 - p0, p2 - p0) + if float(np.dot(geo_n, np.asarray(normal, dtype=np.float64))) < 0.0: + i1, i3 = i3, i1 # reverse winding + faces.extend([i0, i1, i2, i0, i2, i3]) + + # Outer base (z=0), normal -z. + quad( + outer_base_start + 0, + outer_base_start + 1, + outer_base_start + 2, + outer_base_start + 3, + (0.0, 0.0, -1.0), + ) + for k in range(4): + j = (k + 1) % 4 + # Outer side wall (z=0..total_h), outward horizontal normal. + quad( + outer_base_start + k, + outer_base_start + j, + outer_top_start + j, + outer_top_start + k, + _OUTER_WALL_NORMALS[k], + ) + # Top rim trapezoid (z=total_h), normal +z; the outer-to-inner diagonal tiles the frame. + quad( + outer_top_start + k, + outer_top_start + j, + inner_top_start + j, + inner_top_start + k, + (0.0, 0.0, 1.0), + ) + # Inner cavity wall (z=bottom_thickness..total_h), normal points into the cavity. + inner_n = tuple(-c for c in _OUTER_WALL_NORMALS[k]) + quad( + inner_floor_start + k, + inner_floor_start + j, + inner_top_start + j, + inner_top_start + k, + inner_n, + ) + # Cavity floor (z=bottom_thickness), normal +z (into the cavity). + quad( + inner_floor_start + 0, + inner_floor_start + 1, + inner_floor_start + 2, + inner_floor_start + 3, + (0.0, 0.0, 1.0), + ) + + indices = np.asarray(faces, dtype=np.int32) + + # Auto-orient to outward (positive signed volume), matching make_hemisphere_scoop_mesh. + tris = vertices[indices.reshape(-1, 3)] + signed_vol = float(np.einsum("ij,ij->i", tris[:, 0], np.cross(tris[:, 1], tris[:, 2])).sum()) + if signed_vol < 0.0: + indices = indices.reshape(-1, 3)[:, ::-1].reshape(-1).astype(np.int32) + + if validate: + _validate_closed_oriented_mesh(indices, "Cube bowl") + + return vertices, indices + + +def cube_bowl_inner_bounds( + inner_width: float, + inner_depth: float, + cavity_depth: float, + bottom_thickness: float, +) -> tuple[np.ndarray, np.ndarray]: + """Return the open inner-cavity axis-aligned bounds in the bowl local frame. + + Args: + inner_width: Inner cavity size along x [m]. + inner_depth: Inner cavity size along y [m]. + cavity_depth: Inner cavity height [m]. + bottom_thickness: Base thickness below the cavity floor [m]. + + Returns: + ``(lo, hi)`` each ``(3,)`` float32: the cavity floor corner and the rim corner. + """ + ihx, ihy = 0.5 * float(inner_width), 0.5 * float(inner_depth) + bt = float(bottom_thickness) + lo = np.array([-ihx, -ihy, bt], dtype=np.float32) + hi = np.array([ihx, ihy, bt + float(cavity_depth)], dtype=np.float32) + return lo, hi diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/cube_bowl_spawner.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/cube_bowl_spawner.py new file mode 100644 index 000000000000..ad0d6babd4a3 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/cube_bowl_spawner.py @@ -0,0 +1,166 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Procedural USD spawner for Franka pour cube bowls.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from pxr import Gf, Usd, UsdGeom, UsdPhysics, UsdShade + +import isaaclab.sim as sim_utils +from isaaclab.sim.schemas import SchemaFragment + +from .cube_bowl_mesh import make_cube_bowl_mesh + +if TYPE_CHECKING: + from .cube_bowl_spawner_cfg import CubeBowlSpawnerCfg + + +def _fragments(value: object) -> list[SchemaFragment] | None: + """Return a normalized schema-fragment list, or None for a legacy config.""" + if isinstance(value, SchemaFragment): + return [value] + if isinstance(value, (list, tuple)) and all(isinstance(fragment, SchemaFragment) for fragment in value): + return list(value) + return None + + +def _apply_collision_properties(prim_path: str, properties: object, stage: Usd.Stage) -> None: + """Apply legacy or fragment collision properties to a prim.""" + fragments = _fragments(properties) + if fragments is None: + sim_utils.define_collision_properties(prim_path, properties, stage=stage) + else: + sim_utils.apply_collision_properties(prim_path, fragments, stage=stage) + + +def _apply_mass_properties(prim_path: str, properties: object, stage: Usd.Stage) -> None: + """Apply legacy or fragment mass properties to a prim.""" + fragments = _fragments(properties) + if fragments is None: + sim_utils.define_mass_properties(prim_path, properties, stage=stage) + else: + sim_utils.apply_mass_properties(prim_path, fragments, stage=stage) + + +def _apply_rigid_body_properties(prim_path: str, properties: object, stage: Usd.Stage) -> None: + """Apply legacy or fragment rigid-body properties to a prim.""" + fragments = _fragments(properties) + if fragments is None: + sim_utils.define_rigid_body_properties(prim_path, properties, stage=stage) + else: + sim_utils.apply_rigid_body_properties(prim_path, fragments, stage=stage) + + +@sim_utils.clone +def spawn_cube_bowl( + prim_path: str, + cfg: CubeBowlSpawnerCfg, + translation: tuple[float, float, float] | None = None, + orientation: tuple[float, float, float, float] | None = None, + **kwargs: object, +) -> Usd.Prim: + """Spawn a visual hollow cube bowl with an optional grasp collider. + + Args: + prim_path: Absolute USD path for the bowl root. + cfg: Bowl geometry and rigid-object configuration. + translation: Local translation relative to the parent [m]. Defaults to the origin. + orientation: Local quaternion in ``(x, y, z, w)`` order. Defaults to identity. + **kwargs: Additional clone-spawner options. + + Returns: + The spawned root Xform prim. + + Raises: + ValueError: If a prim already exists at ``prim_path``. + """ + del kwargs + stage = sim_utils.get_current_stage() + if stage.GetPrimAtPath(prim_path).IsValid(): + raise ValueError(f"A prim already exists at path: '{prim_path}'.") + + root_prim = sim_utils.create_prim( + prim_path, + prim_type="Xform", + translation=translation, + orientation=orientation, + stage=stage, + ) + geometry_path = f"{prim_path}/geometry" + mesh_path = f"{geometry_path}/mesh" + UsdGeom.Xform.Define(stage, geometry_path) + + vertices, indices = make_cube_bowl_mesh( + inner_width=cfg.inner_width, + inner_depth=cfg.inner_depth, + cavity_depth=cfg.cavity_depth, + wall_thickness=cfg.wall_thickness, + bottom_thickness=cfg.bottom_thickness, + ) + mesh = UsdGeom.Mesh.Define(stage, mesh_path) + mesh.CreatePointsAttr().Set([Gf.Vec3f(*(float(value) for value in point)) for point in vertices]) + mesh.CreateFaceVertexIndicesAttr().Set(indices.tolist()) + mesh.CreateFaceVertexCountsAttr().Set([3] * (indices.size // 3)) + mesh.CreateSubdivisionSchemeAttr().Set(UsdGeom.Tokens.none) + mesh.CreateExtentAttr().Set( + [ + Gf.Vec3f(*(float(value) for value in vertices.min(axis=0))), + Gf.Vec3f(*(float(value) for value in vertices.max(axis=0))), + ] + ) + mesh.CreateDisplayColorPrimvar(UsdGeom.Tokens.constant).Set([Gf.Vec3f(*cfg.display_color)]) + + visual_material_path = f"{geometry_path}/visual_material" + visual_material = sim_utils.PreviewSurfaceCfg(diffuse_color=cfg.display_color) + visual_material.func(visual_material_path, visual_material) + sim_utils.bind_visual_material(mesh_path, visual_material_path, stage=stage) + + grasp_proxy_prim: Usd.Prim | None = None + if cfg.grasp_proxy_half_extents is not None: + half_x, half_y, half_z = cfg.grasp_proxy_half_extents + grasp_proxy_path = f"{geometry_path}/grasp_proxy" + grasp_proxy_prim = sim_utils.create_prim( + grasp_proxy_path, + prim_type="Cube", + translation=(0.0, 0.0, half_z), + scale=(2.0 * half_x, 2.0 * half_y, 2.0 * half_z), + attributes={ + "size": 1.0, + "extent": [Gf.Vec3f(-0.5), Gf.Vec3f(0.5)], + }, + stage=stage, + ) + UsdGeom.Imageable(grasp_proxy_prim).MakeInvisible() + if cfg.collision_props is None: + UsdPhysics.CollisionAPI.Apply(grasp_proxy_prim) + else: + _apply_collision_properties(grasp_proxy_path, cfg.collision_props, stage) + + if cfg.physics_material is not None: + if cfg.physics_material_path.startswith("/"): + physics_material_path = cfg.physics_material_path + else: + physics_material_path = f"{geometry_path}/{cfg.physics_material_path}" + cfg.physics_material.func(physics_material_path, cfg.physics_material) + if grasp_proxy_prim is not None: + sim_utils.bind_physics_material(grasp_proxy_prim.GetPath(), physics_material_path, stage=stage) + else: + material = UsdShade.Material(stage.GetPrimAtPath(physics_material_path)) + binding_api = UsdShade.MaterialBindingAPI.Apply(root_prim) + binding_api.Bind( + material, + bindingStrength=UsdShade.Tokens.strongerThanDescendants, + materialPurpose="physics", + ) + + if cfg.mass_props is not None: + _apply_mass_properties(prim_path, cfg.mass_props, stage) + if cfg.rigid_props is not None: + _apply_rigid_body_properties(prim_path, cfg.rigid_props, stage) + + return root_prim diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/cube_bowl_spawner_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/cube_bowl_spawner_cfg.py new file mode 100644 index 000000000000..4d99311b780c --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/cube_bowl_spawner_cfg.py @@ -0,0 +1,52 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Configuration for the Franka pour cube-bowl USD spawner.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import MISSING + +from isaaclab.sim.spawners.materials import RigidBodyMaterialBaseCfg +from isaaclab.sim.spawners.spawner_cfg import RigidObjectSpawnerCfg +from isaaclab.utils.configclass import configclass + + +@configclass +class CubeBowlSpawnerCfg(RigidObjectSpawnerCfg): + """Configuration for a visual cube bowl with an optional rigid grasp proxy.""" + + func: Callable | str = "{DIR}.cube_bowl_spawner:spawn_cube_bowl" + + inner_width: float = MISSING + """Inner cavity size along x [m].""" + + inner_depth: float = MISSING + """Inner cavity size along y [m].""" + + cavity_depth: float = MISSING + """Cavity height from the inner floor to the rim [m].""" + + wall_thickness: float = MISSING + """Side-wall thickness [m].""" + + bottom_thickness: float = MISSING + """Base thickness below the inner floor [m].""" + + display_color: tuple[float, float, float] = (0.95, 0.82, 0.16) + """RGB display color in linear color space.""" + + grasp_proxy_half_extents: tuple[float, float, float] | None = None + """Optional collision-proxy half extents along x, y, and z [m].""" + + physics_material_path: str = "material" + """Path of the rigid-body physics material. + + A relative path is resolved below the bowl's ``geometry`` prim. + """ + + physics_material: RigidBodyMaterialBaseCfg | None = None + """Optional rigid-body physics material for contact metadata.""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/cup_media.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/cup_media.py new file mode 100644 index 000000000000..c88f27a5f23a --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/cup_media.py @@ -0,0 +1,125 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Config-side MPM media generation for the dynamic hollow-cube source cup.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +from isaaclab_newton.assets import MPMObjectCfg +from isaaclab_newton.sim.spawners.mpm import MPMPointsCfg + +from .cube_bowl_mesh import cube_bowl_inner_bounds +from .media_fill import cube_fill_points + +if TYPE_CHECKING: + from .pour_env_cfg import FrankaPourEnvCfg + +MEDIA_SPAWN_SEED = 7 +"""Fixed seed for media spawn sampling: every env must emit identical particles.""" + + +def particle_spacing(cfg: FrankaPourEnvCfg) -> float: + """Particle lattice spacing [m]: ``voxel_size / particles_per_cell``.""" + return float(cfg.voxel_size) / max(float(cfg.particles_per_cell), 1.0) + + +def particle_mass_and_radius(cfg: FrankaPourEnvCfg) -> tuple[float, float]: + """Return mass [kg] and radius [m] for one MPM lattice cell. + + Newton's implicit MPM backend treats each particle as a cube with volume + ``8 * radius**3``. A radius of half the lattice spacing therefore makes the + represented particle volume exactly match the cell volume used to compute + mass, preserving the configured material density. + """ + spacing = particle_spacing(cfg) + volume = spacing**3 + return float(volume * cfg.media_material.density), float(0.5 * spacing) + + +def cup_cavity_lattice(cfg: FrankaPourEnvCfg) -> tuple[np.ndarray, np.ndarray]: + """Jittered lattice filling the source cup cavity in its local frame. + + Args: + cfg: The pour env config (cup geometry + MPM spacing fields). + + Returns: + ``(points, cell)`` with cup-local points ``(N, 3)`` float32 and the lattice cell size + ``(3,)`` float32 [m] used for per-particle mass/radius derivation. + """ + spacing = particle_spacing(cfg) + # Keep particle centres at least one lattice spacing / collider margin from the wall. The old + # ``max(voxel_size, 3 * margin)`` inset removed 6 mm on every side of a 37 mm cup; consequently + # a requested 70% fill represented only 39% of the cavity volume. + clearance = max(spacing, float(cfg.collider_margin)) + lo, hi = cube_bowl_inner_bounds( + float(cfg.source_cup_inner_width), + float(cfg.source_cup_inner_depth), + float(cfg.source_cup_cavity_depth), + float(cfg.source_cup_bottom_thickness), + ) + # ``cube_fill_points.fill_frac`` controls seed *height* inside an inset footprint, whereas the + # task config describes the represented MPM *volume*. Choose the nearest whole number of z + # layers whose cubic particle volumes match that requested cavity-volume fraction. + spans = np.maximum((hi - lo)[:2] - 2.0 * clearance, 0.0) + nx, ny = (int(np.floor(span / spacing)) + 1 for span in spans) + cavity_volume = float(np.prod(hi - lo)) + target_count = float(cfg.media_fill_frac) * cavity_volume / spacing**3 + nz = max(1, int(round(target_count / max(nx * ny, 1)))) + fill_depth = min((nz - 1) * spacing + 1.0e-6 * spacing, float(hi[2] - lo[2]) - 2.0 * clearance) + seed_height_frac = max(fill_depth, 0.0) / float(hi[2] - lo[2]) + points = cube_fill_points( + lo, + hi, + spacing=spacing, + fill_frac=seed_height_frac, + clearance=clearance, + jitter=0.05, + seed=MEDIA_SPAWN_SEED, + ) + if points.shape[0] == 0: + raise RuntimeError("Cup media initialization produced no particles; reduce voxel size or clearance.") + cell = np.full(3, spacing, dtype=np.float32) + return points.astype(np.float32, copy=False), cell + + +def transform_points(points: np.ndarray, pos, quat_xyzw) -> np.ndarray: + """Rotate + translate cup-local ``points`` ``(N, 3)`` by an xyzw quaternion and a translation.""" + q = np.asarray(quat_xyzw, dtype=np.float64) + q = q / (np.linalg.norm(q) + 1.0e-12) + xyz = q[:3] + v = points.astype(np.float64) + t = 2.0 * np.cross(np.broadcast_to(xyz, v.shape), v) + rotated = v + float(q[3]) * t + np.cross(np.broadcast_to(xyz, v.shape), t) + return (rotated + np.asarray(pos, dtype=np.float64)).astype(np.float32) + + +def build_media_object_cfg(cfg: FrankaPourEnvCfg, cup_pos, cup_quat_xyzw) -> MPMObjectCfg: + """Build the declarative cup-media :class:`MPMObjectCfg` from the env config. + + Args: + cfg: The pour env config. + cup_pos: World position [m] of the cup body at reset (the cup-local frame origin). + cup_quat_xyzw: World orientation (xyzw quaternion) of the cup body at reset. + + Returns: + An :class:`MPMObjectCfg` whose spawn points fill the cup cavity at the reset pose, with + per-particle mass/radius derived from the lattice cell. + """ + local_points, cell = cup_cavity_lattice(cfg) + world_points = transform_points(local_points, cup_pos, cup_quat_xyzw) + mass, radius = particle_mass_and_radius(cfg) + return MPMObjectCfg( + prim_path="{ENV_REGEX_NS}/Media", + spawn=MPMPointsCfg( + positions=world_points.tolist(), + mass=mass, + radius=radius, + material=cfg.media_material, + visual_color=(0.85, 0.72, 0.45), + ), + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/__init__.py new file mode 100644 index 000000000000..c1fc60d3e6c4 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/__init__.py @@ -0,0 +1,108 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""MDP terms for the Franka pour task (grasp a dynamic cup of MPM media and pour).""" + +from isaaclab.envs.mdp import ( # noqa: F401 + AbsBinaryJointPositionActionCfg, + BinaryJointPositionActionCfg, + DifferentialInverseKinematicsActionCfg, + JointPositionActionCfg, + RelativeJointPositionActionCfg, + action_l2, + action_rate_l2, + joint_pos_rel, + joint_vel_l2, + joint_vel_rel, + last_action, + time_out, +) + +from .actions import ( # noqa: F401 + CurriculumGripperPositionAction, + CurriculumGripperPositionActionCfg, + CurriculumJointPositionAction, + CurriculumJointPositionActionCfg, + TrajectoryJointPositionAction, + TrajectoryJointPositionActionCfg, +) +from .curriculums import PourCurriculum # noqa: F401 +from .events import reset_pour_scene # noqa: F401 +from .observations import ( # noqa: F401 + arm_reference_error_obs, + arm_reference_phase_obs, + cup_pose_obs, + cup_to_target_obs, + cup_velocity_obs, + ee_pose_obs, + finger_position_obs, + finger_velocity_obs, + grasp_to_tcp_quat_obs, + gripper_contact_obs, + gripper_target_obs, + gripper_width_obs, + held_delivery_history_obs, + lost_grasp_dwell_obs, + particle_fractions_obs, + particle_transfer_obs, + pour_target_fraction_obs, + success_dwell_obs, + target_position_c_obs, + target_pose_obs, + tcp_pose_obs, + tcp_to_grasp_obs, + tcp_to_grasp_position_c_obs, + time_remaining_obs, + trajectory_status_obs, +) +from .rewards import ( # noqa: F401 + AlignProgress, + ApproachProgress, + GraspLiftProgress, + HeldDeliveryProgress, + LiftProgress, + NewlyDeliveredParticles, + NewlySpilledParticles, + PourReferenceProgress, + PourTaskProgress, + PourTiltProgress, + align_command_progress, + align_cup_over_target, + finite_joint_velocity_l2, + grasp_cup, + lift_command_progress, + lift_cup, + media_target_distance_tanh, + particles_in_source, + particles_in_target, + pour_success_bonus, + sustained_pour_success, + reach_cup, + spilled_particles, + terminal_failure, + tcp_cup_distance_tanh, + tilt_command_progress, + tilt_over_target, +) +from .reset_dataset import ( # noqa: F401 + PourResetDatasetCurriculum, + reset_dataset_difficulty, +) +from .reset_mixture import ( # noqa: F401 + RESET_MIXTURE_REGION_NAMES, + RESET_MIXTURE_STAGE_NAMES, + PourResetMixture, +) +from .terminations import ( # noqa: F401 + excessive_spill, + extreme_rigid_state, + immediate_pour_success, + lost_lifted_grasp, + nonterminating_stable_pour_success, + nonfinite_failure, + particle_out_of_bounds, + stable_pour_success, + unsuccessful_time_out, +) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/actions.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/actions.py new file mode 100644 index 000000000000..20cc8109a077 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/actions.py @@ -0,0 +1,906 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Reset-relative arm and continuous symmetric-gripper actions.""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from dataclasses import MISSING + +import torch + +from isaaclab.envs.mdp.actions.actions_cfg import JointPositionActionCfg +from isaaclab.envs.mdp.actions.joint_actions import JointPositionAction +from isaaclab.managers import ActionTerm, ActionTermCfg +from isaaclab.utils.configclass import configclass + +_GRIPPER_POSITION_TOLERANCE = 1.0e-6 + + +def _bilateral_gripper_preload( + joint_position: torch.Tensor, + joint_velocity: torch.Tensor, + joint_target: torch.Tensor, + *, + min_deflection: float, + max_velocity: float, + max_command: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """Return per-finger drive deflection and a stable bilateral-contact mask.""" + finite = torch.isfinite(joint_position) & torch.isfinite(joint_velocity) & torch.isfinite(joint_target) + deflection = torch.where(finite, torch.clamp(joint_position - joint_target, min=0.0), 0.0) + bilateral = ( + finite.all(dim=-1) + & (deflection.amin(dim=-1) >= float(min_deflection)) + & (joint_velocity.abs().amax(dim=-1) <= float(max_velocity)) + & (joint_target.amax(dim=-1) <= float(max_command) + _GRIPPER_POSITION_TOLERANCE) + ) + return deflection, bilateral + + +class CurriculumJointPositionAction(JointPositionAction): + """Reset-relative joint-position action with an optional early reference manifold. + + During the first supplied-grasp curriculum stages, unconstrained independent joint targets can + leave the physically validated held-pour trajectory and destabilize the light source cup. The + optional reference projection retains the policy's seven-dimensional interface, but projects + its command onto the line from the reset offset to a validated joint target. Later stages use + the ordinary full-rank joint-position action unchanged. + """ + + cfg: CurriculumJointPositionActionCfg + + def __init__(self, cfg: CurriculumJointPositionActionCfg, env) -> None: + super().__init__(cfg, env) + self._alpha = float(cfg.alpha) + if not 0.0 < self._alpha <= 1.0: + raise ValueError(f"Moving-average weight must lie in (0, 1], got {self._alpha}.") + self._project_reference_through_stage = int(cfg.project_reference_through_stage) + if self._project_reference_through_stage < -1: + raise ValueError("project_reference_through_stage must be at least -1.") + self._reference_action_magnitude = float(cfg.reference_action_magnitude) + if not math.isfinite(self._reference_action_magnitude) or self._reference_action_magnitude <= 0.0: + raise ValueError("reference_action_magnitude must be finite and positive.") + self._reference_action_index = int(cfg.reference_action_index) + if self._reference_action_index < 0 or self._reference_action_index >= self.action_dim: + raise ValueError(f"reference_action_index must lie in [0, {self.action_dim - 1}].") + if cfg.reference_target: + if len(cfg.reference_target) != self.action_dim: + raise ValueError( + f"reference_target must contain {self.action_dim} joint positions, got {len(cfg.reference_target)}." + ) + if any(not math.isfinite(value) for value in cfg.reference_target): + raise ValueError("reference_target must contain only finite joint positions.") + self._reference_target = torch.tensor(cfg.reference_target, device=self.device).repeat(self.num_envs, 1) + else: + self._reference_target = None + if self._project_reference_through_stage >= 0 and self._reference_target is None: + raise ValueError("reference_target is required when reference projection is enabled.") + self._previous_target = self._processed_actions.clone() + + @property + def action_offset(self) -> torch.Tensor: + """Per-environment joint-position action offset [rad].""" + if not isinstance(self._offset, torch.Tensor): + raise RuntimeError("Curriculum joint-position actions require a tensor action offset.") + return self._offset + + @property + def action_scale(self) -> torch.Tensor | float: + """Joint-position displacement represented by one policy-action unit [rad].""" + return self._scale + + def set_action_offset( + self, + offset: torch.Tensor, + env_ids: Sequence[int] | torch.Tensor | slice | None = None, + ) -> None: + """Set selected environments' zero-action joint targets [rad]. + + Args: + offset: Joint targets with shape ``(len(env_ids), action_dim)`` [rad]. + env_ids: Environments to update. If omitted, update every environment. + + Raises: + ValueError: If :paramref:`offset` does not match the selected offset shape. + """ + selected = slice(None) if env_ids is None else env_ids + expected_shape = self.action_offset[selected].shape + if offset.shape != expected_shape: + raise ValueError(f"Action offset shape {tuple(offset.shape)} does not match {tuple(expected_shape)}.") + offset = offset.to(device=self._offset.device, dtype=self._offset.dtype) + self._offset[selected] = offset + if hasattr(self, "_previous_target"): + self._previous_target[selected] = offset + + def process_actions(self, actions: torch.Tensor) -> None: + """Map raw actions to low-pass-filtered joint targets [rad].""" + self._raw_actions.copy_(actions) + effective_actions = actions + if self._project_reference_through_stage >= 0: + projected_worlds = self._env.curriculum_stage <= self._project_reference_through_stage + scale = self._scale + reference_action = (self._reference_target - self.action_offset) / scale + reference_norm_sq = reference_action.square().sum(dim=-1).clamp_min(1.0e-12) + # A fixed policy coordinate preserves the meaning of "pour" when the curriculum reset + # pose changes. The action term expands that scalar into the stage-specific correlated + # joint trajectory; the remaining coordinates are ignored only in these supplied-grasp + # stages and regain their ordinary joint semantics afterward. + commanded_phase = actions[:, self._reference_action_index] / self._reference_action_magnitude + previous_action = (self._previous_target - self.action_offset) / scale + previous_phase = (previous_action * reference_action).sum(dim=-1) / reference_norm_sq + # Pouring is one-way in this supplied-grasp stage. Never command a phase behind the + # observed filtered target: this prevents a stochastic action from reversing the cup + # after transfer, while requiring continued positive commands to advance the EMA. + phase = torch.maximum(commanded_phase, previous_phase).clamp(0.0, 1.0) + projected_actions = reference_action * phase.unsqueeze(-1) + effective_actions = torch.where(projected_worlds.unsqueeze(-1), projected_actions, actions) + self._processed_actions = effective_actions * self._scale + self._offset + if self.cfg.clip is not None: + self._processed_actions = torch.clamp( + self._processed_actions, + min=self._clip[:, :, 0], + max=self._clip[:, :, 1], + ) + self._processed_actions.lerp_(self._previous_target, 1.0 - self._alpha) + self._previous_target.copy_(self._processed_actions) + + def reset(self, env_ids: Sequence[int] | torch.Tensor | slice | None = None) -> None: + """Clear selected raw actions and align their buffered targets with the reset origin.""" + selected = slice(None) if env_ids is None else env_ids + self._raw_actions[selected] = 0.0 + self._processed_actions[selected] = self.action_offset[selected] + self._previous_target[selected] = self.action_offset[selected] + + +@configclass +class CurriculumJointPositionActionCfg(JointPositionActionCfg): + """Configuration for :class:`CurriculumJointPositionAction`.""" + + alpha: float = 0.2 + """Weight of the current joint target in the exponential moving average.""" + + reference_target: tuple[float, ...] = () + """Validated absolute joint target used by the early reference projection.""" + + project_reference_through_stage: int = -1 + """Last curriculum stage projected onto the reset-to-reference segment, or ``-1`` to disable.""" + + reference_action_magnitude: float = 1.0 + """Policy-space scalar magnitude that commands the complete reference segment.""" + + reference_action_index: int = 0 + """Fixed policy-action coordinate used as the early-stage reference phase.""" + + class_type: type[CurriculumJointPositionAction] = CurriculumJointPositionAction + + +class TrajectoryJointPositionAction(ActionTerm): + """Filtered joint-position residuals around a monotonic per-environment reference trajectory. + + The first policy coordinate always controls forward progress and the remaining coordinates + always control joint residuals. Curriculum resets only change the initial progress value and + reference waypoints; they never repurpose policy coordinates. This avoids the distribution + shift caused by making a pour-phase action become an ordinary arm-joint action after a stage + promotion. + """ + + cfg: TrajectoryJointPositionActionCfg + + def __init__(self, cfg: TrajectoryJointPositionActionCfg, env) -> None: + super().__init__(cfg, env) + self._joint_ids, self._joint_names = self._asset.find_joints( + cfg.joint_names, + preserve_order=cfg.preserve_order, + ) + self._num_joints = len(self._joint_ids) + if self._num_joints == 0: + raise ValueError("Trajectory joint-position action resolved no joints.") + if int(cfg.waypoint_count) < 2: + raise ValueError("Trajectory joint-position action requires at least two waypoints.") + self._waypoint_count = int(cfg.waypoint_count) + self._alpha = float(cfg.alpha) + self._phase_rate = float(cfg.phase_rate) + self._approach_phase_rate = float(cfg.approach_phase_rate) + self._transport_phase_rate = float(cfg.transport_phase_rate) + self._waypoint_phases = torch.as_tensor(cfg.waypoint_phases, device=self.device, dtype=torch.float32) + if self._waypoint_phases.shape != (self._waypoint_count,): + raise ValueError(f"waypoint_phases must contain {self._waypoint_count} values.") + if ( + not bool(torch.isfinite(self._waypoint_phases).all()) + or float(self._waypoint_phases[0]) != 0.0 + or float(self._waypoint_phases[-1]) != 1.0 + or not bool(torch.all(self._waypoint_phases[1:] > self._waypoint_phases[:-1])) + ): + raise ValueError("waypoint_phases must increase strictly from 0 to 1.") + milestone_indices = ( + int(cfg.approach_waypoint), + int(cfg.grasp_waypoint), + int(cfg.lift_waypoint), + int(cfg.align_waypoint), + ) + if not ( + 0 + < milestone_indices[0] + < milestone_indices[1] + < milestone_indices[2] + < milestone_indices[3] + < self._waypoint_count + ): + raise ValueError("Approach, grasp, lift, and align waypoints must be strictly ordered interior points.") + self._approach_phase = float(self._waypoint_phases[milestone_indices[0]]) + self._grasp_phase = float(self._waypoint_phases[milestone_indices[1]]) + self._lift_phase = float(self._waypoint_phases[milestone_indices[2]]) + self._align_phase = float(self._waypoint_phases[milestone_indices[3]]) + self._grasp_gate_stage = int(cfg.grasp_gate_stage) + self._grasp_dwell_steps = int(cfg.grasp_dwell_steps) + self._grasp_max_tcp_distance = float(cfg.grasp_max_tcp_distance) + self._grasp_max_linear_velocity = float(cfg.grasp_max_linear_velocity) + self._grasp_max_angular_velocity = float(cfg.grasp_max_angular_velocity) + self._approach_max_lateral_distance = float(cfg.approach_max_lateral_distance) + self._approach_max_joint_error = float(cfg.approach_max_joint_error) + self._approach_dwell_steps = int(cfg.approach_dwell_steps) + self._approach_max_linear_velocity = float(cfg.approach_max_linear_velocity) + self._approach_max_angular_velocity = float(cfg.approach_max_angular_velocity) + self._align_max_distance = float(cfg.align_max_distance) + if not 0.0 < self._alpha <= 1.0: + raise ValueError(f"Moving-average weight must lie in (0, 1], got {self._alpha}.") + if not math.isfinite(self._phase_rate) or self._phase_rate <= 0.0: + raise ValueError("phase_rate must be finite and positive.") + if not math.isfinite(self._approach_phase_rate) or self._approach_phase_rate <= 0.0: + raise ValueError("approach_phase_rate must be finite and positive.") + if not math.isfinite(self._transport_phase_rate) or self._transport_phase_rate <= 0.0: + raise ValueError("transport_phase_rate must be finite and positive.") + if self._grasp_dwell_steps <= 0: + raise ValueError("grasp_dwell_steps must be positive.") + if self._approach_dwell_steps <= 0: + raise ValueError("approach_dwell_steps must be positive.") + if not all( + math.isfinite(value) and value > 0.0 + for value in ( + self._grasp_max_tcp_distance, + self._grasp_max_linear_velocity, + self._grasp_max_angular_velocity, + self._approach_max_lateral_distance, + self._approach_max_joint_error, + self._approach_max_linear_velocity, + self._approach_max_angular_velocity, + self._align_max_distance, + ) + ): + raise ValueError("Grasp stability limits must be finite and positive.") + + residual_scale = torch.as_tensor(cfg.residual_scale, device=self.device, dtype=torch.float32) + if residual_scale.ndim == 0: + residual_scale = residual_scale.repeat(self._num_joints) + if residual_scale.shape != (self._num_joints,): + raise ValueError( + f"residual_scale must be scalar or contain {self._num_joints} values, " + f"got shape {tuple(residual_scale.shape)}." + ) + if not bool(torch.isfinite(residual_scale).all()) or bool(torch.any(residual_scale < 0.0)): + raise ValueError("residual_scale must contain finite nonnegative values.") + self._residual_scale = residual_scale + + self._raw_actions = torch.zeros((self.num_envs, self._num_joints + 1), device=self.device) + self._processed_actions = torch.zeros((self.num_envs, self._num_joints), device=self.device) + self._filtered_residual = torch.zeros_like(self._processed_actions) + self._reference_waypoints = torch.zeros( + (self.num_envs, self._waypoint_count, self._num_joints), + device=self.device, + ) + self._reference_phase = torch.zeros(self.num_envs, device=self.device) + self._minimum_phase = torch.zeros(self.num_envs, device=self.device) + self._grasp_dwell_count = torch.zeros(self.num_envs, device=self.device, dtype=torch.long) + self._approach_dwell_count = torch.zeros(self.num_envs, device=self.device, dtype=torch.long) + self._grasp_unlocked = torch.zeros(self.num_envs, device=self.device, dtype=torch.bool) + self._approach_unlocked = torch.zeros(self.num_envs, device=self.device, dtype=torch.bool) + self._lift_unlocked = torch.zeros(self.num_envs, device=self.device, dtype=torch.bool) + self._align_unlocked = torch.zeros(self.num_envs, device=self.device, dtype=torch.bool) + limits = self._asset.data.soft_joint_pos_limits.torch[:, self._joint_ids] + self._lower_limits = limits[..., 0] + self._upper_limits = limits[..., 1] + + @property + def action_dim(self) -> int: + return self._num_joints + 1 + + @property + def raw_actions(self) -> torch.Tensor: + return self._raw_actions + + @property + def processed_actions(self) -> torch.Tensor: + return self._processed_actions + + @property + def reference_phase(self) -> torch.Tensor: + """Current monotonic trajectory progress in ``[0, 1]``.""" + return self._reference_phase + + @property + def reference_target(self) -> torch.Tensor: + """Joint target on the reference trajectory before policy residuals [rad].""" + return self._interpolate_reference(self._reference_phase) + + @property + def reference_error(self) -> torch.Tensor: + """Applied joint-target error relative to the physical arm joints [rad].""" + arm_q = self._asset.data.joint_pos.torch[:, self._joint_ids] + return self._processed_actions - arm_q + + @property + def milestone_status(self) -> torch.Tensor: + """Observable task latches plus approach and grasp dwell progress.""" + approach_dwell = self._approach_dwell_count.float() / max(self._approach_dwell_steps, 1) + dwell = self._grasp_dwell_count.float() / max(self._grasp_dwell_steps, 1) + return torch.stack( + ( + self._approach_unlocked.float(), + self._grasp_unlocked.float(), + self._lift_unlocked.float(), + self._align_unlocked.float(), + torch.clamp(approach_dwell, 0.0, 1.0), + torch.clamp(dwell, 0.0, 1.0), + ), + dim=-1, + ) + + @property + def lift_unlocked(self) -> torch.Tensor: + """Whether each environment has demonstrated a held physical lift.""" + return self._lift_unlocked + + @property + def residual_scale(self) -> torch.Tensor: + """Joint displacement represented by one residual-action unit [rad].""" + return self._residual_scale + + def set_reference( + self, + waypoints: torch.Tensor, + phase: torch.Tensor, + initial_target: torch.Tensor, + env_ids: Sequence[int] | torch.Tensor | slice | None = None, + ) -> None: + """Set selected trajectory waypoints, progress, and filtered target. + + Args: + waypoints: Joint waypoints with shape ``(N, waypoint_count, num_joints)`` [rad]. + phase: Initial trajectory progress with shape ``(N,)``. + initial_target: Initial filtered joint targets with shape ``(N, num_joints)`` [rad]. + env_ids: Environments to update. If omitted, update every environment. + """ + selected = slice(None) if env_ids is None else env_ids + expected_waypoints = self._reference_waypoints[selected].shape + expected_phase = self._reference_phase[selected].shape + expected_target = self._processed_actions[selected].shape + if waypoints.shape != expected_waypoints: + raise ValueError(f"Waypoint shape {tuple(waypoints.shape)} does not match {tuple(expected_waypoints)}.") + if phase.shape != expected_phase: + raise ValueError(f"Phase shape {tuple(phase.shape)} does not match {tuple(expected_phase)}.") + if initial_target.shape != expected_target: + raise ValueError( + f"Initial target shape {tuple(initial_target.shape)} does not match {tuple(expected_target)}." + ) + if not bool(torch.isfinite(waypoints).all()): + raise ValueError("Reference waypoints must be finite.") + if not bool(torch.isfinite(phase).all()) or bool(torch.any((phase < 0.0) | (phase > 1.0))): + raise ValueError("Reference phase must be finite and lie in [0, 1].") + self._reference_waypoints[selected] = waypoints + self._reference_phase[selected] = phase + self._minimum_phase[selected] = phase + self._grasp_dwell_count[selected] = 0 + self._approach_dwell_count[selected] = 0 + self._approach_unlocked[selected] = phase > self._approach_phase + 1.0e-6 + # Gated stages must demonstrate fresh stable contact even when a supplied carry reset + # begins beyond the grasp waypoint; ungated pour-only resets may continue immediately. + gated_stage = self._env.curriculum_stage[selected] >= self._grasp_gate_stage + self._grasp_unlocked[selected] = (phase > self._grasp_phase + 1.0e-6) & ~gated_stage + self._lift_unlocked[selected] = (phase >= self._lift_phase) & ~gated_stage + self._align_unlocked[selected] = (phase >= self._align_phase) & ~gated_stage + self._processed_actions[selected] = initial_target + self._filtered_residual[selected] = 0.0 + + def _interpolate_reference(self, phase: torch.Tensor) -> torch.Tensor: + lower_index = torch.bucketize(phase.contiguous(), self._waypoint_phases[1:-1]) + lower_phase = self._waypoint_phases[lower_index] + upper_phase = self._waypoint_phases[lower_index + 1] + fraction = ((phase - lower_phase) / (upper_phase - lower_phase)).unsqueeze(-1) + # Smooth target velocity at every waypoint without changing the reference path. + fraction = fraction.square() * (3.0 - 2.0 * fraction) + env_ids = torch.arange(self.num_envs, device=self.device) + lower = self._reference_waypoints[env_ids, lower_index] + upper = self._reference_waypoints[env_ids, lower_index + 1] + return torch.lerp(lower, upper, fraction) + + def _grasp_ready(self) -> torch.Tensor: + held = self._held_ready(require_settled=True) + cup_velocity = self._env.cup_velocity_w() + return ( + held + & (torch.linalg.vector_norm(cup_velocity[:, :3], dim=-1) <= self._grasp_max_linear_velocity) + & (torch.linalg.vector_norm(cup_velocity[:, 3:], dim=-1) <= self._grasp_max_angular_velocity) + ) + + def _approach_ready(self) -> torch.Tensor: + _, cross_track_error = self._env.grasp_approach_error() + cup_velocity = self._env.cup_velocity_w() + return ( + (cross_track_error <= self._approach_max_lateral_distance) + & (torch.linalg.vector_norm(self.reference_error, dim=-1) <= self._approach_max_joint_error) + & (torch.linalg.vector_norm(cup_velocity[:, :3], dim=-1) <= self._approach_max_linear_velocity) + & (torch.linalg.vector_norm(cup_velocity[:, 3:], dim=-1) <= self._approach_max_angular_velocity) + ) + + def _held_ready(self, require_settled: bool = False) -> torch.Tensor: + """Return whether the hand still geometrically holds the preloaded cup.""" + tcp_distance = torch.linalg.vector_norm(self._env.tcp_pos_e() - self._env.cup_grasp_point_e(), dim=-1) + width_error = torch.abs(self._env.gripper_width() - float(self._env.gripper_grasp_width)) + gripper = self._env.action_manager.get_term("gripper_action") + contact = gripper.bilateral_preload if require_settled else gripper.bilateral_contact + return ( + (tcp_distance <= self._grasp_max_tcp_distance) + & (width_error <= float(self._env.cfg.success_max_gripper_width_error)) + & contact + & ( + gripper.commanded_position[:, 0] + <= float(self._env.cfg.gripper_preload_pos) + _GRIPPER_POSITION_TOLERANCE + ) + ) + + def _lift_ready(self) -> torch.Tensor: + return self._held_ready() & ( + self._env.cup_pose_e()[:, 2] - float(self._env.cup_reset_height) + >= float(self._env.cfg.success_min_lift_height) + ) + + def _align_ready(self) -> torch.Tensor: + source = self._env.cup_grasp_point_e()[:, :2] + target = self._env.target_pose_e()[:, :2] + desired = target + source.new_tensor(self._env.cfg.pour_source_offset_xy) + return self._lift_ready() & (torch.linalg.vector_norm(source - desired, dim=-1) <= self._align_max_distance) + + @staticmethod + def _phase_speed_command(phase_action: torch.Tensor) -> torch.Tensor: + """Map a residual phase action to a bounded multiplier around nominal speed.""" + return 1.0 + 0.25 * torch.clamp(phase_action, -1.0, 1.0) + + @staticmethod + def _monotonic_gate_limit(current_phase: torch.Tensor, gate_phase: float) -> torch.Tensor: + """Hold at a milestone without rewinding curriculum resets that start after it.""" + return torch.maximum(current_phase, torch.full_like(current_phase, float(gate_phase))) + + def process_actions(self, actions: torch.Tensor) -> None: + """Advance the reference phase and add bounded joint residuals.""" + self._raw_actions.copy_(actions) + # Zero action is the validated nominal trajectory. The policy only adjusts timing within a + # contact-safe range, matching the residual semantics of the seven joint coordinates. + phase_command = self._phase_speed_command(actions[:, 0]) + approaching = self._reference_phase < self._grasp_phase + transporting = self._grasp_unlocked & (self._reference_phase < self._align_phase) + phase_rate = torch.where( + approaching, + torch.full_like(self._reference_phase, self._approach_phase_rate), + torch.where( + transporting, + torch.full_like(self._reference_phase, self._transport_phase_rate), + torch.full_like(self._reference_phase, self._phase_rate), + ), + ) + proposed_phase = self._reference_phase + float(self._env.step_dt) * phase_rate * phase_command + proposed_phase = torch.maximum(proposed_phase, self._minimum_phase).clamp_max_(1.0) + gated_stage = self._env.curriculum_stage >= self._grasp_gate_stage + at_approach = self._reference_phase >= self._approach_phase - 1.0e-6 + approach_dwell = gated_stage & at_approach & self._approach_ready() & ~self._approach_unlocked + self._approach_dwell_count.copy_( + torch.where( + approach_dwell, + self._approach_dwell_count + 1, + torch.zeros_like(self._approach_dwell_count), + ) + ) + self._approach_unlocked |= self._approach_dwell_count >= self._approach_dwell_steps + at_grasp = self._reference_phase >= self._grasp_phase - 1.0e-6 + grasp_ready = self._grasp_ready() + dwell_active = gated_stage & at_grasp & grasp_ready & ~self._grasp_unlocked + self._grasp_dwell_count.copy_( + torch.where(dwell_active, self._grasp_dwell_count + 1, torch.zeros_like(self._grasp_dwell_count)) + ) + self._grasp_unlocked |= self._grasp_dwell_count >= self._grasp_dwell_steps + self._lift_unlocked |= self._grasp_unlocked & self._lift_ready() + self._align_unlocked |= self._lift_unlocked & self._align_ready() + + advancing = proposed_phase > self._reference_phase + approach_gate = advancing & gated_stage & ~self._approach_unlocked & (proposed_phase > self._approach_phase) + approach_limit = self._monotonic_gate_limit(self._reference_phase, self._approach_phase) + proposed_phase = torch.where( + approach_gate, + approach_limit, + proposed_phase, + ) + grasp_gate = advancing & gated_stage & ~self._grasp_unlocked & (proposed_phase > self._grasp_phase) + grasp_limit = self._monotonic_gate_limit(self._reference_phase, self._grasp_phase) + proposed_phase = torch.where(grasp_gate, grasp_limit, proposed_phase) + lift_gate = advancing & ~self._lift_unlocked & (proposed_phase > self._lift_phase) + lift_limit = torch.maximum(self._reference_phase, torch.full_like(proposed_phase, self._lift_phase)) + proposed_phase = torch.where(lift_gate, lift_limit, proposed_phase) + align_gate = advancing & ~self._align_unlocked & (proposed_phase > self._align_phase) + align_limit = torch.maximum(self._reference_phase, torch.full_like(proposed_phase, self._align_phase)) + proposed_phase = torch.where(align_gate, align_limit, proposed_phase) + self._reference_phase.copy_(proposed_phase) + + residual = torch.tanh(actions[:, 1:]) * self._residual_scale + # Hold the validated approach and preload geometry stationary while bilateral finger + # contact settles. Residual control resumes immediately after the grasp latch. + pending_preload = gated_stage & at_approach & ~self._grasp_unlocked + residual = torch.where(pending_preload.unsqueeze(-1), torch.zeros_like(residual), residual) + self._filtered_residual.lerp_(residual, self._alpha) + target = self._interpolate_reference(self._reference_phase) + self._filtered_residual + target = torch.clamp(target, min=self._lower_limits, max=self._upper_limits) + self._processed_actions.copy_(target) + + def apply_actions(self) -> None: + self._asset.set_joint_position_target_index(target=self._processed_actions, joint_ids=self._joint_ids) + + def reset(self, env_ids: Sequence[int] | torch.Tensor | slice | None = None) -> None: + """Clear selected raw residual commands while preserving the reset reference.""" + selected = slice(None) if env_ids is None else env_ids + self._raw_actions[selected] = 0.0 + + +@configclass +class TrajectoryJointPositionActionCfg(ActionTermCfg): + """Configuration for :class:`TrajectoryJointPositionAction`.""" + + joint_names: list[str] = MISSING + preserve_order: bool = False + waypoint_count: int = 6 + residual_scale: float | tuple[float, ...] = 0.05 + alpha: float = 0.2 + """Exponential smoothing weight applied only to policy joint residuals.""" + phase_rate: float = 1.0 / 3.5 + approach_phase_rate: float = 1.0 / 3.5 + """Phase rate before the grasp waypoint, kept slower for a contact-safe approach [1/s].""" + transport_phase_rate: float = 0.5 + """Phase rate after a validated grasp and before the receiver-aligned pour [1/s].""" + waypoint_phases: tuple[float, ...] = (0.0, 0.12, 0.24, 0.40, 0.62, 1.0) + approach_waypoint: int = 1 + grasp_waypoint: int = 2 + lift_waypoint: int = 3 + align_waypoint: int = 4 + grasp_gate_stage: int = 3 + approach_max_lateral_distance: float = 0.01 + """Maximum cross-track TCP error perpendicular to the grasp-approach axis [m].""" + approach_max_joint_error: float = 0.08 + approach_dwell_steps: int = 10 + """Consecutive centered, stationary steps required before the guarded grasp approach.""" + approach_max_linear_velocity: float = 0.01 + approach_max_angular_velocity: float = 0.1 + align_max_distance: float = 0.06 + """Maximum source-grasp-point error from the receiver-side pour pose [m].""" + grasp_dwell_steps: int = 15 + grasp_max_tcp_distance: float = 0.01 + grasp_max_linear_velocity: float = 0.05 + grasp_max_angular_velocity: float = 0.5 + class_type: type[TrajectoryJointPositionAction] = TrajectoryJointPositionAction + + +class CurriculumGripperPositionAction(ActionTerm): + """Filtered symmetric finger-position command with residual, incremental, and binary modes.""" + + cfg: CurriculumGripperPositionActionCfg + + def __init__(self, cfg: CurriculumGripperPositionActionCfg, env) -> None: + super().__init__(cfg, env) + self._joint_ids, self._joint_names = self._asset.find_joints(cfg.joint_names, preserve_order=True) + self._num_joints = len(self._joint_ids) + if self._num_joints == 0: + raise ValueError("CurriculumGripperPositionAction resolved no joints.") + + self._scale = float(cfg.scale) + self._alpha = float(cfg.alpha) + if not isinstance(cfg.use_incremental_target, bool): + raise TypeError("use_incremental_target must be a bool.") + self._use_incremental_target = cfg.use_incremental_target + self._binary_threshold = cfg.binary_threshold + self._close_position = float(cfg.close_position) + self._neutral_position = float(cfg.neutral_position) + self._open_position = float(cfg.open_position) + self._default_position = self._close_position if cfg.default_position is None else float(cfg.default_position) + self._contact_min_deflection = float(cfg.contact_min_deflection) + self._contact_max_velocity = float(cfg.contact_max_velocity) + self._force_open_stage = int(cfg.force_open_before_phase_stage) + self._force_open_phase = float(cfg.force_open_before_phase) + self._capture_max_lateral_distance = float(cfg.capture_max_lateral_distance) + self._capture_max_vertical_distance = float(cfg.capture_max_vertical_distance) + self._capture_max_joint_error = float(cfg.capture_max_joint_error) + self._capture_dwell_steps = int(cfg.capture_dwell_steps) + self._capture_max_linear_velocity = float(cfg.capture_max_linear_velocity) + self._capture_max_angular_velocity = float(cfg.capture_max_angular_velocity) + if not math.isfinite(self._scale) or self._scale <= 0.0: + raise ValueError("Curriculum gripper action scale must be finite and positive.") + if not 0.0 < self._alpha <= 1.0: + raise ValueError(f"Moving-average weight must lie in (0, 1], got {self._alpha}.") + if self._binary_threshold is not None: + if ( + isinstance(self._binary_threshold, bool) + or not math.isfinite(self._binary_threshold) + or not -1.0 < self._binary_threshold < 1.0 + ): + raise ValueError("Binary gripper threshold must be finite and lie strictly between -1 and 1.") + if self._use_incremental_target: + raise ValueError("Binary and incremental gripper targets are mutually exclusive.") + if ( + not math.isfinite(self._close_position) + or not math.isfinite(self._neutral_position) + or not math.isfinite(self._open_position) + or not self._close_position <= self._neutral_position <= self._open_position + ): + raise ValueError( + "Curriculum gripper positions must be finite with close_position <= neutral_position <= open_position." + ) + if not math.isfinite(self._default_position) or not ( + self._close_position <= self._default_position <= self._neutral_position + ): + raise ValueError("default_position must lie in [close_position, neutral_position].") + if self._force_open_stage < -1: + raise ValueError("force_open_before_phase_stage must be at least -1.") + if not 0.0 <= self._force_open_phase <= 1.0: + raise ValueError("force_open_before_phase must lie in [0, 1].") + if not all( + math.isfinite(value) and value > 0.0 + for value in (self._capture_max_lateral_distance, self._capture_max_vertical_distance) + ): + raise ValueError("Capture position limits must be finite and positive.") + if not math.isfinite(self._capture_max_joint_error) or self._capture_max_joint_error <= 0.0: + raise ValueError("capture_max_joint_error must be finite and positive.") + if self._capture_dwell_steps <= 0: + raise ValueError("capture_dwell_steps must be positive.") + if not all( + math.isfinite(value) and value > 0.0 + for value in (self._capture_max_linear_velocity, self._capture_max_angular_velocity) + ): + raise ValueError("Capture velocity limits must be finite and positive.") + if not math.isfinite(self._contact_min_deflection) or self._contact_min_deflection <= 0.0: + raise ValueError("contact_min_deflection must be finite and positive.") + if not math.isfinite(self._contact_max_velocity) or self._contact_max_velocity <= 0.0: + raise ValueError("contact_max_velocity must be finite and positive.") + self._raw_actions = torch.zeros((self.num_envs, 1), device=self.device) + self._action_offset = torch.full( + (self.num_envs, 1), + self._default_position, + device=self.device, + ) + self._processed_actions = self._action_offset.expand(-1, self._num_joints).clone() + self._capture_unlocked = torch.ones(self.num_envs, device=self.device, dtype=torch.bool) + self._capture_dwell_count = torch.zeros(self.num_envs, device=self.device, dtype=torch.long) + + @property + def action_dim(self) -> int: + return 1 + + @property + def raw_actions(self) -> torch.Tensor: + return self._raw_actions + + @property + def processed_actions(self) -> torch.Tensor: + return self._processed_actions + + @property + def action_offset(self) -> torch.Tensor: + """Per-environment residual-mode target and initialization position [m].""" + return self._action_offset + + @property + def commanded_position(self) -> torch.Tensor: + """Current symmetric per-finger position target [m].""" + return self._processed_actions[:, :1] + + @property + def contact_deflection(self) -> torch.Tensor: + """Per-finger position-drive deflection caused by contact [m].""" + joint_position = self._asset.data.joint_pos.torch[:, self._joint_ids] + joint_velocity = self._asset.data.joint_vel.torch[:, self._joint_ids] + deflection, _ = _bilateral_gripper_preload( + joint_position, + joint_velocity, + self._processed_actions, + min_deflection=self._contact_min_deflection, + max_velocity=self._contact_max_velocity, + max_command=self._neutral_position, + ) + return deflection + + @property + def bilateral_preload(self) -> torch.Tensor: + """Whether both fingers have settled against the commanded preload.""" + joint_position = self._asset.data.joint_pos.torch[:, self._joint_ids] + joint_velocity = self._asset.data.joint_vel.torch[:, self._joint_ids] + _, bilateral = _bilateral_gripper_preload( + joint_position, + joint_velocity, + self._processed_actions, + min_deflection=self._contact_min_deflection, + max_velocity=self._contact_max_velocity, + max_command=self._neutral_position, + ) + return bilateral + + @property + def bilateral_contact(self) -> torch.Tensor: + """Whether both fingers remain deflected against the commanded cup contact.""" + joint_position = self._asset.data.joint_pos.torch[:, self._joint_ids] + joint_velocity = self._asset.data.joint_vel.torch[:, self._joint_ids] + deflection, _ = _bilateral_gripper_preload( + joint_position, + joint_velocity, + self._processed_actions, + min_deflection=self._contact_min_deflection, + max_velocity=self._contact_max_velocity, + max_command=self._neutral_position, + ) + finite = torch.isfinite(joint_position).all(dim=-1) & torch.isfinite(self._processed_actions).all(dim=-1) + command_valid = self._processed_actions.amax(dim=-1) <= (self._neutral_position + _GRIPPER_POSITION_TOLERANCE) + return finite & command_valid & (deflection.amin(dim=-1) >= self._contact_min_deflection) + + @property + def contact_quality(self) -> torch.Tensor: + """Smooth bilateral-contact quality in ``[0, 1]``.""" + joint_velocity = self._asset.data.joint_vel.torch[:, self._joint_ids] + deflection = self.contact_deflection + deflection_quality = torch.clamp(deflection.amin(dim=-1) / self._contact_min_deflection, 0.0, 1.0) + velocity_quality = torch.clamp( + 1.0 - joint_velocity.abs().amax(dim=-1) / self._contact_max_velocity, + 0.0, + 1.0, + ) + command_valid = self._processed_actions.amax(dim=-1) <= (self._neutral_position + _GRIPPER_POSITION_TOLERANCE) + return deflection_quality * velocity_quality * command_valid.float() + + @property + def capture_status(self) -> torch.Tensor: + """Observable gripper-capture latch and dwell progress.""" + dwell = self._capture_dwell_count.float() / max(self._capture_dwell_steps, 1) + return torch.stack((self._capture_unlocked.float(), torch.clamp(dwell, 0.0, 1.0)), dim=-1) + + @property + def IO_descriptor(self): + descriptor = super().IO_descriptor + descriptor.shape = (1,) + descriptor.dtype = str(self.raw_actions.dtype) + descriptor.action_type = "JointAction" + descriptor.joint_names = self._joint_names + descriptor.scale = self._scale + return descriptor + + def set_action_offset( + self, + offset: torch.Tensor, + env_ids: Sequence[int] | torch.Tensor | slice | None = None, + ) -> None: + """Set selected environments' residual-mode target and initialization position [m].""" + selected = slice(None) if env_ids is None else env_ids + expected_shape = self._action_offset[selected].shape + if offset.shape != expected_shape: + raise ValueError(f"Action offset shape {tuple(offset.shape)} does not match {tuple(expected_shape)}.") + offset = offset.to(device=self._action_offset.device, dtype=self._action_offset.dtype) + self._action_offset[selected] = offset + self._processed_actions[selected] = offset.expand(-1, self._num_joints) + + def set_reset_position( + self, + position: torch.Tensor, + env_ids: Sequence[int] | torch.Tensor | slice | None = None, + ) -> None: + """Align the filtered target with selected physical reset positions [m].""" + selected = slice(None) if env_ids is None else env_ids + expected_shape = self._action_offset[selected].shape + if position.shape != expected_shape: + raise ValueError(f"Reset-position shape {tuple(position.shape)} does not match {tuple(expected_shape)}.") + position = position.to(device=self.device, dtype=self._processed_actions.dtype) + expanded = position.expand(-1, self._num_joints) + self._processed_actions[selected] = expanded + + def process_actions(self, actions: torch.Tensor) -> None: + self._raw_actions.copy_(actions) + bounded_actions = torch.clamp(actions, -1.0, 1.0) + if self._binary_threshold is None: + action_base = self._processed_actions[:, :1] if self._use_incremental_target else self._action_offset + target = torch.clamp( + action_base + self._scale * bounded_actions, + min=self._close_position, + max=self._neutral_position, + ) + else: + target = torch.where( + bounded_actions < self._binary_threshold, + torch.full_like(bounded_actions, self._close_position), + torch.full_like(bounded_actions, self._neutral_position), + ) + if self._force_open_stage >= 0: + arm_action = self._env.action_manager.get_term("arm_action") + joint_error = torch.linalg.vector_norm(arm_action.reference_error, dim=-1) + axial_error, cross_track_error = self._env.grasp_approach_error() + cup_velocity = self._env.cup_velocity_w() + capture_required = self._env.curriculum_stage >= self._force_open_stage + capture_ready = ( + (arm_action.reference_phase >= self._force_open_phase) + & (cross_track_error <= self._capture_max_lateral_distance) + & (torch.abs(axial_error) <= self._capture_max_vertical_distance) + & (joint_error <= self._capture_max_joint_error) + & (torch.linalg.vector_norm(cup_velocity[:, :3], dim=-1) <= self._capture_max_linear_velocity) + & (torch.linalg.vector_norm(cup_velocity[:, 3:], dim=-1) <= self._capture_max_angular_velocity) + ) + dwell_active = capture_required & capture_ready & ~self._capture_unlocked + self._capture_dwell_count.copy_( + torch.where(dwell_active, self._capture_dwell_count + 1, torch.zeros_like(self._capture_dwell_count)) + ) + self._capture_unlocked |= self._capture_dwell_count >= self._capture_dwell_steps + force_open = capture_required & ~self._capture_unlocked + target = torch.where(force_open.unsqueeze(-1), torch.full_like(target, self._open_position), target) + self._processed_actions.lerp_(target.expand(-1, self._num_joints), self._alpha) + + def apply_actions(self) -> None: + self._asset.set_joint_position_target_index(target=self._processed_actions, joint_ids=self._joint_ids) + + def reset(self, env_ids: Sequence[int] | torch.Tensor | slice | None = None) -> None: + selected = slice(None) if env_ids is None else env_ids + self._raw_actions[selected] = 0.0 + if self._force_open_stage < 0: + self._capture_unlocked[selected] = True + else: + self._capture_unlocked[selected] = self._env.curriculum_stage[selected] < self._force_open_stage + self._capture_dwell_count[selected] = 0 + + +@configclass +class CurriculumGripperPositionActionCfg(ActionTermCfg): + """Configuration for :class:`CurriculumGripperPositionAction`.""" + + joint_names: list[str] = MISSING + scale: float = 0.04 + """Per-finger residual or incremental delta per policy-action unit [m]; unused in binary mode.""" + alpha: float = 0.2 + """Interpolation weight applied to the selected finger target.""" + use_incremental_target: bool = False + """Whether actions increment the target by ``alpha * scale`` so zero action holds its position.""" + binary_threshold: float | None = None + """Optional threshold selecting filtered close/maximum targets; values below it close.""" + close_position: float = 0.0 + neutral_position: float = 0.025 + """Largest per-finger command accepted from the action [m].""" + open_position: float = 0.04 + default_position: float | None = None + """Per-finger residual-mode zero command and initial target [m]. ``None`` uses ``close_position``.""" + limit_to_preload: bool = True + """Whether task validation restricts action targets to the contact-safe preload interval.""" + contact_min_deflection: float = 0.001 + """Minimum settled position-drive deflection required on each finger [m].""" + contact_max_velocity: float = 0.005 + """Maximum absolute finger speed accepted as settled bilateral contact [m/s].""" + force_open_before_phase_stage: int = -1 + """First stage that holds the hand open during the approach phase, or ``-1`` to disable.""" + force_open_before_phase: float = 0.25 + """Reference phase below which configured approach stages force the hand open.""" + capture_max_lateral_distance: float = 0.005 + """Maximum cross-track TCP error perpendicular to the grasp-approach axis [m].""" + capture_max_vertical_distance: float = 0.008 + """Maximum absolute TCP error along the grasp-approach axis [m]. + + The field retains its historical name for configuration compatibility. + """ + capture_max_joint_error: float = 0.08 + """Maximum reference-to-physical arm-joint error that releases the interlock [rad].""" + capture_dwell_steps: int = 5 + """Consecutive centered, stationary steps required before finger closure is enabled.""" + capture_max_linear_velocity: float = 0.02 + """Maximum source-cup linear speed during capture qualification [m/s].""" + capture_max_angular_velocity: float = 0.2 + """Maximum source-cup angular speed during capture qualification [rad/s].""" + class_type: type[CurriculumGripperPositionAction] = CurriculumGripperPositionAction diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/curriculums.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/curriculums.py new file mode 100644 index 000000000000..842fe533d860 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/curriculums.py @@ -0,0 +1,181 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Success-driven curriculum for the Franka Pour task.""" + +from __future__ import annotations + +import math +from collections import deque +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import torch + +from isaaclab.managers import CurriculumTermCfg +from isaaclab.managers.manager_base import ManagerTermBase + +if TYPE_CHECKING: + from ..pour_env import FrankaPourEnv + + +class PourCurriculum(ManagerTermBase): + """Advance one shared stage-and-randomization frontier from completed episodes.""" + + def __init__(self, cfg: CurriculumTermCfg, env: FrankaPourEnv): + super().__init__(cfg, env) + stage_count = len(env.cfg.curriculum_stage_names) + if stage_count == 0 or len(env.cfg.curriculum_target_frac) != stage_count: + raise ValueError("Curriculum stage names and target fractions must have equal nonzero length.") + + self.max_stage = stage_count - 1 + self.stage = int(env.cfg.curriculum_start_stage) + if self.stage < 0 or self.stage > self.max_stage: + raise ValueError(f"curriculum_start_stage must lie in [0, {self.max_stage}].") + randomization_levels = env.cfg.curriculum_randomization_extent_levels + if not randomization_levels: + raise ValueError("curriculum_randomization_extent_levels must not be empty.") + self.max_randomization_level = len(randomization_levels) - 1 + self.randomization_level = int(env.cfg.curriculum_randomization_start_level) + if self.randomization_level < 0 or self.randomization_level > self.max_randomization_level: + raise ValueError(f"curriculum_randomization_start_level must lie in [0, {self.max_randomization_level}].") + + window_size = self._success_window_size(env) + self._minimum_completed_episodes(env, window_size) + self._success_window: deque[bool] = deque() + self._success_count = 0 + self.success_rate = 0.0 + self.resets_in_stage = 0 + env.set_curriculum_stage(slice(None), self.stage) + env.set_curriculum_randomization_level(slice(None), self.randomization_level) + + def __call__( + self, + env: FrankaPourEnv, + env_ids: Sequence[int] | torch.Tensor | slice, + ) -> dict[str, float]: + """Update stage statistics and assign the next episodes.""" + ids = self._as_tensor(env, env_ids) + window_size = self._success_window_size(env) + minimum_completed_episodes = self._minimum_completed_episodes(env, window_size) + self._trim_success_window(window_size) + mean_peak_target_frac: float | None = None + + if ids.numel() > 0: + # Count only completed episodes from the active frontier. This excludes the initial + # reset, replay episodes, and asynchronous old-stage episodes after a promotion. + completed = (env.episode_length_buf[ids] > 0) & (env.curriculum_stage[ids] == self.stage) + if self.stage == self.max_stage: + completed &= env.curriculum_randomization_level[ids] == self.randomization_level + completed_ids = ids[completed] + count = int(completed_ids.numel()) + if count > 0: + outcomes = env.episode_succeeded[completed_ids].detach().to(device="cpu", dtype=torch.bool).tolist() + self._update_success_window(outcomes, window_size) + self.resets_in_stage += count + mean_peak_target_frac = float(env.ep_max_target_frac[completed_ids].mean().item()) + + promotion_threshold = float(env.cfg.curriculum_success_threshold) + if self.stage == self.max_stage and self.randomization_level < self.max_randomization_level: + promotion_threshold = float(env.cfg.curriculum_randomization_promotion_threshold) + frontier_mastered = ( + not env.cfg.curriculum_freeze + and self.resets_in_stage >= minimum_completed_episodes + and self.success_rate >= promotion_threshold + ) + if frontier_mastered: + if self.stage < self.max_stage: + self.stage += 1 + self._clear_stage_statistics() + elif self.randomization_level < self.max_randomization_level: + self.randomization_level += 1 + self._clear_stage_statistics() + + env.set_curriculum_stage(ids, self.stage) + env.set_curriculum_randomization_level(ids, self.randomization_level) + replay_fraction = self._previous_frontier_replay_fraction(env, minimum_completed_episodes) + if not env.cfg.curriculum_freeze and self.stage > 0 and replay_fraction > 0.0: + replay = torch.rand(ids.numel(), device=env.device) < replay_fraction + if self.stage == self.max_stage and self.randomization_level > 0: + env.set_curriculum_randomization_level(ids[replay], self.randomization_level - 1) + else: + env.set_curriculum_stage(ids[replay], self.stage - 1) + + mastered = ( + self.stage == self.max_stage + and self.randomization_level == self.max_randomization_level + and self.resets_in_stage >= minimum_completed_episodes + and self.success_rate >= env.cfg.curriculum_success_threshold + ) + metrics = { + "stage": float(self.stage), + "randomization_level": float(self.randomization_level), + "success_rate": float(self.success_rate), + "completed_episodes": float(self.resets_in_stage), + "required_completed_episodes": float(minimum_completed_episodes), + "mastered": float(mastered), + } + if mean_peak_target_frac is not None: + metrics["mean_peak_target_frac"] = mean_peak_target_frac + return metrics + + def _previous_frontier_replay_fraction(self, env: FrankaPourEnv, minimum_completed_episodes: int) -> float: + """Return predecessor sampling probability at the current frontier.""" + if env.cfg.curriculum_freeze or self.stage == 0: + return 0.0 + retention = float(env.cfg.curriculum_previous_stage_replay_fraction) + entry = float(env.cfg.curriculum_frontier_entry_replay_fraction) + evidence_fraction = min(float(self.resets_in_stage) / float(minimum_completed_episodes), 1.0) + return entry + evidence_fraction * (retention - entry) + + @staticmethod + def _success_window_size(env: FrankaPourEnv) -> int: + """Return the configured number of recent frontier episodes used for promotion.""" + window_size = int(env.cfg.curriculum_min_resets_per_stage) + if window_size <= 0: + raise ValueError("curriculum_min_resets_per_stage must be positive.") + return window_size + + @staticmethod + def _minimum_completed_episodes(env: FrankaPourEnv, window_size: int) -> int: + """Return the exposure floor before promotion, scaled by vectorized environment count.""" + cohort_count = float(env.cfg.curriculum_min_reset_cohorts_per_stage) + if not math.isfinite(cohort_count) or cohort_count < 0.0: + raise ValueError("curriculum_min_reset_cohorts_per_stage must be finite and nonnegative.") + cohort_episodes = math.ceil(cohort_count * int(env.num_envs)) + return max(window_size, cohort_episodes) + + def _update_success_window(self, outcomes: list[bool], window_size: int) -> None: + """Append episode outcomes and update the exact recent-window success rate.""" + for outcome in outcomes: + if len(self._success_window) == window_size: + self._success_count -= int(self._success_window.popleft()) + outcome = bool(outcome) + self._success_window.append(outcome) + self._success_count += int(outcome) + self.success_rate = self._success_count / len(self._success_window) + + def _trim_success_window(self, window_size: int) -> None: + """Apply a runtime window-size change without retaining stale outcomes.""" + while len(self._success_window) > window_size: + self._success_count -= int(self._success_window.popleft()) + self.success_rate = self._success_count / len(self._success_window) if self._success_window else 0.0 + + def _clear_stage_statistics(self) -> None: + """Clear frontier evidence after promotion to a new stage.""" + self._success_window.clear() + self._success_count = 0 + self.success_rate = 0.0 + self.resets_in_stage = 0 + + @staticmethod + def _as_tensor( + env: FrankaPourEnv, + env_ids: Sequence[int] | torch.Tensor | slice, + ) -> torch.Tensor: + """Normalize manager environment indices to a one-dimensional device tensor.""" + if isinstance(env_ids, slice): + return torch.arange(env.num_envs, device=env.device, dtype=torch.long)[env_ids] + return torch.as_tensor(env_ids, device=env.device, dtype=torch.long).flatten() diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/events.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/events.py new file mode 100644 index 000000000000..3f81b41a2e14 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/events.py @@ -0,0 +1,22 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Event (reset) terms for the two-cup pour task.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from ..pour_env import FrankaPourEnv + + +def reset_pour_scene(env: FrankaPourEnv, env_ids: torch.Tensor) -> None: + """Reset the arm to home and refill the source cup with media for ``env_ids``.""" + if not isinstance(env_ids, torch.Tensor): + env_ids = torch.as_tensor(list(env_ids), device=env.device, dtype=torch.long) + env.reset_pour_scene(env_ids.long()) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/observations.py new file mode 100644 index 000000000000..21406c036d39 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/observations.py @@ -0,0 +1,206 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Observation terms for the physical-grasp, two-cup Franka pour task.""" + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING + +import torch + +import isaaclab.utils.math as math_utils + +if TYPE_CHECKING: + from ..pour_env import FrankaPourEnv + + +def _canonical_pose(pose: torch.Tensor) -> torch.Tensor: + """Return a finite pose with the quaternion mapped to its unique hemisphere.""" + pose = torch.nan_to_num(pose) + return torch.cat((pose[:, :3], math_utils.quat_unique(pose[:, 3:7])), dim=-1) + + +def tcp_pose_obs(env: FrankaPourEnv) -> torch.Tensor: + """Tool-centre pose in the robot-base frame with a canonical XYZW quaternion.""" + return _canonical_pose(env.tcp_pose_e()) + + +def ee_pose_obs(env: FrankaPourEnv) -> torch.Tensor: + """Backward-compatible hand-pose observation used by older diagnostic scripts.""" + return torch.nan_to_num(env.ee_pose_e()) + + +def cup_pose_obs(env: FrankaPourEnv) -> torch.Tensor: + """Source-cup pose in the robot-base frame with a canonical XYZW quaternion.""" + return _canonical_pose(env.cup_pose_e()) + + +def target_pose_obs(env: FrankaPourEnv) -> torch.Tensor: + return _canonical_pose(env.target_pose_e()) + + +def tcp_to_grasp_position_c_obs(env: FrankaPourEnv) -> torch.Tensor: + """Desired grasp point minus TCP position, expressed in the source-cup frame [m].""" + cup_pose = env.cup_pose_e() + tcp_position_c = math_utils.quat_apply_inverse( + cup_pose[:, 3:7], + env.tcp_pos_e() - cup_pose[:, :3], + ) + desired_position_c = torch.zeros_like(tcp_position_c) + desired_position_c[:, 2] = float(env.cfg.cup_grasp_height) + return torch.nan_to_num(desired_position_c - tcp_position_c) + + +def grasp_to_tcp_quat_obs(env: FrankaPourEnv) -> torch.Tensor: + """TCP orientation relative to the desired source-cup grasp frame as a canonical quaternion.""" + cup_quat = env.cup_pose_e()[:, 3:7] + desired_quat = math_utils.quat_mul(cup_quat, env.desired_grasp_tcp_quat_c()) + tcp_quat = env.tcp_pose_e()[:, 3:7] + error_quat = math_utils.quat_mul(math_utils.quat_conjugate(desired_quat), tcp_quat) + return torch.nan_to_num(math_utils.quat_unique(error_quat)) + + +def target_position_c_obs(env: FrankaPourEnv) -> torch.Tensor: + """Receiving-cup center expressed in the source-cup frame [m].""" + cup_pose = env.cup_pose_e() + target_offset = env.target_pose_e()[:, :3] - cup_pose[:, :3] + return torch.nan_to_num(math_utils.quat_apply_inverse(cup_pose[:, 3:7], target_offset)) + + +def tcp_to_grasp_obs(env: FrankaPourEnv) -> torch.Tensor: + return torch.nan_to_num(env.cup_grasp_point_e() - env.tcp_pos_e()) + + +def cup_to_target_obs(env: FrankaPourEnv) -> torch.Tensor: + return torch.nan_to_num(env.target_pose_e()[:, :3] - env.cup_pose_e()[:, :3]) + + +def gripper_width_obs(env: FrankaPourEnv) -> torch.Tensor: + return torch.nan_to_num(env.gripper_width()).unsqueeze(-1) + + +def finger_position_obs(env: FrankaPourEnv) -> torch.Tensor: + """Individual finger-joint positions [m].""" + return torch.nan_to_num(env.finger_joint_pos()) + + +def finger_velocity_obs(env: FrankaPourEnv) -> torch.Tensor: + """Individual finger-joint velocities [m/s].""" + return torch.nan_to_num(env.finger_joint_vel()) + + +def arm_reference_phase_obs(env: FrankaPourEnv) -> torch.Tensor: + """Monotonic arm-reference progress in ``[0, 1]``.""" + action = env.action_manager.get_term("arm_action") + phase = getattr(action, "reference_phase", None) + if phase is None: + phase = torch.zeros(env.num_envs, device=env.device) + return torch.nan_to_num(phase).unsqueeze(-1) + + +def arm_reference_error_obs(env: FrankaPourEnv) -> torch.Tensor: + """Current applied trajectory target minus measured arm position [rad]. + + The action term low-pass filters policy residuals, so its applied target contains persistent + controller state that cannot be reconstructed from only the latest raw action. Exposing that + target error keeps the policy observation Markov under randomized reset-bank references. + """ + action = env.action_manager.get_term("arm_action") + error = getattr(action, "reference_error", None) + if error is None: + error = action.processed_actions - env.arm_joint_pos() + return torch.nan_to_num(error) + + +def trajectory_status_obs(env: FrankaPourEnv) -> torch.Tensor: + """Controller milestone, dwell, and capture state that affects future transitions.""" + arm = env.action_manager.get_term("arm_action") + gripper = env.action_manager.get_term("gripper_action") + arm_status = getattr(arm, "milestone_status", torch.zeros((env.num_envs, 6), device=env.device)) + capture_status = getattr(gripper, "capture_status", torch.ones((env.num_envs, 2), device=env.device)) + return torch.nan_to_num(torch.cat((arm_status, capture_status), dim=-1)).clamp_(0.0, 1.0) + + +def success_dwell_obs(env: FrankaPourEnv) -> torch.Tensor: + """Fraction of the stable-success dwell already satisfied.""" + dwell_steps = max(1, math.ceil(float(env.cfg.success_dwell_time_s) / max(float(env.step_dt), 1.0e-6))) + return torch.clamp(env._success_dwell_count.float() / dwell_steps, 0.0, 1.0).unsqueeze(-1) + + +def time_remaining_obs(env: FrankaPourEnv) -> torch.Tensor: + """Normalized finite-horizon time remaining in ``[0, 1]``.""" + progress = env.episode_length_buf.float() / max(int(env.max_episode_length), 1) + return torch.clamp(1.0 - progress, 0.0, 1.0).unsqueeze(-1) + + +def lost_grasp_dwell_obs(env: FrankaPourEnv) -> torch.Tensor: + """Fraction of the consecutive post-lift grasp-loss dwell already accumulated.""" + dwell_steps = max( + 1, + math.ceil(float(env.cfg.lost_grasp_dwell_time_s) / max(float(env.step_dt), 1.0e-6)), + ) + return torch.clamp(env._lost_grasp_dwell_count.float() / dwell_steps, 0.0, 1.0).unsqueeze(-1) + + +def pour_target_fraction_obs(env: FrankaPourEnv) -> torch.Tensor: + """Current episode's required held-delivery fraction in ``[0, 1]``.""" + return torch.nan_to_num(env.pour_target_frac).clamp_(0.0, 1.0).unsqueeze(-1) + + +def gripper_target_obs(env: FrankaPourEnv) -> torch.Tensor: + """Filtered symmetric per-finger target [m].""" + return torch.nan_to_num(env.action_manager.get_term("gripper_action").commanded_position) + + +def gripper_contact_obs(env: FrankaPourEnv) -> torch.Tensor: + """Per-finger position-drive deflection caused by contact [m].""" + gripper = env.action_manager.get_term("gripper_action") + deflection = getattr(gripper, "contact_deflection", torch.zeros((env.num_envs, 2), device=env.device)) + return torch.nan_to_num(deflection) + + +def cup_velocity_obs(env: FrankaPourEnv) -> torch.Tensor: + """Source-cup world-frame linear and angular velocity [m/s, rad/s].""" + return torch.nan_to_num(env.cup_velocity_w()) + + +def particle_fractions_obs(env: FrankaPourEnv) -> torch.Tensor: + """Source, target, spill, and held-qualified target fractions.""" + scale = max(env.num_particles, 1) + source = env.count_in_source() / scale + target = env.count_in_target() / scale + spilled = env.count_spilled() / scale + held_target = env.current_held_delivered_mask().sum(dim=1).float() / scale + return torch.stack((source, target, spilled, held_target), dim=-1) + + +def particle_transfer_obs(env: FrankaPourEnv) -> torch.Tensor: + """Airborne-media centroid/velocity relative to the receiver plus airborne fraction. + + Counts alone cannot tell the privileged critic whether a stream is moving toward or past the + receiver. This compact summary supplies that value-estimation signal without exposing + per-particle state to the actor. + """ + source, target, spilled = env.particle_region_masks() + airborne = ~(source | target | spilled) + weight = airborne.unsqueeze(-1).to(dtype=torch.float32) + count = weight.sum(dim=1) + denominator = count.clamp_min(1.0) + centroid = (env.particle_pos_e() * weight).sum(dim=1) / denominator + velocity = (env.particle_vel_e() * weight).sum(dim=1) / denominator + receiver = env.target_pose_e()[:, :3] + centroid_relative = torch.where(count > 0.0, centroid - receiver, torch.zeros_like(centroid)) + velocity = torch.where(count > 0.0, velocity, torch.zeros_like(velocity)) + fraction = count / max(env.num_particles, 1) + summary = torch.cat((centroid_relative / 0.30, velocity / 2.0, fraction), dim=-1) + return torch.nan_to_num(summary).clamp_(-2.0, 2.0) + + +def held_delivery_history_obs(env: FrankaPourEnv) -> torch.Tensor: + """Fraction of particles that have entered the receiver during a qualified held pour.""" + fraction = env.held_delivered_mask().sum(dim=1).float() / max(env.num_particles, 1) + return torch.nan_to_num(fraction).clamp_(0.0, 1.0).unsqueeze(-1) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/reset_dataset.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/reset_dataset.py new file mode 100644 index 000000000000..fe5322385b0c --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/reset_dataset.py @@ -0,0 +1,183 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Adaptive curriculum over a validated Franka Pour reset-state dataset.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING + +import torch + +from isaaclab.managers import CurriculumTermCfg +from isaaclab.managers.manager_base import ManagerTermBase + +from isaaclab_tasks.utils.adaptive_reset_sampler import AdaptiveResetSampler + +if TYPE_CHECKING: + from ..pour_env import FrankaPourEnv + + +def reset_dataset_difficulty( + states: Mapping[str, torch.Tensor], + task_contract: Mapping[str, object], +) -> torch.Tensor: + """Return a normalized reverse-curriculum difficulty for every dataset row. + + Grasping rows occupy the easy half of the range and are ordered by their task objective. + Non-grasping rows occupy the hard half and are graded from information already stored in the + dataset: normalized arm displacement, source/receiver displacement, and hand closure. This + avoids the arbitrary ordering that results when every non-grasping objective is fixed at -1. + + Args: + states: Row-aligned Franka Pour reset-state tensors. + task_contract: Geometry and robot contract stored with the dataset. + + Returns: + One difficulty per row in ``[0, 1]``, where lower is easier. + """ + category = states["category"] + objective = states["objective"] + if category.ndim != 1 or objective.shape != category.shape or category.numel() == 0: + raise ValueError("Reset dataset category and objective must be aligned non-empty vectors.") + if not bool(torch.isfinite(objective).all()) or not bool(((category == 0) | (category == 1)).all()): + raise ValueError("Reset dataset categories and objectives must be finite and valid.") + + difficulty = torch.empty_like(objective, dtype=torch.float32) + grasping = category == 1 + difficulty[grasping] = 0.5 * (1.0 - objective[grasping].float()).clamp(0.0, 1.0) + non_grasping = ~grasping + if not bool(non_grasping.any()): + return difficulty + + arm_position = states["arm_joint_position"][non_grasping].float() + arm_home = torch.as_tensor(task_contract["arm_home"], device=arm_position.device, dtype=arm_position.dtype) + arm_limits = torch.as_tensor( + task_contract["arm_joint_limits"], device=arm_position.device, dtype=arm_position.dtype + ) + if arm_home.shape != (arm_position.shape[1],) or arm_limits.shape != (arm_position.shape[1], 2): + raise ValueError("Reset dataset arm contract does not match its joint-state shape.") + arm_span = (arm_limits[:, 1] - arm_limits[:, 0]).clamp_min(torch.finfo(arm_position.dtype).eps) + arm_score = (2.0 * torch.abs(arm_position - arm_home) / arm_span).mean(dim=-1).clamp(0.0, 1.0) + + source_xy = states["source_root_pose"][non_grasping, :2].float() + source_center = torch.as_tensor( + task_contract["source_region_center"], device=source_xy.device, dtype=source_xy.dtype + )[:2] + target_xy = states["target_root_pose"][non_grasping, :2].float() + target_center = torch.as_tensor(task_contract["target_center_xy"], device=target_xy.device, dtype=target_xy.dtype) + support_lower = torch.as_tensor( + task_contract["tabletop_support_lower_xy"], device=source_xy.device, dtype=source_xy.dtype + ) + support_upper = torch.as_tensor( + task_contract["tabletop_support_upper_xy"], device=source_xy.device, dtype=source_xy.dtype + ) + support_diagonal = torch.linalg.vector_norm(support_upper - support_lower).clamp_min(1.0e-6) + source_score = (torch.linalg.vector_norm(source_xy - source_center, dim=-1) / support_diagonal).clamp(0.0, 1.0) + target_score = (torch.linalg.vector_norm(target_xy - target_center, dim=-1) / support_diagonal).clamp(0.0, 1.0) + + fingers = states["finger_joint_position"][non_grasping].float().mean(dim=-1) + gripper_lower, gripper_upper = ( + float(value) + for value in task_contract["gripper_position_range"] # type: ignore[arg-type] + ) + gripper_span = max(gripper_upper - gripper_lower, 1.0e-6) + closed_score = ((gripper_upper - fingers) / gripper_span).clamp(0.0, 1.0) + + non_grasp_score = 0.45 * arm_score + 0.30 * source_score + 0.15 * target_score + 0.10 * closed_score + difficulty[non_grasping] = 0.5 + 0.5 * non_grasp_score + return difficulty + + +class PourResetDatasetCurriculum(ManagerTermBase): + """Record reset outcomes and sample a stable easy-to-hard dataset frontier.""" + + def __init__(self, cfg: CurriculumTermCfg, env: FrankaPourEnv): + super().__init__(cfg, env) + if not getattr(env, "_uses_reset_dataset", False): + raise RuntimeError("PourResetDatasetCurriculum requires the reset-dataset environment variant.") + + states = env._reset_dataset_states + self._difficulty = reset_dataset_difficulty(states, env._reset_dataset_metadata["task_contract"]) + difficulty_order = torch.argsort(self._difficulty, stable=True) + sampler_cfg = env.cfg.reset_dataset_sampler.copy() + self._sampler = AdaptiveResetSampler( + difficulty_order, + sampler_cfg, + ) + + row_count = states["category"].numel() + self._frozen_rows = torch.arange(row_count, device=env.device, dtype=torch.long) + top_grasp_count = env.cfg.reset_dataset_top_grasp_count + if top_grasp_count is not None: + grasp_rows = torch.nonzero(states["category"] == 1, as_tuple=False).flatten() + if top_grasp_count > grasp_rows.numel(): + raise ValueError( + "reset_dataset_top_grasp_count exceeds the dataset grasp count: " + f"{top_grasp_count} > {grasp_rows.numel()}." + ) + grasp_order = torch.argsort(states["objective"][grasp_rows], descending=True, stable=True) + self._frozen_rows = grasp_rows[grasp_order[:top_grasp_count]] + if row_count == 0 or self._frozen_rows.numel() == 0: + raise RuntimeError("The reset dataset must expose at least one playback row.") + + @staticmethod + def _env_ids( + env: FrankaPourEnv, + env_ids: Sequence[int] | torch.Tensor | slice, + ) -> torch.Tensor: + """Normalize manager-provided environment IDs on the simulation device.""" + if isinstance(env_ids, slice): + return torch.arange(env.num_envs, device=env.device, dtype=torch.long)[env_ids] + return torch.as_tensor(env_ids, device=env.device, dtype=torch.long).flatten() + + def __call__( + self, + env: FrankaPourEnv, + env_ids: Sequence[int] | torch.Tensor | slice, + ) -> dict[str, float]: + """Record completed episodes, select exact raw rows, and report compact progress.""" + ids = self._env_ids(env, env_ids) + if ids.numel() > 0: + completed = (env.episode_length_buf[ids] > 0) & (env.reset_dataset_row_id[ids] >= 0) + completed_ids = ids[completed] + if completed_ids.numel() > 0 and not env.cfg.curriculum_freeze: + completed_rows = env.reset_dataset_row_id[completed_ids] + self._sampler.record(completed_rows, env.episode_succeeded[completed_ids]) + + forced_rows = env._forced_reset_dataset_row[ids] + if env.cfg.curriculum_freeze: + slots = torch.randint(self._frozen_rows.numel(), (ids.numel(),), device=env.device) + rows = self._frozen_rows[slots] + use_forced = forced_rows >= 0 + if bool(torch.any(use_forced)): + # Reuse the generic sampler's exact-row validation without changing frozen + # sampling for the remaining environments. + rows = torch.where( + use_forced, + self._sampler.sample(ids.numel(), forced_row_ids=forced_rows), + rows, + ) + else: + rows = self._sampler.sample(ids.numel(), forced_row_ids=forced_rows) + env.reset_dataset_row_id[ids] = rows + env.pour_target_frac[ids] = float(env.cfg.pour_target_frac) + + if env.cfg.curriculum_freeze: + return { + "frozen_pool_fraction": float(self._frozen_rows.numel() / self._difficulty.numel()), + "frozen_pool_size": float(self._frozen_rows.numel()), + } + + metrics = self._sampler.metrics() + return { + "predicted_success_rate": metrics["predicted_success_rate"], + "observed_success_rate": metrics["bounded_success_rate"], + "dataset_success_rate": metrics["cache_success_rate"], + "dataset_ever_solved_fraction": metrics["ever_solved_fraction"], + "frontier_fraction": metrics["frontier_fraction"], + "effective_pool_size": metrics["effective_pool_size"], + } diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/reset_mixture.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/reset_mixture.py new file mode 100644 index 000000000000..c09396b96b47 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/reset_mixture.py @@ -0,0 +1,17 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Deprecated compatibility names for the Franka Pour reset-dataset curriculum.""" + +from .reset_dataset import PourResetDatasetCurriculum + +RESET_MIXTURE_REGION_NAMES = ("reaching", "near_object", "grasped", "near_goal") +"""Deprecated region names retained for import compatibility.""" + +RESET_MIXTURE_STAGE_NAMES = ("randomized", "grasp", "carry", "tilt") +"""Deprecated stage names retained for import compatibility.""" + +PourResetMixture = PourResetDatasetCurriculum +"""Deprecated compatibility name for :class:`PourResetDatasetCurriculum`.""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/rewards.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/rewards.py new file mode 100644 index 000000000000..d97786abc48c --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/rewards.py @@ -0,0 +1,1031 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Bounded progress and outcome rewards for grasping and pouring MPM media.""" + +from __future__ import annotations + +import math +import warnings +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import torch + +from isaaclab.managers import ManagerTermBase, RewardTermCfg, SceneEntityCfg +from isaaclab.utils import math as math_utils + +from .terminations import source_grasp_milestones + +if TYPE_CHECKING: + from ..pour_env import FrankaPourEnv + + +def finite_joint_velocity_l2( + env: FrankaPourEnv, + asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"), + max_velocity: float = 20.0, +) -> torch.Tensor: + """Return a bounded, finite joint-velocity penalty. + + Non-finite velocities receive the maximum penalty so numerical failures cannot evade this term, + while clamping keeps the terminal transition consumable by PPO. + """ + velocity = env.scene[asset_cfg.name].data.joint_vel.torch[:, asset_cfg.joint_ids] + velocity = torch.nan_to_num( + velocity, + nan=float(max_velocity), + posinf=float(max_velocity), + neginf=-float(max_velocity), + ) + velocity = torch.clamp(velocity, min=-float(max_velocity), max=float(max_velocity)) + return torch.sum(torch.square(velocity), dim=-1) + + +def tcp_cup_distance_tanh(env: FrankaPourEnv, std: float = 1.0) -> torch.Tensor: + """Return bounded TCP-to-source-grasp proximity using distance scale ``std`` [m].""" + if not math.isfinite(std) or std <= 0.0: + raise ValueError("std must be finite and positive.") + distance = torch.linalg.vector_norm(env.tcp_pos_e() - env.cup_grasp_point_e(), dim=-1) + distance = torch.where(torch.isfinite(distance), distance, torch.full_like(distance, torch.inf)) + return 1.0 - torch.tanh(distance / float(std)) + + +def media_target_distance_tanh(env: FrankaPourEnv, std: float = 1.0) -> torch.Tensor: + """Return mean particle proximity to the source-exclusive target set. + + Distance to the receiving cavity supplies transport guidance. For particles still in the + source, distance to its open rim prevents nesting the loaded source inside the receiver from + scoring as delivery. Particles that have irreversibly reached the spill plane receive no + credit. + """ + if not math.isfinite(std) or std <= 0.0: + raise ValueError("std must be finite and positive.") + target_pose = env.target_pose_e() + source_pose = env.cup_pose_e() + particle_position = env.particle_pos_e() + target_quat = target_pose[:, None, 3:7].expand(-1, particle_position.shape[1], -1) + particle_target = math_utils.quat_apply_inverse( + target_quat, + particle_position - target_pose[:, None, :3], + ) + source_quat = source_pose[:, None, 3:7].expand_as(target_quat) + particle_source = math_utils.quat_apply_inverse( + source_quat, + particle_position - source_pose[:, None, :3], + ) + margin = float(env.cfg.particle_count_margin) + lower = env._target_inner_lo_t - margin + upper = env._target_inner_hi_t + margin + outside = torch.maximum(lower - particle_target, particle_target - upper).clamp_(min=0.0) + target_distance = torch.linalg.vector_norm(outside, dim=-1) + in_source, _, spilled = env.particle_region_masks() + source_exit_distance = torch.where( + in_source, + torch.clamp(env._source_inner_hi_t[2] + margin - particle_source[..., 2], min=0.0), + torch.zeros_like(target_distance), + ) + goal_distance = torch.maximum(target_distance, source_exit_distance) + valid = torch.isfinite(goal_distance) & ~spilled + quality = torch.where(valid, 1.0 - torch.tanh(goal_distance / float(std)), 0.0) + return quality.mean(dim=1) + + +def terminal_failure(env: FrankaPourEnv, include_time_out: bool = True) -> torch.Tensor: + """Return one unit-integral pulse for an unsuccessful terminal transition.""" + if not isinstance(include_time_out, bool): + raise TypeError("include_time_out must be a bool.") + success = _success_terminal(env) + completed = env.termination_manager.dones if include_time_out else env.termination_manager.terminated + failed = completed & ~success + # RewardManager multiplies terms by step_dt. Dividing here makes the configured weight the + # exact one-time episode penalty, independent of the control frequency. + return failed.float() / max(float(env.step_dt), 1.0e-6) + + +def _reach_quality(env: FrankaPourEnv, std: float) -> torch.Tensor: + distance = torch.linalg.norm(env.tcp_pos_e() - env.cup_grasp_point_e(), dim=-1) + return 1.0 - torch.tanh(distance / max(float(std), 1.0e-6)) + + +def _open_quality(env: FrankaPourEnv) -> torch.Tensor: + travel = max(float(env.gripper_open_width) - float(env.gripper_grasp_width), 1.0e-4) + return torch.clamp((env.gripper_width() - float(env.gripper_grasp_width)) / travel, 0.0, 1.0) + + +def _approach_potential( + env: FrankaPourEnv, + position_std: float, + orientation_std: float, + open_hand_fraction: float, +) -> torch.Tensor: + """Return bounded grasp-pose approach with recoverable open-hand coordination.""" + if not math.isfinite(position_std) or position_std <= 0.0: + raise ValueError("position_std must be finite and positive.") + if not math.isfinite(orientation_std) or orientation_std <= 0.0: + raise ValueError("orientation_std must be finite and positive.") + if not math.isfinite(open_hand_fraction) or not 0.0 <= open_hand_fraction <= 1.0: + raise ValueError("open_hand_fraction must lie in [0, 1].") + + distance = torch.linalg.vector_norm(env.tcp_pos_e() - env.cup_grasp_point_e(), dim=-1) + position_quality = 1.0 - torch.tanh(distance / float(position_std)) + + cup_quat = env.cup_pose_e()[:, 3:7] + desired_quat = math_utils.quat_mul(cup_quat, env.desired_grasp_tcp_quat_c()) + tcp_quat = env.tcp_pose_e()[:, 3:7] + error_quat = math_utils.quat_mul(math_utils.quat_conjugate(desired_quat), tcp_quat) + error_quat = math_utils.quat_unique(error_quat) + orientation_error = torch.linalg.vector_norm(math_utils.axis_angle_from_quat(error_quat), dim=-1) + orientation_quality = 1.0 - torch.tanh(orientation_error / float(orientation_std)) + + # Retain half of the Cartesian gradient even under a poor tool orientation. Multiplying two + # narrow kernels would make an independently sampled arm receive almost no useful reach signal. + pose_quality = position_quality * (0.5 + 0.5 * orientation_quality) + # Opening is useful while approaching, but its bonus fades to zero at the grasp pose. Closing + # prematurely therefore loses potential without erasing the position/orientation gradient that + # lets the policy recover and finish the approach. + potential = pose_quality + float(open_hand_fraction) * (1.0 - pose_quality) * _open_quality(env) + return torch.nan_to_num(potential, nan=0.0, posinf=1.0, neginf=0.0).clamp_(0.0, 1.0) + + +def _grasp_width_quality(env: FrankaPourEnv, preload_position: float) -> torch.Tensor: + """Return grasp quality requiring commanded preload and non-empty physical closure.""" + width = env.gripper_width() + open_width = float(env.gripper_open_width) + grasp_width = float(env.gripper_grasp_width) + travel = max(open_width - grasp_width, 1.0e-4) + close_progress = torch.clamp((open_width - width) / travel, 0.0, 1.0) + # Contact can settle slightly inside the nominal geometry, but an empty fully closed hand + # must not look grasped. This factor stays one throughout the intended open-to-contact path + # and falls linearly only after the fingers move inside the nominal cup width. + not_empty = torch.clamp(width / max(grasp_width, 1.0e-4), 0.0, 1.0) + open_position = 0.5 * open_width + preload_travel = max(open_position - float(preload_position), 1.0e-4) + gripper = env.action_manager.get_term("gripper_action") + command = gripper.commanded_position[:, 0] + commanded_preload = torch.clamp((open_position - command) / preload_travel, 0.0, 1.0) + # ``minimum`` is a conjunctive crossfade without the quadratic dead zone produced by a + # product: coordinated physical closure and commanded preload remain monotonic from open to + # contact, while either an open command or an empty fully closed hand still yields zero. + contact_quality = getattr(gripper, "contact_quality", not_empty) + return torch.minimum(close_progress, commanded_preload) * contact_quality + + +def _grasp_lift_potential( + env: FrankaPourEnv, + target_height: float, + grasp_reach_std: float, + grasp_preload_position: float, + grasp_fraction: float, +) -> torch.Tensor: + """Return bounded contact-qualified grasp and lift completion in ``[0, 1]``.""" + if not math.isfinite(target_height) or target_height <= 0.0: + raise ValueError("target_height must be finite and positive.") + if not math.isfinite(grasp_reach_std) or grasp_reach_std <= 0.0: + raise ValueError("grasp_reach_std must be finite and positive.") + if not math.isfinite(grasp_preload_position): + raise ValueError("grasp_preload_position must be finite.") + if not math.isfinite(grasp_fraction) or not 0.0 <= grasp_fraction <= 1.0: + raise ValueError("grasp_fraction must lie in [0, 1].") + + distance = torch.linalg.vector_norm(env.tcp_pos_e() - env.cup_grasp_point_e(), dim=-1) + proximity = torch.clamp(1.0 - distance / float(grasp_reach_std), 0.0, 1.0) + # Compact-support smoothstep prevents closing an empty hand at stand-off from earning grasp + # credit. The broader approach potential supplies the gradient until this contact neighborhood. + proximity = proximity.square() * (3.0 - 2.0 * proximity) + grasp = proximity * _grasp_width_quality(env, grasp_preload_position) + height = torch.clamp( + (env.cup_pose_e()[:, 2] - float(env.cup_reset_height)) / float(target_height), + 0.0, + 1.0, + ) + potential = grasp * (float(grasp_fraction) + (1.0 - float(grasp_fraction)) * height) + return torch.nan_to_num(potential, nan=0.0, posinf=1.0, neginf=0.0).clamp_(0.0, 1.0) + + +def _lift_potential( + env: FrankaPourEnv, + target_height: float, + reach_std: float, + grasp_reach_std: float, + grasp_preload_position: float, + approach_fraction: float, + grasp_fraction: float, +) -> torch.Tensor: + """Return bounded open-approach, contact-grasp, and lift completion in ``[0, 1]``.""" + if approach_fraction < 0.0 or grasp_fraction < 0.0 or approach_fraction + grasp_fraction > 1.0: + raise ValueError("approach_fraction and grasp_fraction must be nonnegative and sum to at most one.") + reach = _reach_quality(env, reach_std) + # Width alone cannot distinguish a real grasp from an empty hand closed to the cup width. + # Gate grasp and lift completion much more sharply than the broad approach shaping so that + # closing at a stand-off pose always loses potential, while closing at the grasp point gains it. + grasp_reach = _reach_quality(env, grasp_reach_std) + height = torch.clamp( + (env.cup_pose_e()[:, 2] - float(env.cup_reset_height)) / max(float(target_height), 1.0e-6), + 0.0, + 1.0, + ) + grasp = _grasp_width_quality(env, grasp_preload_position) + lift_fraction = 1.0 - float(approach_fraction) - float(grasp_fraction) + potential = ( + float(approach_fraction) * reach * _open_quality(env) + + float(grasp_fraction) * grasp_reach * grasp + + lift_fraction * height * grasp_reach * grasp + ) + return torch.nan_to_num(potential, nan=0.0, posinf=1.0, neginf=0.0).clamp_(0.0, 1.0) + + +def _align_potential( + env: FrankaPourEnv, + lift_height: float, + std: float, + source_offset_xy: Sequence[float], + grasp_reach_std: float, + grasp_preload_position: float, +) -> torch.Tensor: + """Return bounded held-source-to-receiver alignment in ``[0, 1]``.""" + if len(source_offset_xy) != 2 or any(not math.isfinite(value) for value in source_offset_xy): + raise ValueError("source_offset_xy must contain two finite values.") + cup = env.cup_pose_e()[:, :3] + grasp_point = env.cup_grasp_point_e() + target = env.target_pose_e()[:, :3] + lifted = torch.clamp((cup[:, 2] - float(env.cup_reset_height)) / max(float(lift_height), 1.0e-6), 0.0, 1.0) + desired_grasp_xy = target[:, :2] + cup.new_tensor(source_offset_xy) + distance_xy = torch.linalg.norm(grasp_point[:, :2] - desired_grasp_xy, dim=-1) + aligned = 1.0 - torch.tanh(distance_xy / max(float(std), 1.0e-6)) + grasp = _grasp_width_quality(env, grasp_preload_position) + grasp_reach = _reach_quality(env, grasp_reach_std) + potential = lifted * aligned * grasp_reach * grasp + return torch.nan_to_num(potential, nan=0.0, posinf=1.0, neginf=0.0).clamp_(0.0, 1.0) + + +def _pour_tilt_potential( + env: FrankaPourEnv, + target_tilt: float, + pour_direction_xy: Sequence[float], + source_mouth_height: float, + alignment_radius: float, + active_through_stage: int, + min_lift_height: float, + max_tcp_distance: float, + max_gripper_width_error: float, + max_gripper_command: float, +) -> torch.Tensor: + """Return bounded early-stage progress toward the authored physical pour in ``[0, 1]``.""" + if not math.isfinite(target_tilt) or not 0.0 < target_tilt < math.pi: + raise ValueError(f"target_tilt must lie in (0, pi), got {target_tilt}.") + if len(pour_direction_xy) != 2 or any(not math.isfinite(value) for value in pour_direction_xy): + raise ValueError("pour_direction_xy must contain two finite values.") + direction_norm = math.hypot(float(pour_direction_xy[0]), float(pour_direction_xy[1])) + if direction_norm <= 0.0: + raise ValueError("pour_direction_xy must be nonzero.") + if not math.isfinite(source_mouth_height) or source_mouth_height < 0.0: + raise ValueError(f"source_mouth_height must be finite and nonnegative, got {source_mouth_height}.") + if not math.isfinite(alignment_radius) or alignment_radius <= 0.0: + raise ValueError(f"alignment_radius must be finite and positive, got {alignment_radius}.") + if active_through_stage < 0: + raise ValueError(f"active_through_stage must be nonnegative, got {active_through_stage}.") + + cup_pose = env.cup_pose_e() + target_pose = env.target_pose_e() + local_open_axis = torch.zeros_like(cup_pose[:, :3]) + local_open_axis[:, 2] = 1.0 + open_axis = math_utils.quat_apply(cup_pose[:, 3:7], local_open_axis) + direction_x = float(pour_direction_xy[0]) / direction_norm + direction_y = float(pour_direction_xy[1]) / direction_norm + directed_opening = open_axis[:, 0] * direction_x + open_axis[:, 1] * direction_y + # ``atan2`` remains monotonic after the rim passes horizontal, unlike a sine projection. This + # lets the curriculum teach the validated deep drain while rejecting rotation away from + # the receiver and sideways inversion. + directed_angle = torch.atan2(torch.clamp(directed_opening, min=0.0), open_axis[:, 2]) + directed_angle = torch.where(directed_opening > 0.0, directed_angle, 0.0) + tilt = torch.clamp(directed_angle / float(target_tilt), 0.0, 1.0) + + local_mouth = torch.zeros_like(cup_pose[:, :3]) + local_mouth[:, 2] = float(source_mouth_height) + mouth_position = cup_pose[:, :3] + math_utils.quat_apply(cup_pose[:, 3:7], local_mouth) + distance_xy = torch.linalg.vector_norm(mouth_position[:, :2] - target_pose[:, :2], dim=-1) + proximity = torch.clamp(1.0 - distance_xy / float(alignment_radius), 0.0, 1.0) + # Compact-support smoothstep prevents the carry stage from earning tilt credit before the cup + # reaches the receiver, while avoiding a discontinuity at the alignment boundary. + aligned = proximity.square() * (3.0 - 2.0 * proximity) + _, _, held_pour = source_grasp_milestones( + env, + min_lift_height=min_lift_height, + max_tcp_distance=max_tcp_distance, + max_gripper_width_error=max_gripper_width_error, + max_gripper_command=max_gripper_command, + ) + active_stage = env.curriculum_stage <= int(active_through_stage) + potential = tilt * aligned * held_pour.float() * active_stage.float() + return torch.nan_to_num(potential, nan=0.0, posinf=1.0, neginf=0.0).clamp_(0.0, 1.0) + + +def _pour_reference_potential( + env: FrankaPourEnv, + start_q: Sequence[float], + target_q: Sequence[float], + active_stage: int, + min_lift_height: float, + max_tcp_distance: float, + max_gripper_width_error: float, + max_gripper_command: float, +) -> torch.Tensor: + """Return held progress along the validated stage-zero arm trajectory in ``[0, 1]``.""" + arm_q = env.arm_joint_pos() + if len(start_q) != arm_q.shape[1] or len(target_q) != arm_q.shape[1]: + raise ValueError(f"start_q and target_q must each contain {arm_q.shape[1]} joint positions.") + if any(not math.isfinite(value) for value in (*start_q, *target_q)): + raise ValueError("start_q and target_q must contain finite joint positions.") + reference_distance = math.sqrt(sum((target - start) ** 2 for start, target in zip(start_q, target_q, strict=True))) + if reference_distance <= 0.0: + raise ValueError("start_q and target_q must be distinct.") + target = arm_q.new_tensor(target_q) + distance = torch.linalg.vector_norm(arm_q - target, dim=-1) + progress = torch.clamp(1.0 - distance / reference_distance, 0.0, 1.0) + _, _, held_pour = source_grasp_milestones( + env, + min_lift_height=min_lift_height, + max_tcp_distance=max_tcp_distance, + max_gripper_width_error=max_gripper_width_error, + max_gripper_command=max_gripper_command, + ) + active = env.curriculum_stage == int(active_stage) + potential = progress * held_pour.float() * active.float() + return torch.nan_to_num(potential, nan=0.0, posinf=1.0, neginf=0.0).clamp_(0.0, 1.0) + + +class _SignedPotentialProgress(ManagerTermBase): + """Track independent per-environment potential history.""" + + def __init__(self, cfg: RewardTermCfg, env: FrankaPourEnv): + super().__init__(cfg, env) + self._previous_potential = torch.zeros(env.num_envs, device=env.device) + self._initialized = torch.zeros(env.num_envs, device=env.device, dtype=torch.bool) + + def _reset_potential( + self, + potential: torch.Tensor, + env_ids: Sequence[int] | torch.Tensor | slice | None, + ) -> None: + if env_ids is None: + env_ids = slice(None) + self._previous_potential[env_ids] = potential[env_ids] + self._initialized[env_ids] = True + + def _signed_progress(self, potential: torch.Tensor, step_dt: float) -> torch.Tensor: + progress = torch.where(self._initialized, potential - self._previous_potential, 0.0) + self._previous_potential.copy_(potential) + self._initialized.fill_(True) + return progress / max(float(step_dt), 1.0e-6) + + def _discounted_progress( + self, + potential: torch.Tensor, + step_dt: float, + discount_factor: float, + terminal: torch.Tensor, + ) -> torch.Tensor: + """Return policy-invariant discounted potential shaping. + + Treating the post-transition potential as zero on terminal states preserves the standard + episodic potential-shaping convention. A forward/reverse cycle then has exactly the same + discounted shaping return as holding the original state for the same duration. + """ + if not 0.0 < float(discount_factor) <= 1.0: + raise ValueError("discount_factor must lie in (0, 1].") + next_potential = torch.where(terminal, torch.zeros_like(potential), potential) + progress = torch.where( + self._initialized, + float(discount_factor) * next_potential - self._previous_potential, + torch.zeros_like(potential), + ) + self._previous_potential.copy_(potential) + self._initialized.fill_(True) + return progress / max(float(step_dt), 1.0e-6) + + +class PourTaskProgress(_SignedPotentialProgress): + """Signed progress through approach, grasp, lift, align, and tilt milestones. + + The single ordered physical potential replaces independently weighted, stage-switched shaping + terms. Its episode integral telescopes, so holding an intermediate pose earns nothing and + reversing or dropping the cup repays the corresponding forward progress. + """ + + def reset(self, env_ids: Sequence[int] | torch.Tensor | slice | None = None) -> None: + """Start progress accounting from the physical state supplied by the reset.""" + params = dict(self.cfg.params or {}) + params.pop("discount_factor", None) + potential = self._potential(self._env, **params) + self._reset_potential(potential, env_ids) + + @staticmethod + def _potential( + env: FrankaPourEnv, + target_height: float, + reach_std: float, + grasp_reach_std: float, + grasp_preload_position: float, + lift_height: float, + align_std: float, + source_offset_xy: Sequence[float], + target_tilt: float, + pour_direction_xy: Sequence[float], + source_mouth_height: float, + alignment_radius: float, + active_through_stage: int, + min_lift_height: float, + max_tcp_distance: float, + max_gripper_width_error: float, + max_gripper_command: float, + ) -> torch.Tensor: + approach_grasp_lift = _lift_potential( + env, + target_height=target_height, + reach_std=reach_std, + grasp_reach_std=grasp_reach_std, + grasp_preload_position=grasp_preload_position, + approach_fraction=0.20, + grasp_fraction=0.30, + ) + align = _align_potential( + env, + lift_height=lift_height, + std=align_std, + source_offset_xy=source_offset_xy, + grasp_reach_std=grasp_reach_std, + grasp_preload_position=grasp_preload_position, + ) + tilt = _pour_tilt_potential( + env, + target_tilt=target_tilt, + pour_direction_xy=pour_direction_xy, + source_mouth_height=source_mouth_height, + alignment_radius=alignment_radius, + active_through_stage=active_through_stage, + min_lift_height=min_lift_height, + max_tcp_distance=max_tcp_distance, + max_gripper_width_error=max_gripper_width_error, + max_gripper_command=max_gripper_command, + ) + return 0.45 * approach_grasp_lift + 0.20 * align + 0.35 * tilt + + def __call__( + self, + env: FrankaPourEnv, + target_height: float, + reach_std: float, + grasp_reach_std: float, + grasp_preload_position: float, + lift_height: float, + align_std: float, + source_offset_xy: Sequence[float], + target_tilt: float, + pour_direction_xy: Sequence[float], + source_mouth_height: float, + alignment_radius: float, + active_through_stage: int, + min_lift_height: float, + max_tcp_distance: float, + max_gripper_width_error: float, + max_gripper_command: float, + discount_factor: float, + ) -> torch.Tensor: + potential = self._potential( + env, + target_height=target_height, + reach_std=reach_std, + grasp_reach_std=grasp_reach_std, + grasp_preload_position=grasp_preload_position, + lift_height=lift_height, + align_std=align_std, + source_offset_xy=source_offset_xy, + target_tilt=target_tilt, + pour_direction_xy=pour_direction_xy, + source_mouth_height=source_mouth_height, + alignment_radius=alignment_radius, + active_through_stage=active_through_stage, + min_lift_height=min_lift_height, + max_tcp_distance=max_tcp_distance, + max_gripper_width_error=max_gripper_width_error, + max_gripper_command=max_gripper_command, + ) + # The five-second manipulation attempt is a finite-horizon episode. Close the shaping + # potential on both task terminations and the deadline so its discounted return remains + # policy-invariant and failed timeouts cannot retain bootstrap value. + terminal = env.termination_manager.dones + return self._discounted_progress( + potential, + env.step_dt, + discount_factor=discount_factor, + terminal=terminal, + ) + + +class ApproachProgress(_SignedPotentialProgress): + """Discounted progress toward an open, correctly oriented source-cup grasp pose.""" + + def reset(self, env_ids: Sequence[int] | torch.Tensor | slice | None = None) -> None: + """Start progress accounting from each environment's randomized arm pose.""" + params = dict(self.cfg.params or {}) + params.pop("discount_factor", None) + active_from_stage = int(params.pop("active_from_stage")) + potential = _approach_potential(self._env, **params) + potential *= (self._env.curriculum_stage >= active_from_stage).float() + self._reset_potential(potential, env_ids) + + def __call__( + self, + env: FrankaPourEnv, + position_std: float, + orientation_std: float, + open_hand_fraction: float, + active_from_stage: int, + discount_factor: float, + ) -> torch.Tensor: + potential = _approach_potential( + env, + position_std=position_std, + orientation_std=orientation_std, + open_hand_fraction=open_hand_fraction, + ) + potential *= (env.curriculum_stage >= int(active_from_stage)).float() + return self._discounted_progress( + potential, + env.step_dt, + discount_factor=discount_factor, + terminal=env.termination_manager.dones, + ) + + +class GraspLiftProgress(_SignedPotentialProgress): + """Discounted progress from near-contact closure through lifting the source cup.""" + + def reset(self, env_ids: Sequence[int] | torch.Tensor | slice | None = None) -> None: + """Start progress accounting from each environment's reset grasp state.""" + params = dict(self.cfg.params or {}) + params.pop("discount_factor", None) + active_from_stage = int(params.pop("active_from_stage")) + potential = _grasp_lift_potential(self._env, **params) + potential *= (self._env.curriculum_stage >= active_from_stage).float() + self._reset_potential(potential, env_ids) + + def __call__( + self, + env: FrankaPourEnv, + target_height: float, + grasp_reach_std: float, + grasp_preload_position: float, + grasp_fraction: float, + active_from_stage: int, + discount_factor: float, + ) -> torch.Tensor: + potential = _grasp_lift_potential( + env, + target_height=target_height, + grasp_reach_std=grasp_reach_std, + grasp_preload_position=grasp_preload_position, + grasp_fraction=grasp_fraction, + ) + potential *= (env.curriculum_stage >= int(active_from_stage)).float() + return self._discounted_progress( + potential, + env.step_dt, + discount_factor=discount_factor, + terminal=env.termination_manager.dones, + ) + + +class LiftProgress(_SignedPotentialProgress): + """Signed change in bounded grasp-and-lift completion from the lift stage onward. + + The term is divided by the environment step interval because + :class:`~isaaclab.managers.RewardManager` multiplies every reward by that interval. Its integrated + episode contribution therefore telescopes: reversing a lift repays the earlier positive reward. + """ + + def reset(self, env_ids: Sequence[int] | torch.Tensor | slice | None = None) -> None: + params = self.cfg.params or {} + potential = _lift_potential( + self._env, + target_height=float(params.get("target_height", 0.12)), + reach_std=float(params.get("reach_std", 0.07)), + grasp_reach_std=float(params.get("grasp_reach_std", 0.015)), + grasp_preload_position=float(params.get("grasp_preload_position", 0.025)), + approach_fraction=float(params.get("approach_fraction", 0.2)), + grasp_fraction=float(params.get("grasp_fraction", 0.3)), + ) + self._reset_potential(potential, env_ids) + + def __call__( + self, + env: FrankaPourEnv, + target_height: float = 0.12, + reach_std: float = 0.07, + grasp_reach_std: float = 0.015, + grasp_preload_position: float = 0.025, + approach_fraction: float = 0.2, + grasp_fraction: float = 0.3, + ) -> torch.Tensor: + potential = _lift_potential( + env, + target_height=target_height, + reach_std=reach_std, + grasp_reach_std=grasp_reach_std, + grasp_preload_position=grasp_preload_position, + approach_fraction=approach_fraction, + grasp_fraction=grasp_fraction, + ) + progress = self._signed_progress(potential, env.step_dt) + return progress * (env.curriculum_stage >= 2) + + +class AlignProgress(_SignedPotentialProgress): + """Signed change in held-source-to-receiver alignment from the carry stage onward. + + Returning to an earlier alignment produces the exact negative of the corresponding forward + progress, preventing cyclic motion from accumulating reward. + """ + + def reset(self, env_ids: Sequence[int] | torch.Tensor | slice | None = None) -> None: + params = self.cfg.params or {} + potential = _align_potential( + self._env, + lift_height=float(params.get("lift_height", 0.06)), + std=float(params.get("std", 0.12)), + source_offset_xy=params.get("source_offset_xy", (0.0, 0.0)), + grasp_reach_std=float(params.get("grasp_reach_std", 0.015)), + grasp_preload_position=float(params.get("grasp_preload_position", 0.025)), + ) + self._reset_potential(potential, env_ids) + + def __call__( + self, + env: FrankaPourEnv, + lift_height: float = 0.06, + std: float = 0.12, + source_offset_xy: Sequence[float] = (0.0, 0.0), + grasp_reach_std: float = 0.015, + grasp_preload_position: float = 0.025, + ) -> torch.Tensor: + potential = _align_potential( + env, + lift_height=lift_height, + std=std, + source_offset_xy=source_offset_xy, + grasp_reach_std=grasp_reach_std, + grasp_preload_position=grasp_preload_position, + ) + progress = self._signed_progress(potential, env.step_dt) + return progress * (env.curriculum_stage >= 1) + + +class PourTiltProgress(_SignedPotentialProgress): + """Signed early-curriculum progress toward tilting a held cup over the receiver. + + This term is active only through the configured curriculum stage. Its episode integral + telescopes, so reversing a tilt or releasing the cup repays the corresponding positive reward + and cyclic wiggling cannot accumulate return. + """ + + def reset(self, env_ids: Sequence[int] | torch.Tensor | slice | None = None) -> None: + params = self.cfg.params or {} + potential = _pour_tilt_potential( + self._env, + target_tilt=float(params.get("target_tilt", math.radians(150.0))), + pour_direction_xy=params.get("pour_direction_xy", (0.0, -1.0)), + source_mouth_height=float(params.get("source_mouth_height", 0.0)), + alignment_radius=float(params.get("alignment_radius", 0.10)), + active_through_stage=int(params.get("active_through_stage", 1)), + min_lift_height=float(params.get("min_lift_height", 0.05)), + max_tcp_distance=float(params.get("max_tcp_distance", 0.015)), + max_gripper_width_error=float(params.get("max_gripper_width_error", 0.012)), + max_gripper_command=float(params.get("max_gripper_command", 0.025)), + ) + self._reset_potential(potential, env_ids) + + def __call__( + self, + env: FrankaPourEnv, + target_tilt: float = math.radians(150.0), + pour_direction_xy: Sequence[float] = (0.0, -1.0), + source_mouth_height: float = 0.0, + alignment_radius: float = 0.10, + active_through_stage: int = 1, + min_lift_height: float = 0.05, + max_tcp_distance: float = 0.015, + max_gripper_width_error: float = 0.012, + max_gripper_command: float = 0.025, + ) -> torch.Tensor: + potential = _pour_tilt_potential( + env, + target_tilt=target_tilt, + pour_direction_xy=pour_direction_xy, + source_mouth_height=source_mouth_height, + alignment_radius=alignment_radius, + active_through_stage=active_through_stage, + min_lift_height=min_lift_height, + max_tcp_distance=max_tcp_distance, + max_gripper_width_error=max_gripper_width_error, + max_gripper_command=max_gripper_command, + ) + return self._signed_progress(potential, env.step_dt) + + +class PourReferenceProgress(_SignedPotentialProgress): + """Signed stage-zero progress along the validated held-pour joint trajectory.""" + + def reset(self, env_ids: Sequence[int] | torch.Tensor | slice | None = None) -> None: + params = self.cfg.params or {} + potential = _pour_reference_potential( + self._env, + start_q=params.get("start_q", ()), + target_q=params.get("target_q", ()), + active_stage=int(params.get("active_stage", 0)), + min_lift_height=float(params.get("min_lift_height", 0.05)), + max_tcp_distance=float(params.get("max_tcp_distance", 0.015)), + max_gripper_width_error=float(params.get("max_gripper_width_error", 0.012)), + max_gripper_command=float(params.get("max_gripper_command", 0.025)), + ) + self._reset_potential(potential, env_ids) + + def __call__( + self, + env: FrankaPourEnv, + start_q: Sequence[float], + target_q: Sequence[float], + active_stage: int = 0, + min_lift_height: float = 0.05, + max_tcp_distance: float = 0.015, + max_gripper_width_error: float = 0.012, + max_gripper_command: float = 0.025, + ) -> torch.Tensor: + potential = _pour_reference_potential( + env, + start_q=start_q, + target_q=target_q, + active_stage=active_stage, + min_lift_height=min_lift_height, + max_tcp_distance=max_tcp_distance, + max_gripper_width_error=max_gripper_width_error, + max_gripper_command=max_gripper_command, + ) + return self._signed_progress(potential, env.step_dt) + + +def _warn_legacy_reward(name: str) -> None: + warnings.warn( + f"mdp.{name} is deprecated; use the bounded progress and outcome terms from FrankaPourEnvCfg instead.", + DeprecationWarning, + stacklevel=2, + ) + + +def _legacy_closure_quality(env: FrankaPourEnv) -> torch.Tensor: + travel = max(float(env.gripper_open_width) - float(env.gripper_grasp_width), 1.0e-4) + return torch.clamp((float(env.gripper_open_width) - env.gripper_width()) / travel, 0.0, 1.0) + + +def reach_cup(env: FrankaPourEnv, std: float = 0.10) -> torch.Tensor: + """Deprecated Cartesian reach reward retained for downstream compatibility.""" + _warn_legacy_reward("reach_cup") + return _reach_quality(env, std) + + +def grasp_cup(env: FrankaPourEnv, reach_std: float = 0.06) -> torch.Tensor: + """Deprecated grasp reward retained for downstream compatibility.""" + _warn_legacy_reward("grasp_cup") + return _reach_quality(env, reach_std) * _legacy_closure_quality(env) + + +def lift_cup(env: FrankaPourEnv, target_height: float = 0.12, reach_std: float = 0.07) -> torch.Tensor: + """Deprecated lift reward retained for downstream compatibility.""" + _warn_legacy_reward("lift_cup") + height = torch.clamp( + (env.cup_pose_e()[:, 2] - float(env.cup_reset_height)) / max(float(target_height), 1.0e-6), + 0.0, + 1.0, + ) + return height * _reach_quality(env, reach_std) * _legacy_closure_quality(env) + + +def lift_command_progress( + env: FrankaPourEnv, + target_height: float = 0.12, + reach_std: float = 0.07, +) -> torch.Tensor: + """Deprecated Cartesian lift-command reward retained for downstream compatibility.""" + _warn_legacy_reward("lift_command_progress") + height = torch.clamp( + (env.cup_pose_e()[:, 2] - float(env.cup_reset_height)) / max(float(target_height), 1.0e-6), + 0.0, + 1.0, + ) + upward = torch.clamp(env.action_manager.action[:, 2], 0.0, 1.0) + grasp = _reach_quality(env, reach_std) * _legacy_closure_quality(env) + return grasp * (1.0 - height) * upward + + +def align_cup_over_target( + env: FrankaPourEnv, + lift_height: float = 0.06, + std: float = 0.12, +) -> torch.Tensor: + """Deprecated source-alignment reward retained for downstream compatibility.""" + _warn_legacy_reward("align_cup_over_target") + cup = env.cup_pose_e()[:, :3] + target = env.target_pose_e()[:, :3] + lifted = torch.clamp( + (cup[:, 2] - float(env.cup_reset_height)) / max(float(lift_height), 1.0e-6), + 0.0, + 1.0, + ) + distance_xy = torch.linalg.vector_norm(cup[:, :2] - target[:, :2], dim=-1) + return lifted * (1.0 - torch.tanh(distance_xy / max(float(std), 1.0e-6))) + + +def align_command_progress( + env: FrankaPourEnv, + lift_height: float = 0.06, + std: float = 0.12, +) -> torch.Tensor: + """Deprecated Cartesian alignment-command reward retained for downstream compatibility.""" + _warn_legacy_reward("align_command_progress") + cup = env.cup_pose_e()[:, :3] + target = env.target_pose_e()[:, :3] + lifted = torch.clamp( + (cup[:, 2] - float(env.cup_reset_height)) / max(float(lift_height), 1.0e-6), + 0.0, + 1.0, + ) + delta_xy = target[:, :2] - cup[:, :2] + distance = torch.linalg.vector_norm(delta_xy, dim=-1) + direction = delta_xy / torch.clamp(distance[:, None], min=1.0e-6) + toward = torch.clamp(torch.sum(env.action_manager.action[:, :2] * direction, dim=-1), 0.0, 1.0) + return lifted * torch.tanh(distance / max(float(std), 1.0e-6)) * toward + + +def _legacy_cup_up_z(env: FrankaPourEnv) -> torch.Tensor: + quat = env.cup_pose_e()[:, 3:7] + up = torch.zeros((quat.shape[0], 3), device=quat.device, dtype=quat.dtype) + up[:, 2] = 1.0 + xyz = quat[:, :3] + cross = 2.0 * torch.cross(xyz, up, dim=-1) + rotated = up + quat[:, 3:4] * cross + torch.cross(xyz, cross, dim=-1) + return rotated[:, 2] + + +def tilt_over_target( + env: FrankaPourEnv, + lift_height: float = 0.06, + align_std: float = 0.10, +) -> torch.Tensor: + """Deprecated source-tilt reward retained for downstream compatibility.""" + _warn_legacy_reward("tilt_over_target") + cup = env.cup_pose_e()[:, :3] + lifted = torch.clamp( + (cup[:, 2] - float(env.cup_reset_height)) / max(float(lift_height), 1.0e-6), + 0.0, + 1.0, + ) + distance_xy = torch.linalg.vector_norm(cup[:, :2] - env.target_pose_e()[:, :2], dim=-1) + aligned = 1.0 - torch.tanh(distance_xy / max(float(align_std), 1.0e-6)) + tilt = torch.clamp( + (math.cos(math.pi / 3.0) - _legacy_cup_up_z(env)) / math.cos(math.pi / 3.0), + 0.0, + 1.0, + ) + return lifted * aligned * tilt + + +def tilt_command_progress( + env: FrankaPourEnv, + lift_height: float = 0.06, + align_std: float = 0.10, +) -> torch.Tensor: + """Deprecated Cartesian tilt-command reward retained for downstream compatibility.""" + _warn_legacy_reward("tilt_command_progress") + cup = env.cup_pose_e()[:, :3] + lifted = torch.clamp( + (cup[:, 2] - float(env.cup_reset_height)) / max(float(lift_height), 1.0e-6), + 0.0, + 1.0, + ) + distance_xy = torch.linalg.vector_norm(cup[:, :2] - env.target_pose_e()[:, :2], dim=-1) + aligned = 1.0 - torch.tanh(distance_xy / max(float(align_std), 1.0e-6)) + tilt = torch.clamp( + (math.cos(math.pi / 3.0) - _legacy_cup_up_z(env)) / math.cos(math.pi / 3.0), + 0.0, + 1.0, + ) + rotate_toward_pour = torch.clamp(env.action_manager.action[:, 3], 0.0, 1.0) + return lifted * aligned * (1.0 - tilt) * rotate_toward_pour + + +def particles_in_target(env: FrankaPourEnv) -> torch.Tensor: + return env.count_in_target() / max(env.num_particles, 1) + + +def particles_in_source(env: FrankaPourEnv) -> torch.Tensor: + return env.count_in_source() / max(env.num_particles, 1) + + +def _success_terminal(env: FrankaPourEnv) -> torch.Tensor: + """Return this step's managed success, or false when a replay tool removed that term.""" + if "success" not in env.termination_manager.active_terms: + return torch.zeros(env.num_envs, device=env.device, dtype=torch.bool) + return env.termination_manager.get_term("success") + + +class HeldDeliveryProgress(ManagerTermBase): + """Reward signed progress toward the active held-delivery target. + + Credit is capped at each environment's success threshold, decreases when qualified particles + leave the receiver, and is fully clawed back when an episode ends without success. A successful + terminal transition retains its final credit because the separate success term rewards the same + predicate used by curriculum progression. + """ + + def __init__(self, cfg: RewardTermCfg, env: FrankaPourEnv): + super().__init__(cfg, env) + self._previous_credit = torch.zeros(env.num_envs, device=env.device) + + def reset(self, env_ids: Sequence[int] | torch.Tensor | None = None) -> None: + if env_ids is None: + env_ids = slice(None) + self._previous_credit[env_ids] = 0.0 + + def __call__( + self, + env: FrankaPourEnv, + min_lift_height: float = 0.05, + max_tcp_distance: float = 0.015, + max_gripper_width_error: float = 0.012, + max_gripper_command: float = 0.025, + ) -> torch.Tensor: + _, _, held_pour = source_grasp_milestones( + env, + min_lift_height=min_lift_height, + max_tcp_distance=max_tcp_distance, + max_gripper_width_error=max_gripper_width_error, + max_gripper_command=max_gripper_command, + ) + env.update_held_delivery_tracker(held_pour) + current_fraction = env.current_held_delivered_mask().sum(dim=1).float() / max(env.num_particles, 1) + credit = torch.minimum(current_fraction, env.pour_target_frac) + success_terminal = _success_terminal(env) + unsuccessful_done = env.termination_manager.dones & ~success_terminal + retained_credit = torch.where(unsuccessful_done, torch.zeros_like(credit), credit) + progress = retained_credit - self._previous_credit + self._previous_credit.copy_(retained_credit) + return progress / max(float(env.step_dt), 1.0e-6) + + +# Backward-compatible name retained for downstream configurations and imports. +NewlyDeliveredParticles = HeldDeliveryProgress + + +class NewlySpilledParticles(ManagerTermBase): + """Penalize each particle's first spill as a time-step-independent pulse.""" + + def __init__(self, cfg: RewardTermCfg, env: FrankaPourEnv): + super().__init__(cfg, env) + self._spilled = torch.zeros( + (env.num_envs, env.num_particles), + device=env.device, + dtype=torch.bool, + ) + + def reset(self, env_ids: Sequence[int] | torch.Tensor | None = None) -> None: + if env_ids is None: + env_ids = slice(None) + self._spilled[env_ids] = False + + def __call__(self, env: FrankaPourEnv) -> torch.Tensor: + spilled = env.particles_spilled_mask() + newly_spilled = spilled & ~self._spilled + self._spilled |= spilled + fraction = newly_spilled.sum(dim=1).float() / max(env.num_particles, 1) + return fraction / max(float(env.step_dt), 1.0e-6) + + +def spilled_particles(env: FrankaPourEnv) -> torch.Tensor: + """Fraction irrecoverably spilled onto or below the table.""" + return env.spilled_fraction() + + +def pour_success_bonus(env: FrankaPourEnv) -> torch.Tensor: + """Reward exactly the stable-success termination used by curriculum progression.""" + success = _success_terminal(env) + return success.float() / max(float(env.step_dt), 1.0e-6) + + +def sustained_pour_success(env: FrankaPourEnv, dwell_time_s: float = 0.15) -> torch.Tensor: + """Reward the current dwell-qualified success state without a terminal pulse.""" + if not math.isfinite(dwell_time_s) or dwell_time_s <= 0.0: + raise ValueError("dwell_time_s must be finite and positive.") + dwell_steps = max(1, math.ceil(float(dwell_time_s) / max(float(env.step_dt), 1.0e-6))) + return (env._success_dwell_count >= dwell_steps).float() diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/terminations.py new file mode 100644 index 000000000000..4f59f9182df1 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/terminations.py @@ -0,0 +1,313 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Termination terms for the pour task. + +In addition to the standard timeout, instability guards reset non-finite states and finite media +that leave the task workspace before they can expand the allocating sparse grid without bound. +""" + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from ..pour_env import FrankaPourEnv + + +_GRIPPER_POSITION_TOLERANCE = 1.0e-6 + + +def _state_finite( + robot_joint_pos: torch.Tensor, + robot_joint_vel: torch.Tensor, + tcp_body_q: torch.Tensor, + cup_body_q: torch.Tensor, + cup_lin_vel: torch.Tensor, + cup_ang_vel: torch.Tensor, + particle_pos: torch.Tensor, +) -> torch.Tensor: + """Return a per-environment finite-state mask using unsanitized simulation tensors.""" + robot_ok = torch.isfinite(robot_joint_pos).all(dim=-1) & torch.isfinite(robot_joint_vel).all(dim=-1) + tcp_ok = torch.isfinite(tcp_body_q).all(dim=-1) + cup_ok = ( + torch.isfinite(cup_body_q).all(dim=-1) + & torch.isfinite(cup_lin_vel).all(dim=-1) + & torch.isfinite(cup_ang_vel).all(dim=-1) + ) + media_ok = torch.isfinite(particle_pos).all(dim=(1, 2)) + return robot_ok & tcp_ok & cup_ok & media_ok + + +def _rigid_state_in_bounds( + robot_joint_pos: torch.Tensor, + robot_joint_vel: torch.Tensor, + joint_pos_limits: torch.Tensor, + tcp_body_q: torch.Tensor, + cup_body_q: torch.Tensor, + cup_lin_vel: torch.Tensor, + cup_ang_vel: torch.Tensor, + env_origins: torch.Tensor, + lower_bound: tuple[float, float, float] | torch.Tensor, + upper_bound: tuple[float, float, float] | torch.Tensor, + joint_position_margin: float, + max_joint_velocity: float, + max_cup_linear_velocity: float, + max_cup_angular_velocity: float, +) -> torch.Tensor: + """Return whether every rigid state that feeds actor observations is physically bounded.""" + lower = torch.as_tensor(lower_bound, device=robot_joint_pos.device, dtype=robot_joint_pos.dtype) + upper = torch.as_tensor(upper_bound, device=robot_joint_pos.device, dtype=robot_joint_pos.dtype) + joint_lower = joint_pos_limits[..., 0] - float(joint_position_margin) + joint_upper = joint_pos_limits[..., 1] + float(joint_position_margin) + joint_position_ok = ((robot_joint_pos >= joint_lower) & (robot_joint_pos <= joint_upper)).all(dim=-1) + joint_velocity_ok = (torch.abs(robot_joint_vel) <= float(max_joint_velocity)).all(dim=-1) + + tcp_position = tcp_body_q[:, :3] - env_origins + cup_position = cup_body_q[:, :3] - env_origins + tcp_position_ok = ((tcp_position >= lower) & (tcp_position <= upper)).all(dim=-1) + cup_position_ok = ((cup_position >= lower) & (cup_position <= upper)).all(dim=-1) + tcp_quat_norm = torch.linalg.vector_norm(tcp_body_q[:, 3:7], dim=-1) + cup_quat_norm = torch.linalg.vector_norm(cup_body_q[:, 3:7], dim=-1) + pose_ok = ( + torch.isfinite(tcp_body_q).all(dim=-1) + & torch.isfinite(cup_body_q).all(dim=-1) + & (torch.abs(tcp_quat_norm - 1.0) <= 0.1) + & (torch.abs(cup_quat_norm - 1.0) <= 0.1) + ) + + cup_linear_velocity_ok = torch.linalg.vector_norm(cup_lin_vel, dim=-1) <= float(max_cup_linear_velocity) + cup_angular_velocity_ok = torch.linalg.vector_norm(cup_ang_vel, dim=-1) <= float(max_cup_angular_velocity) + return ( + joint_position_ok + & joint_velocity_ok + & tcp_position_ok + & cup_position_ok + & pose_ok + & cup_linear_velocity_ok + & cup_angular_velocity_ok + ) + + +def _particles_in_workspace( + particle_pos_e: torch.Tensor, + lower_bound: tuple[float, float, float] | torch.Tensor, + upper_bound: tuple[float, float, float] | torch.Tensor, +) -> torch.Tensor: + """Return whether every particle lies inside its environment-local workspace.""" + lower = torch.as_tensor(lower_bound, device=particle_pos_e.device, dtype=particle_pos_e.dtype) + upper = torch.as_tensor(upper_bound, device=particle_pos_e.device, dtype=particle_pos_e.dtype) + return ((particle_pos_e >= lower) & (particle_pos_e <= upper)).all(dim=(1, 2)) + + +def _spilled_particle_mask( + particle_pos_e: torch.Tensor, + in_source: torch.Tensor, + in_target: torch.Tensor, + max_height: float, +) -> torch.Tensor: + """Classify particles outside both cups that have reached the table [m].""" + outside_cups = ~in_source & ~in_target + return outside_cups & (particle_pos_e[..., 2] <= float(max_height)) + + +def _delivered_particle_mask(in_source: torch.Tensor, in_target: torch.Tensor) -> torch.Tensor: + """Classify particles inside the receiver only after they have left the source cup.""" + return in_target & ~in_source + + +def nonfinite_failure(env: FrankaPourEnv) -> torch.Tensor: + """Terminate on non-finite simulation state (instability guard).""" + return ~env.state_finite() + + +def extreme_rigid_state(env: FrankaPourEnv) -> torch.Tensor: + """Terminate finite rigid states before extreme values reach observation normalization.""" + return ~env.rigid_state_in_bounds() + + +def particle_out_of_bounds(env: FrankaPourEnv) -> torch.Tensor: + """Terminate an environment when any media particle escapes its workspace.""" + return ~env.particles_in_workspace() + + +def excessive_spill(env: FrankaPourEnv, terminate: bool = True) -> torch.Tensor: + """Track when strictly more than the allowed media fraction is spilled.""" + if not isinstance(terminate, bool): + raise TypeError("terminate must be a bool.") + spilled = env.spilled_fraction() > float(env.cfg.max_spill_fraction) + return spilled if terminate else torch.zeros_like(spilled) + + +def unsuccessful_time_out(env: FrankaPourEnv) -> torch.Tensor: + """Truncate at the finite-horizon deadline unless success fired on the same step.""" + deadline = env.episode_length_buf >= env.max_episode_length + return deadline & ~env.episode_succeeded + + +def lost_lifted_grasp( + env: FrankaPourEnv, + dwell_time_s: float = 0.05, + max_tcp_distance: float = 0.015, + max_gripper_width_error: float = 0.012, + max_gripper_command: float = 0.025, + terminate: bool = True, +) -> torch.Tensor: + """Track when a demonstrated lift loses the physical cup grasp continuously. + + A short dwell rejects isolated contact-deflection flicker from the coupled rigid solver while + still allowing the reverse curriculum to terminate a genuinely dropped cup promptly. Static + reset-dataset training can disable termination while retaining the state monitor for metrics. + """ + if not math.isfinite(dwell_time_s) or dwell_time_s <= 0.0: + raise ValueError(f"dwell_time_s must be finite and positive, got {dwell_time_s}.") + if not isinstance(terminate, bool): + raise TypeError("terminate must be a bool.") + _, preloaded_grasp, lifted_grasp = source_grasp_milestones( + env, + min_lift_height=max(float(env.cfg.success_min_lift_height), 1.0e-6), + max_tcp_distance=max_tcp_distance, + max_gripper_width_error=max_gripper_width_error, + max_gripper_command=max_gripper_command, + ) + env._lifted_grasp_seen |= lifted_grasp + lost = env._lifted_grasp_seen & ~preloaded_grasp + dwell_steps = max(1, math.ceil(float(dwell_time_s) / max(float(env.step_dt), 1.0e-6))) + env._lost_grasp_dwell_count[:] = torch.where( + lost, + torch.clamp(env._lost_grasp_dwell_count + 1, max=dwell_steps), + 0, + ) + dwell_qualified_loss = lost & (env._lost_grasp_dwell_count >= dwell_steps) + return dwell_qualified_loss if terminate else torch.zeros_like(dwell_qualified_loss) + + +def source_grasp_milestones( + env: FrankaPourEnv, + min_lift_height: float, + max_tcp_distance: float, + max_gripper_width_error: float, + max_gripper_command: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return reached, actively preloaded, and lifted source-grasp masks.""" + cup_lift = env.cup_pose_e()[:, 2] - float(env.cup_reset_height) + tcp_distance = torch.linalg.vector_norm(env.tcp_pos_e() - env.cup_grasp_point_e(), dim=-1) + gripper_width_error = torch.abs(env.gripper_width() - float(env.gripper_grasp_width)) + gripper = env.action_manager.get_term("gripper_action") + gripper_command = gripper.commanded_position[:, 0] + bilateral_contact = getattr( + gripper, + "bilateral_contact", + torch.ones(env.num_envs, device=env.device, dtype=torch.bool), + ) + reached_grasp = tcp_distance <= float(max_tcp_distance) + preloaded_grasp = ( + reached_grasp + & (gripper_width_error <= float(max_gripper_width_error)) + & (gripper_command <= float(max_gripper_command) + _GRIPPER_POSITION_TOLERANCE) + & bilateral_contact + ) + lifted_grasp = preloaded_grasp & (cup_lift >= float(min_lift_height)) + return reached_grasp, preloaded_grasp, lifted_grasp + + +def stable_pour_success( + env: FrankaPourEnv, + dwell_time_s: float = 0.15, + min_lift_height: float = 0.05, + max_tcp_distance: float = 0.015, + max_gripper_width_error: float = 0.012, + max_gripper_command: float = 0.025, +) -> torch.Tensor: + """Terminate after a delivered pour remains held, preloaded, and lifted for a dwell interval. + + This is deliberately a plain manager function rather than a class term. Isaac Lab's generic + record/replay tools remove the success termination from the manager and invoke its configured + function directly; keeping the per-world counter on the task preserves that standard workflow. + The success term must follow every failure predicate so ``terminated`` contains them, and must + precede :func:`unsuccessful_time_out` so a valid deadline transfer is classified only as + success. + """ + if not math.isfinite(dwell_time_s) or dwell_time_s <= 0.0: + raise ValueError(f"dwell_time_s must be finite and positive, got {dwell_time_s}.") + for name, value in ( + ("min_lift_height", min_lift_height), + ("max_tcp_distance", max_tcp_distance), + ("max_gripper_width_error", max_gripper_width_error), + ): + if not math.isfinite(value) or value <= 0.0: + raise ValueError(f"{name} must be finite and positive, got {value}.") + if not math.isfinite(max_gripper_command) or max_gripper_command < 0.0: + raise ValueError(f"max_gripper_command must be finite and nonnegative, got {max_gripper_command}.") + dwell_steps = max(1, math.ceil(float(dwell_time_s) / max(float(env.step_dt), 1.0e-6))) + target_fraction = env.count_in_target() / max(env.num_particles, 1) + env.ep_max_target_frac[:] = torch.maximum(env.ep_max_target_frac, target_fraction) + _, _, held_pour = source_grasp_milestones( + env, + min_lift_height=min_lift_height, + max_tcp_distance=max_tcp_distance, + max_gripper_width_error=max_gripper_width_error, + max_gripper_command=max_gripper_command, + ) + env.update_held_delivery_tracker(held_pour) + held_target_fraction = env.current_held_delivered_mask().sum(dim=1).float() / max(env.num_particles, 1) + within_spill_limit = env.spilled_fraction() <= float(env.cfg.max_spill_fraction) + candidate = ( + (held_target_fraction >= env.pour_target_frac) + & held_pour + & within_spill_limit + & ~env.termination_manager.terminated + ) + env._success_dwell_count[:] = torch.where( + candidate, + torch.clamp(env._success_dwell_count + 1, max=dwell_steps), + 0, + ) + success = candidate & (env._success_dwell_count >= dwell_steps) + env.episode_succeeded |= success + return success + + +def immediate_pour_success(env: FrankaPourEnv) -> torch.Tensor: + """Terminate immediately when the current target-bowl fraction reaches its threshold. + + Reset-cache particles always start in the source cup, so current target occupancy directly + measures task progress without also requiring a retained grasp, lift milestone, delivery + history, or dwell interval. Failure terms run before success and therefore retain precedence. + """ + target_fraction = env.count_in_target() / max(env.num_particles, 1) + env.ep_max_target_frac[:] = torch.maximum(env.ep_max_target_frac, target_fraction) + success = (target_fraction >= env.pour_target_frac) & ~env.termination_manager.terminated + env._success_dwell_count[:] = success.to(dtype=env._success_dwell_count.dtype) + env.episode_succeeded |= success + return success + + +def nonterminating_stable_pour_success( + env: FrankaPourEnv, + dwell_time_s: float = 0.15, + min_lift_height: float = 0.05, + max_tcp_distance: float = 0.015, + max_gripper_width_error: float = 0.012, + max_gripper_command: float = 0.025, +) -> torch.Tensor: + """Track stable success without ending training episodes, while remaining replay-compatible.""" + success = stable_pour_success( + env, + dwell_time_s=dwell_time_s, + min_lift_height=min_lift_height, + max_tcp_distance=max_tcp_distance, + max_gripper_width_error=max_gripper_width_error, + max_gripper_command=max_gripper_command, + ) + # Standard record/replay tools temporarily remove the managed success term and invoke this + # configured function directly. Return the real predicate in that context; suppress only the + # live TerminationManager's done signal used by fixed-horizon reset-dataset training. + if "success" not in env.termination_manager.active_terms: + return success + return torch.zeros_like(success) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/media_fill.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/media_fill.py new file mode 100644 index 000000000000..dab1ca57ddea --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/media_fill.py @@ -0,0 +1,103 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Granular-media fill for the Franka pour source bowl. + +:func:`cube_fill_points` builds a deterministic, jittered axis-aligned lattice +clipped to the analytic hollow-cube bowl's inner cavity. + +Particles must be seeded as a jittered lattice at ``spacing = voxel_size / particles_per_cell`` and +inset from the walls by ``clearance``: a particle spawned inside the grid-level collider shell is +ejected on the first solve, and any overlap explodes at the near-incompressible MPM stiffness. +""" + +from __future__ import annotations + +import numpy as np + +# Default wall inset: at least one particle spacing, and clear of the collider margin band. +_DEFAULT_MARGIN = 0.002 + + +def _resolve_clearance(spacing: float, clearance: float | None) -> float: + return float(clearance) if clearance is not None else max(float(spacing), 3.0 * _DEFAULT_MARGIN) + + +def _fill_region( + inner_lo: np.ndarray, + inner_hi: np.ndarray, + spacing: float, + fill_frac: float, + clearance: float, +) -> tuple[np.ndarray, np.ndarray]: + """Return the inset ``(region_lo, region_hi)`` the lattice fills.""" + inner_lo = np.asarray(inner_lo, dtype=np.float64) + inner_hi = np.asarray(inner_hi, dtype=np.float64) + cavity_h = float(inner_hi[2] - inner_lo[2]) + fill_depth = max(0.0, min(float(fill_frac) * cavity_h, cavity_h - 2.0 * clearance)) + region_lo = np.array([inner_lo[0] + clearance, inner_lo[1] + clearance, inner_lo[2] + clearance]) + region_hi = np.array([inner_hi[0] - clearance, inner_hi[1] - clearance, inner_lo[2] + clearance + fill_depth]) + return region_lo, region_hi + + +def _axis_samples(lo: float, hi: float, spacing: float) -> np.ndarray: + """Regularly spaced samples in ``[lo, hi]`` (at least one, centred when the span is short).""" + span = hi - lo + if span <= 0.0: + return np.array([0.5 * (lo + hi)], dtype=np.float64) + n = int(np.floor(span / spacing)) + coords = lo + spacing * np.arange(n + 1, dtype=np.float64) + # Centre the lattice in the span so both walls get equal clearance. + coords = coords + 0.5 * (span - spacing * n) + return coords + + +def expected_fill_count( + inner_lo: np.ndarray, + inner_hi: np.ndarray, + spacing: float, + fill_frac: float = 1.0, + clearance: float | None = None, +) -> int: + """Analytic particle count :func:`cube_fill_points` will produce (for sizing/asserts).""" + clr = _resolve_clearance(spacing, clearance) + region_lo, region_hi = _fill_region(inner_lo, inner_hi, spacing, fill_frac, clr) + counts = [len(_axis_samples(region_lo[a], region_hi[a], float(spacing))) for a in range(3)] + return int(counts[0] * counts[1] * counts[2]) + + +def cube_fill_points( + inner_lo: np.ndarray, + inner_hi: np.ndarray, + spacing: float, + fill_frac: float = 1.0, + clearance: float | None = None, + jitter: float = 0.05, + seed: int = 7, +) -> np.ndarray: + """Build a jittered lattice of particle positions filling a box cavity. + + Args: + inner_lo: Cavity floor corner ``(3,)`` [m] (e.g. from + :func:`.cube_bowl_mesh.cube_bowl_inner_bounds`). + inner_hi: Cavity rim corner ``(3,)`` [m]. + spacing: Particle lattice spacing [m] (``voxel_size / particles_per_cell``). + fill_frac: Fraction of the cavity height to fill (1.0 = up to the rim, capped to leave + ``clearance`` below the rim). + clearance: Wall/floor inset [m]; defaults to ``max(spacing, 3 * 0.002)``. + jitter: Uniform per-particle jitter as a fraction of ``spacing`` (0 = a perfect lattice). + seed: RNG seed; identical ``seed`` gives identical points (per-env determinism). + + Returns: + ``(K, 3)`` float32 particle positions in the bowl local frame. + """ + clr = _resolve_clearance(spacing, clearance) + region_lo, region_hi = _fill_region(inner_lo, inner_hi, spacing, fill_frac, clr) + axes = [_axis_samples(region_lo[a], region_hi[a], float(spacing)) for a in range(3)] + grid = np.stack(np.meshgrid(axes[0], axes[1], axes[2], indexing="ij"), axis=-1).reshape(-1, 3) + if jitter > 0.0: + rng = np.random.default_rng(int(seed)) + grid = grid + (rng.random(grid.shape) - 0.5) * 2.0 * float(jitter) * float(spacing) + return grid.astype(np.float32) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/pour_env.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/pour_env.py new file mode 100644 index 000000000000..6cb157a04c40 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/pour_env.py @@ -0,0 +1,2610 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Manager-based RL environment for a Franka pouring MPM media between two cups. + +The visible dynamic source cup and kinematic receiving cup are scene-owned rigid objects. Their +USD bowl meshes are visual-only, while the source also owns an invisible rigid grasp proxy for +Newton-generated finger contacts. A narrow per-world Newton hook attaches cached hollow +particle-only colliders to both scene bodies and adds only one hidden solver object: a particle-only +spill floor. + +A Newton :class:`~isaaclab_contrib.coupling.CouplerProxyCfg` advances the robot and both cups +in the ``arm`` MJWarp entry and the particles and spill floor in the implicit ``media`` entry. Proxy +coupling makes both cups' particle colliders available to MPM without assigning one body to two +entries. The policy commands arm joint positions and a continuous symmetric finger target; all +observable and reset state flows through the scene assets' public APIs. +""" + +from __future__ import annotations + +import logging +import math +from pathlib import Path +from typing import TYPE_CHECKING + +import newton +import numpy as np +import torch +import warp as wp +from isaaclab_newton.cloner import copy_newton_source_builder, newton_builder_world_hook +from isaaclab_newton.ik.newton_ik_objectives_cfg import NewtonIKJointLimitObjectiveCfg, NewtonIKPoseObjectiveCfg +from isaaclab_newton.ik.newton_ik_solver import NewtonIKSolver +from isaaclab_newton.ik.newton_ik_solver_cfg import NewtonIKSolverCfg +from isaaclab_newton.physics import NewtonManager + +import isaaclab.sim as sim_utils +from isaaclab.cloner import resolve_clone_plan_source +from isaaclab.envs import ManagerBasedRLEnv +from isaaclab.utils import math as math_utils + +from .cube_bowl_mesh import cube_bowl_inner_bounds, make_cube_bowl_mesh +from .cup_media import cup_cavity_lattice +from .mdp.terminations import ( + _delivered_particle_mask, + _particles_in_workspace, + _rigid_state_in_bounds, + _spilled_particle_mask, + _state_finite, +) +from .reset_dataset_generator import ( + GRASPING_CATEGORY, + build_franka_pour_reset_task_contract, + validate_production_reset_dataset, + validate_reset_dataset, +) +from .reset_utils import ( + asymmetric_reset_offset_samples, + balanced_cyclic_permutations, + boolean_selection_mask, + polar_workspace_cells, + reset_rotation_vector_samples, + sample_index_pools, + scale_randomization_rows_by_extent, + target_xy_behind_source, +) + +logger = logging.getLogger(__name__) + +if TYPE_CHECKING: + from isaaclab_newton.assets import MPMObject + + from .pour_env_cfg import FrankaPourEnvCfg + +ARM_JOINTS = [f"panda_joint{i}" for i in range(1, 8)] +FINGER_JOINTS = ["panda_finger_joint1", "panda_finger_joint2"] +_RESET_PATH_SEGMENT_SUBDIVISIONS = 8 + + +@wp.kernel(enable_backward=False) +def _mark_penetrating_self_contacts( + contact_count: wp.array(dtype=wp.int32), + contact_max: int, + contact_shape0: wp.array(dtype=wp.int32), + contact_shape1: wp.array(dtype=wp.int32), + contact_point0: wp.array(dtype=wp.vec3), + contact_point1: wp.array(dtype=wp.vec3), + contact_normal: wp.array(dtype=wp.vec3), + contact_margin0: wp.array(dtype=wp.float32), + contact_margin1: wp.array(dtype=wp.float32), + shape_body: wp.array(dtype=wp.int32), + body_world: wp.array(dtype=wp.int32), + body_q: wp.array(dtype=wp.transform), + penetration_tolerance: float, + colliding_worlds: wp.array(dtype=wp.int32), +): + """Mark replicated IK candidates containing a penetrating robot self-contact.""" + contact_index = wp.tid() + if contact_index >= contact_max or contact_index >= contact_count[0]: + return + shape0 = contact_shape0[contact_index] + shape1 = contact_shape1[contact_index] + body0 = shape_body[shape0] + body1 = shape_body[shape1] + if body0 < 0 or body1 < 0: + return + world0 = body_world[body0] + world1 = body_world[body1] + if world0 < 0 or world0 != world1: + return + point0_w = wp.transform_point(body_q[body0], contact_point0[contact_index]) + point1_w = wp.transform_point(body_q[body1], contact_point1[contact_index]) + separation = wp.dot(contact_normal[contact_index], point1_w - point0_w) + separation = separation - contact_margin0[contact_index] - contact_margin1[contact_index] + if separation < -penetration_tolerance: + wp.atomic_max(colliding_worlds, world0, 1) + + +class FrankaPourEnv(ManagerBasedRLEnv): + """Franka grasping a dynamic cup of MPM media (Newton proxy-coupled MPM), pouring by tilting.""" + + cfg: FrankaPourEnvCfg + + def __init__(self, cfg: FrankaPourEnvCfg, render_mode: str | None = None, **kwargs): + resolved_cfg = cfg.finalize() + self._prepare_newton_extras(resolved_cfg) + with newton_builder_world_hook(self._add_pour_world_to_builder): + super().__init__(resolved_cfg, render_mode, **kwargs) + + def load_managers(self) -> None: + self._setup_after_physics() + super().load_managers() + + # ------------------------------------------------------------------ build + def _prepare_newton_extras(self, cfg: FrankaPourEnvCfg) -> None: + """Bake task-local Newton collision geometry from the resolved scene config. + + Runs before ``super().__init__`` so the per-world builder hook has the + geometry and contact values available while the scene is imported. + """ + # Watertight cube-cup collision meshes. The source's outer extents exactly match the solid + # grasp box, so its visible walls and the rigid finger contacts no longer disagree. + self._cup_vertices, self._cup_indices = make_cube_bowl_mesh( + inner_width=float(cfg.source_cup_inner_width), + inner_depth=float(cfg.source_cup_inner_depth), + wall_thickness=float(cfg.source_cup_wall_thickness), + cavity_depth=float(cfg.source_cup_cavity_depth), + bottom_thickness=float(cfg.source_cup_bottom_thickness), + ) + self._target_vertices, self._target_indices = make_cube_bowl_mesh( + inner_width=float(cfg.target_cup_inner_width), + inner_depth=float(cfg.target_cup_inner_depth), + wall_thickness=float(cfg.target_cup_wall_thickness), + cavity_depth=float(cfg.target_cup_cavity_depth), + bottom_thickness=float(cfg.target_cup_bottom_thickness), + ) + self._source_collider_mesh = newton.Mesh( + self._cup_vertices, + self._cup_indices, + compute_inertia=False, + is_solid=False, + ) + self._target_collider_mesh = newton.Mesh( + self._target_vertices, + self._target_indices, + compute_inertia=False, + is_solid=False, + ) + self._source_inner_lo, self._source_inner_hi = cube_bowl_inner_bounds( + cfg.source_cup_inner_width, + cfg.source_cup_inner_depth, + cfg.source_cup_cavity_depth, + cfg.source_cup_bottom_thickness, + ) + self._target_inner_lo, self._target_inner_hi = cube_bowl_inner_bounds( + cfg.target_cup_inner_width, + cfg.target_cup_inner_depth, + cfg.target_cup_cavity_depth, + cfg.target_cup_bottom_thickness, + ) + + # Cup-local media lattice (the env transforms it by the live cup pose every reset). + self._media_local_points, _ = cup_cavity_lattice(cfg) + + # Cup reset pose (env frame): resting on the table, opening up (cup-local +z is world +z). + self._cup_reset_pos = np.asarray(cfg.cup_reset_pos, dtype=np.float64) + self._grasp_contact_ke = float(cfg.grasp_contact_ke) + self._grasp_contact_kd = float(cfg.grasp_contact_kd) + self._grasp_contact_kf = float(cfg.grasp_contact_kf) + self._grasp_contact_mu = float(cfg.cup_grasp_box_friction) + + self._source_cup_friction = float(cfg.source_cup_friction) + self._target_cup_friction = float(cfg.target_cup_friction) + self._collider_margin = float(cfg.collider_margin) + self._particle_max_velocity = float(cfg.particle_max_velocity) + + def _add_pour_world_to_builder(self, builder, env_id: int, position, quaternion) -> None: + """Add only solver-specific collision representations to one imported scene world.""" + builder.particle_max_velocity = self._particle_max_velocity + env_root = f"/World/envs/env_{env_id}" + body_ids = self._current_world_range(builder, "body", env_id) + shape_ids = self._current_world_range(builder, "shape", env_id) + self._disable_robot_particle_collision(builder, body_ids, shape_ids) + self._configure_finger_contact_material(builder, body_ids, shape_ids) + + source_body = self._find_world_body(builder, body_ids, env_id, "SourceCup") + target_body = self._find_world_body(builder, body_ids, env_id, "TargetCup") + self._add_kinematic_rigid_object_articulation(builder, target_body) + grasp_proxy = self._find_world_shape( + builder, + shape_ids, + env_id, + "/SourceCup/geometry/grasp_proxy", + body_id=source_body, + ) + self._configure_grasp_proxy(builder, grasp_proxy) + + self._add_particle_collider( + builder, + body_id=source_body, + mesh=self._source_collider_mesh, + friction=self._source_cup_friction, + label=f"{env_root}/SourceCup/ParticleCollider", + ) + self._add_particle_collider( + builder, + body_id=target_body, + mesh=self._target_collider_mesh, + friction=self._target_cup_friction, + label=f"{env_root}/TargetCup/ParticleCollider", + ) + + self._add_rigid_collider( + builder, + body_id=target_body, + mesh=self._target_collider_mesh, + friction=self._target_cup_friction, + label=f"{env_root}/TargetCup/Collision", + ) + + world_xform = wp.transform( + wp.vec3(*[float(value) for value in position]), + wp.quat(*[float(value) for value in quaternion]), + ) + spill_floor = builder.add_body( + xform=world_xform, + mass=0.0, + inertia=wp.mat33(), + is_kinematic=True, + lock_inertia=True, + label=f"{env_root}/SpillFloor", + ) + spill_shape = builder.add_shape_plane( + body=spill_floor, + xform=wp.transform_identity(), + width=0.0, + length=0.0, + cfg=newton.ModelBuilder.ShapeConfig( + mu=0.8, + margin=self._collider_margin, + has_shape_collision=False, + has_particle_collision=True, + ), + color=(0.3, 0.3, 0.3), + label=f"{env_root}/SpillFloor/Collision", + ) + self._set_shape_roles(builder, spill_shape, rigid=False, particles=True, visible=False) + + @staticmethod + def _current_world_range(builder, prefix: str, env_id: int) -> range: + """Return the contiguous tail added for the currently open Newton world.""" + worlds = getattr(builder, f"{prefix}_world", None) + if worlds is None: + raise RuntimeError(f"Newton builder does not expose {prefix}_world assignments.") + stop = len(worlds) + start = stop + while start > 0 and int(worlds[start - 1]) == env_id: + start -= 1 + if start == stop: + raise RuntimeError(f"Newton builder contains no {prefix} entries for open world {env_id}.") + return range(start, stop) + + def _find_world_body(self, builder, body_ids: range, env_id: int, body_name: str) -> int: + """Resolve exactly one imported body by Newton world and exact final path component.""" + matches = [body_id for body_id in body_ids if str(builder.body_label[body_id]).rsplit("/", 1)[-1] == body_name] + if len(matches) != 1: + labels = [str(builder.body_label[index]) for index in matches] + raise RuntimeError( + f"Expected exactly one {body_name!r} body in Newton world {env_id}, " + f"found ids={matches}, labels={labels}." + ) + return matches[0] + + @staticmethod + def _add_kinematic_rigid_object_articulation(builder, body_id: int) -> None: + """Expose an imported kinematic body through Newton's articulation-based rigid view.""" + body_label = str(builder.body_label[body_id]) + child_joints = [joint_id for _, joint_id in builder.joint_parents.get(body_id, ())] + if not child_joints: + joint_id = builder.add_joint_free(child=body_id, label=f"{body_label}/FreeJoint") + builder.add_articulation([joint_id], label=body_label) + elif len(child_joints) == 1: + joint_id = child_joints[0] + articulation_id = int(builder.joint_articulation[joint_id]) + if articulation_id < 0 or str(builder.articulation_label[articulation_id]) != body_label: + raise RuntimeError( + f"Kinematic rigid body {body_label!r} has an unexpected joint/articulation association." + ) + else: + raise RuntimeError( + f"Kinematic rigid body {body_label!r} must have at most one root joint, found {child_joints}." + ) + + builder.body_flags[body_id] = int(newton.BodyFlags.KINEMATIC) + builder.body_mass[body_id] = 0.0 + builder.body_inv_mass[body_id] = 0.0 + builder.body_inertia[body_id] = wp.mat33() + builder.body_inv_inertia[body_id] = wp.mat33() + + def _find_world_shape(self, builder, shape_ids: range, env_id: int, label_suffix: str, *, body_id: int) -> int: + """Resolve exactly one imported shape by owning body and exact scene-relative path.""" + matches = [ + shape_id + for shape_id in shape_ids + if int(builder.shape_body[shape_id]) == body_id + and str(builder.shape_label[shape_id]).endswith(label_suffix) + ] + if len(matches) != 1: + labels = [str(builder.shape_label[index]) for index in matches] + raise RuntimeError( + f"Expected exactly one shape ending in {label_suffix!r} on body {body_id} " + f"in Newton world {env_id}, found ids={matches}, labels={labels}." + ) + return matches[0] + + def _configure_grasp_proxy(self, builder, shape_id: int) -> None: + """Keep the imported grasp proxy rigid-only, invisible, and contact-tuned.""" + self._set_shape_roles(builder, shape_id, rigid=True, particles=False, visible=False) + builder.shape_margin[shape_id] = self._collider_margin + builder.shape_material_ke[shape_id] = self._grasp_contact_ke + builder.shape_material_kd[shape_id] = self._grasp_contact_kd + builder.shape_material_kf[shape_id] = self._grasp_contact_kf + builder.shape_material_mu[shape_id] = self._grasp_contact_mu + + def _add_particle_collider( + self, + builder, + *, + body_id: int, + mesh: newton.Mesh, + friction: float, + label: str, + ) -> int: + """Attach an invisible hollow particle-only collider to a scene-owned body.""" + shape_id = builder.add_shape_mesh( + body_id, + xform=wp.transform_identity(), + mesh=mesh, + cfg=newton.ModelBuilder.ShapeConfig( + mu=friction, + density=0.0, + margin=self._collider_margin, + has_shape_collision=False, + has_particle_collision=True, + is_visible=False, + ), + label=label, + ) + self._set_shape_roles(builder, shape_id, rigid=False, particles=True, visible=False) + builder.shape_margin[shape_id] = self._collider_margin + builder.shape_material_mu[shape_id] = friction + return shape_id + + def _add_rigid_collider( + self, + builder, + *, + body_id: int, + mesh: newton.Mesh, + friction: float, + label: str, + ) -> int: + """Attach an invisible hollow rigid-only collider to a solver-owned body.""" + shape_id = builder.add_shape_mesh( + body_id, + xform=wp.transform_identity(), + mesh=mesh, + cfg=newton.ModelBuilder.ShapeConfig( + mu=friction, + density=0.0, + ke=self._grasp_contact_ke, + kd=self._grasp_contact_kd, + kf=self._grasp_contact_kf, + margin=self._collider_margin, + has_shape_collision=True, + has_particle_collision=False, + is_visible=False, + ), + label=label, + ) + self._set_shape_roles(builder, shape_id, rigid=True, particles=False, visible=False) + builder.shape_margin[shape_id] = self._collider_margin + builder.shape_material_mu[shape_id] = friction + return shape_id + + @staticmethod + def _set_shape_roles(builder, shape_id: int, *, rigid: bool, particles: bool, visible: bool) -> None: + flags = int(builder.shape_flags[shape_id]) + assignments = ( + (newton.ShapeFlags.COLLIDE_SHAPES, rigid), + (newton.ShapeFlags.COLLIDE_PARTICLES, particles), + (newton.ShapeFlags.VISIBLE, visible), + ) + for flag, enabled in assignments: + if enabled: + flags |= int(flag) + else: + flags &= ~int(flag) + builder.shape_flags[shape_id] = flags + + def _disable_robot_particle_collision(self, builder, body_ids: range, shape_ids: range) -> None: + """The robot shapes must not collide with MPM particles (only the cup cavity mesh does).""" + collide_particles = int(newton.ShapeFlags.COLLIDE_PARTICLES) + for shape_id in shape_ids: + body_id = int(builder.shape_body[shape_id]) + if body_id not in body_ids: + continue + body_label = str(builder.body_label[body_id]) + if "/Robot/" in body_label or body_label.endswith("/Robot"): + builder.shape_flags[shape_id] &= ~collide_particles + + def _configure_finger_contact_material(self, builder, body_ids: range, shape_ids: range) -> None: + """Use a rigid contact material on the two finger collision shapes. + + Applying the same material to the fingers and cup keeps the intended pair response + independent of import-side material defaults. + """ + finger_suffixes = ("/panda_leftfinger", "/panda_rightfinger") + for shape_id in shape_ids: + body_id = int(builder.shape_body[shape_id]) + if body_id not in body_ids: + continue + if str(builder.body_label[body_id]).endswith(finger_suffixes): + builder.shape_material_ke[shape_id] = self._grasp_contact_ke + builder.shape_material_kd[shape_id] = self._grasp_contact_kd + builder.shape_material_kf[shape_id] = self._grasp_contact_kf + builder.shape_material_mu[shape_id] = self._grasp_contact_mu + + # ----------------------------------------------------------- post-physics + def _setup_after_physics(self) -> None: + dev = self.device + self._robot = self.scene["robot"] + self._source_cup = self.scene["source_cup"] + self._target_cup = self.scene["target_cup"] + self._media: MPMObject = self.scene["media"] + + self._arm_joint_ids, _ = self._robot.find_joints(ARM_JOINTS, preserve_order=True) + self._finger_joint_ids, _ = self._robot.find_joints(FINGER_JOINTS, preserve_order=True) + self._joint_pos_limits_t = self._robot.data.joint_pos_limits.torch.clone() + tcp_body_ids, _ = self._robot.find_bodies(self.cfg.tcp_body_name) + if len(tcp_body_ids) != 1: + raise RuntimeError( + f"Expected one TCP parent body named {self.cfg.tcp_body_name!r}, found {len(tcp_body_ids)}." + ) + self._tcp_body_idx = tcp_body_ids[0] + self._tcp_offset_pos = torch.tensor(self.cfg.tcp_offset_pos, device=dev).repeat(self.num_envs, 1) + self._tcp_offset_quat = torch.tensor(self.cfg.tcp_offset_rot, device=dev).repeat(self.num_envs, 1) + approach_axis_c = torch.as_tensor( + self.cfg.curriculum_randomized_reset_tcp_standoff, + device=dev, + dtype=torch.float32, + ) + self._grasp_approach_axis_c = approach_axis_c / torch.linalg.vector_norm(approach_axis_c) + grasp_tcp_quat_c = torch.as_tensor( + self.cfg.cup_grasp_tcp_quat_c, + device=dev, + dtype=torch.float32, + ) + # Keep the grasp frame explicit: tool +Z and jaw +Y remain parallel to the table, and the + # same cup-local grasp is preserved under source-yaw randomization. + self._desired_grasp_tcp_quat_c = math_utils.quat_unique(grasp_tcp_quat_c).repeat(self.num_envs, 1) + + self.env_origins = self.scene.env_origins.to(device=dev, dtype=torch.float32) + self._num_particles = int(self._media.particles_per_object) + self._media_local_points_t = torch.as_tensor(self._media_local_points, device=dev, dtype=torch.float32) + self._particle_workspace_lower_t = torch.as_tensor( + self.cfg.particle_workspace_lower_bound, device=dev, dtype=torch.float32 + ) + self._particle_workspace_upper_t = torch.as_tensor( + self.cfg.particle_workspace_upper_bound, device=dev, dtype=torch.float32 + ) + grasp_stage_index = self.cfg.curriculum_stage_names.index("grasp") + self._curriculum_arm_q_t = torch.as_tensor( + ( + self.cfg.curriculum_drain_arm_q, + self.cfg.curriculum_deep_tilt_arm_q, + self.cfg.curriculum_tilt_arm_q, + self.cfg.curriculum_pour_arm_q, + *self.cfg._curriculum_transport_arm_configs(), + self.cfg.curriculum_carry_arm_q, + *(self.cfg.arm_home for _ in range(len(self.cfg.curriculum_stage_names) - grasp_stage_index)), + ), + device=dev, + dtype=torch.float32, + ) + self._curriculum_cup_quat_t = torch.zeros((len(self.cfg.curriculum_stage_names), 4), device=dev) + self._curriculum_cup_quat_t[:, 3] = 1.0 + contact_position = float(self.cfg.cup_grasp_box_half[1]) + self._curriculum_finger_pos_t = torch.full( + (len(self.cfg.curriculum_stage_names),), + float(self.cfg.gripper_open_pos), + device=dev, + ) + self._curriculum_finger_pos_t[:grasp_stage_index] = contact_position + # These waypoints seed and validate the reset bank only. The policy never receives or + # follows them; every runtime arm command is a direct relative joint-position action. + self._nominal_reference_waypoints_t = torch.as_tensor( + ( + self.cfg.arm_home, + self.cfg.arm_home, + self.cfg.arm_home, + self.cfg.arm_home, + self.cfg.curriculum_carry_arm_q, + self.cfg.curriculum_pour_arm_q, + self.cfg.curriculum_pour_target_arm_q, + ), + device=dev, + dtype=torch.float32, + ) + self._grasp_stage_index = grasp_stage_index + self._approach_stage_index = self.cfg.curriculum_stage_names.index("approach_1") + self._full_stage_index = self.cfg.curriculum_stage_names.index("full") + self._randomized_stage_index = self.cfg.curriculum_stage_names.index("randomized") + self._uses_reset_dataset = getattr(self.cfg.curriculum, "reset_dataset", None) is not None + reset_dataset_path = getattr(self.cfg, "reset_dataset_path", None) + if self._uses_reset_dataset: + if reset_dataset_path is None: + raise ValueError("The reset-dataset task requires reset_dataset_path.") + self._load_reset_dataset(reset_dataset_path) + else: + self._build_randomized_reset_bank() + # The complete-path collision screen samples the midpoint-to-grasp joint segment at + # eighths. Reuse those exact samples as progressively longer open-hand reset frontiers. + approach_fractions = torch.as_tensor( + self.cfg.curriculum_grasp_approach_fractions, + device=dev, + dtype=torch.float32, + ) + self._curriculum_approach_arm_q_t = torch.lerp( + self._reach_midgrasp_arm_q_bank_t[0].expand(approach_fractions.shape[0], -1), + self._reach_grasp_arm_q_bank_t[0].expand(approach_fractions.shape[0], -1), + approach_fractions.unsqueeze(-1), + ) + self._build_independent_reset_fallbacks() + self._last_source_bank_index = torch.full((self.num_envs,), -1, device=dev, dtype=torch.long) + self._last_arm_bank_index = torch.full((self.num_envs,), -1, device=dev, dtype=torch.long) + self._last_target_bank_index = torch.full((self.num_envs,), -1, device=dev, dtype=torch.long) + start_stage = int(self.cfg.curriculum_start_stage) + start_randomization_level = int(self.cfg.curriculum_randomization_start_level) + self.curriculum_stage = torch.full((self.num_envs,), start_stage, device=dev, dtype=torch.long) + self.curriculum_randomization_level = torch.full( + (self.num_envs,), + start_randomization_level, + device=dev, + dtype=torch.long, + ) + self.reset_dataset_row_id = torch.full((self.num_envs,), -1, device=dev, dtype=torch.long) + # Diagnostic tools may request exact raw dataset rows. This override deliberately has + # one meaning in both adaptive training and frozen playback; it is never relative to a + # curriculum pool. + self._forced_reset_dataset_row = torch.full_like(self.reset_dataset_row_id, -1) + self.pour_target_frac = torch.full( + (self.num_envs,), + float(self.cfg.curriculum_target_frac[start_stage]), + device=dev, + ) + self.episode_succeeded = torch.zeros(self.num_envs, device=dev, dtype=torch.bool) + self.ep_max_target_frac = torch.zeros(self.num_envs, device=dev) + self._success_dwell_count = torch.zeros(self.num_envs, device=dev, dtype=torch.long) + self._lost_grasp_dwell_count = torch.zeros(self.num_envs, device=dev, dtype=torch.long) + self._lifted_grasp_seen = torch.zeros(self.num_envs, device=dev, dtype=torch.bool) + self._target_entry_seen = torch.zeros( + (self.num_envs, self.num_particles), + device=dev, + dtype=torch.bool, + ) + self._held_delivered = torch.zeros_like(self._target_entry_seen) + self._held_delivery_tracker_step = -1 + self._source_inner_lo_t = torch.as_tensor(self._source_inner_lo, device=dev) + self._source_inner_hi_t = torch.as_tensor(self._source_inner_hi, device=dev) + self._target_inner_lo_t = torch.as_tensor(self._target_inner_lo, device=dev) + self._target_inner_hi_t = torch.as_tensor(self._target_inner_hi, device=dev) + self._particle_region_cache: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None + self._particle_region_cache_step = -1 + + def _collision_free_ik_candidates( + self, + prototype_builder: newton.ModelBuilder, + candidate_waypoints: tuple[torch.Tensor, ...], + ) -> torch.Tensor: + """Validate top IK branches against Newton self-contact along the executed path. + + Each candidate occupies one isolated Newton world, so the ordinary collision pipeline can + screen the exact imported Franka collision shapes in a compact batch. Sampling every + joint-space segment catches folded branches that have valid endpoint IK costs but cannot be + tracked by the physical articulation because non-adjacent links intersect. + """ + if len(candidate_waypoints) < 2: + raise ValueError("Collision validation requires at least two IK waypoints.") + shape = candidate_waypoints[0].shape + if len(shape) != 3: + raise ValueError(f"IK candidate waypoints must have shape (rows, candidates, q), got {shape}.") + if any(waypoint.shape != shape for waypoint in candidate_waypoints[1:]): + raise ValueError("Every IK collision-validation waypoint must have the same shape.") + row_count, candidate_count, coordinate_count = shape + if coordinate_count != prototype_builder.joint_coord_count: + raise ValueError( + "IK candidate coordinate count does not match the clone prototype: " + f"{coordinate_count} != {prototype_builder.joint_coord_count}." + ) + + validation_builder = newton.ModelBuilder(up_axis=prototype_builder.up_axis) + validation_builder.replicate(prototype_builder, world_count=row_count) + validation_model = validation_builder.finalize(device=self.device) + if validation_model.world_count != row_count: + raise RuntimeError( + f"Expected {row_count} isolated IK collision worlds, got {validation_model.world_count}." + ) + pipeline = newton.CollisionPipeline( + validation_model, + broad_phase="explicit", + soft_contact_max=0, + verify_buffers=True, + ) + contacts = pipeline.contacts() + state = validation_model.state() + colliding_worlds = wp.zeros(row_count, dtype=wp.int32, device=validation_model.device) + validation_coordinate_count = validation_model.joint_coord_count // row_count + if validation_coordinate_count != coordinate_count: + raise RuntimeError( + "IK validation coordinate count must match the prototype; " + f"got {validation_coordinate_count} coordinates after replicating {coordinate_count}." + ) + robot_q = wp.to_torch(validation_model.joint_q).reshape(row_count, coordinate_count) + collision_free = torch.ones((row_count, candidate_count), device=self.device, dtype=torch.bool) + + # The trajectory action uses a monotonic smoothstep along the same joint-space line. Nine + # equally spaced samples per segment include both endpoints and detect the observed folded + # elbow/hand intersections without turning startup validation into a simulation rollout. + segment_fractions = tuple( + sample / _RESET_PATH_SEGMENT_SUBDIVISIONS for sample in range(_RESET_PATH_SEGMENT_SUBDIVISIONS) + ) + for candidate_index in range(candidate_count): + colliding_worlds.zero_() + for lower, upper in zip(candidate_waypoints[:-1], candidate_waypoints[1:], strict=True): + lower_q = lower[:, candidate_index] + upper_q = upper[:, candidate_index] + for fraction in segment_fractions: + robot_q.copy_(torch.lerp(lower_q, upper_q, fraction).contiguous()) + newton.eval_fk( + validation_model, + validation_model.joint_q, + validation_model.joint_qd, + state, + ) + pipeline.collide(state, contacts) + wp.launch( + _mark_penetrating_self_contacts, + dim=contacts.rigid_contact_max, + inputs=[ + contacts.rigid_contact_count, + contacts.rigid_contact_max, + contacts.rigid_contact_shape0, + contacts.rigid_contact_shape1, + contacts.rigid_contact_point0, + contacts.rigid_contact_point1, + contacts.rigid_contact_normal, + contacts.rigid_contact_margin0, + contacts.rigid_contact_margin1, + validation_model.shape_body, + validation_model.body_world, + state.body_q, + 1.0e-4, + ], + outputs=[colliding_worlds], + device=validation_model.device, + ) + robot_q.copy_(candidate_waypoints[-1][:, candidate_index].contiguous()) + newton.eval_fk( + validation_model, + validation_model.joint_q, + validation_model.joint_qd, + state, + ) + pipeline.collide(state, contacts) + wp.launch( + _mark_penetrating_self_contacts, + dim=contacts.rigid_contact_max, + inputs=[ + contacts.rigid_contact_count, + contacts.rigid_contact_max, + contacts.rigid_contact_shape0, + contacts.rigid_contact_shape1, + contacts.rigid_contact_point0, + contacts.rigid_contact_point1, + contacts.rigid_contact_normal, + contacts.rigid_contact_margin0, + contacts.rigid_contact_margin1, + validation_model.shape_body, + validation_model.body_world, + state.body_q, + 1.0e-4, + ], + outputs=[colliding_worlds], + device=validation_model.device, + ) + collision_free[:, candidate_index] = wp.to_torch(colliding_worlds) == 0 + + return collision_free + + def _load_reset_dataset(self, configured_path: str) -> None: + """Load and stage a validated direct-state dataset on the simulation device.""" + cache_path = Path(configured_path).expanduser().resolve() + if not cache_path.is_file(): + raise FileNotFoundError(f"Franka Pour reset dataset not found: {cache_path}") + payload = torch.load(cache_path, map_location="cpu", weights_only=True) + self._validate_loaded_reset_dataset(payload) + expected_hash = self.cfg.reset_dataset_content_sha256 + if expected_hash is not None and payload["content_sha256"] != expected_hash: + raise RuntimeError( + "Franka Pour reset dataset content hash does not match the configured hash: " + f"{payload['content_sha256']} != {expected_hash}." + ) + + metadata = payload["metadata"] + expected_joint_names = tuple(ARM_JOINTS + FINGER_JOINTS) + if tuple(metadata["joint_names"]) != expected_joint_names: + raise RuntimeError("Reset-dataset joint order does not match the Franka runtime joint order.") + if metadata["frame"] != "environment" or metadata["quaternion_order"] != "xyzw": + raise RuntimeError("Reset-dataset poses must use environment-frame XYZW representation.") + layouts = payload["particle_layouts"] + local_position = layouts["local_position"].to(device=self.device, dtype=torch.float32) + local_velocity = layouts["local_velocity"].to(device=self.device, dtype=torch.float32) + if local_position.shape[1:] != self._media_local_points_t.shape: + raise RuntimeError( + "Reset-dataset particle layout does not match the runtime media shape: " + f"{tuple(local_position.shape[1:])} != {tuple(self._media_local_points_t.shape)}." + ) + if not bool(torch.allclose(local_position[0], self._media_local_points_t, atol=1.0e-7, rtol=0.0)): + raise RuntimeError("Reset-dataset particle layout does not match the current cup fill lattice.") + + self._reset_dataset_metadata = metadata + self._reset_dataset_states = { + name: value.to(device=self.device, non_blocking=True) for name, value in payload["states"].items() + } + self._reset_dataset_particle_local_position = local_position + self._reset_dataset_particle_local_velocity = local_velocity + logger.info( + "Loaded reset dataset %s (%s) with %d rows.", + cache_path, + payload["content_sha256"], + metadata["state_count"], + ) + + def _validate_loaded_reset_dataset(self, payload: dict) -> None: + """Require dynamic-validation provenance for normal task execution.""" + validate_production_reset_dataset( + payload, + expected_task_contract=build_franka_pour_reset_task_contract(self), + ) + + def _build_randomized_reset_bank(self) -> None: + """Build a small Newton-IK bank for collision-safe randomized pre-grasp resets.""" + plan = sim_utils.SimulationContext.instance().get_clone_plan() + resolved = resolve_clone_plan_source(self._robot.cfg.prim_path, plan) if plan is not None else None + if resolved is None: + raise RuntimeError(f"Could not resolve clone-plan source for {self._robot.cfg.prim_path!r}.") + source_path = resolved[0] + prototype_origin = -self.env_origins[0] + prototype_xform = wp.transform( + wp.vec3(*prototype_origin.tolist()), + wp.quat_identity(), + ) + + def local_prototype_builder() -> newton.ModelBuilder: + """Copy the clone source into an environment-local IK frame.""" + source_builder = copy_newton_source_builder(source_path) + local_builder = newton.ModelBuilder(up_axis=source_builder.up_axis) + local_builder.add_builder(source_builder, xform=prototype_xform) + return local_builder + + # Clone layouts are centered around the world origin, so the source environment can be + # tens of metres away for large training batches. Solve IK and validate collisions in the + # environment-local frame to make the reset bank independent of ``num_envs`` and avoid + # losing millimetre-scale joint clearance to float32 cancellation. + prototype_builder = local_prototype_builder() + model = local_prototype_builder().finalize(device=self.device) + + hand_matches = [ + body_id + for body_id, label in enumerate(model.body_label) + if str(label).rsplit("/", 1)[-1] == self.cfg.tcp_body_name + ] + if len(hand_matches) != 1: + raise RuntimeError( + f"Expected one {self.cfg.tcp_body_name!r} body in the IK prototype, found {hand_matches}." + ) + hand_id = hand_matches[0] + joint_labels = [str(label).rsplit("/", 1)[-1] for label in model.joint_label] + joint_q_start = wp.to_torch(model.joint_q_start).to(device=self.device, dtype=torch.long) + + def coordinate_id(joint_name: str) -> int: + matches = [joint_id for joint_id, label in enumerate(joint_labels) if label == joint_name] + if len(matches) != 1: + raise RuntimeError(f"Expected one {joint_name!r} joint in the IK prototype, found {matches}.") + return int(joint_q_start[matches[0]].item()) + + arm_coordinate_ids = torch.tensor( + [coordinate_id(joint_name) for joint_name in ARM_JOINTS], + device=self.device, + dtype=torch.long, + ) + finger_coordinate_ids = torch.tensor( + [coordinate_id(joint_name) for joint_name in FINGER_JOINTS], + device=self.device, + dtype=torch.long, + ) + + def tcp_pose_for_arm_q(arm_q_values: tuple[float, ...]) -> torch.Tensor: + """Evaluate one prototype arm configuration and return its TCP world pose.""" + joint_q = wp.to_torch(model.joint_q).to(device=self.device, dtype=torch.float32).clone() + joint_q[arm_coordinate_ids] = torch.as_tensor(arm_q_values, device=self.device) + joint_q[finger_coordinate_ids] = float(self.cfg.gripper_preload_pos) + state = model.state() + newton.eval_fk( + model, + wp.from_torch(joint_q.contiguous(), dtype=wp.float32), + model.joint_qd, + state, + ) + hand_pose = wp.to_torch(state.body_q)[hand_id : hand_id + 1] + tcp_pos, tcp_quat = math_utils.combine_frame_transforms( + hand_pose[:, :3], + hand_pose[:, 3:7], + self._tcp_offset_pos[0:1], + self._tcp_offset_quat[0:1], + ) + return torch.cat((tcp_pos, tcp_quat), dim=-1)[0].clone() + + nominal_carry_tcp_pose = tcp_pose_for_arm_q(self.cfg.curriculum_carry_arm_q) + nominal_pour_tcp_pose = tcp_pose_for_arm_q(self.cfg.curriculum_pour_arm_q) + nominal_tilt_tcp_pose = tcp_pose_for_arm_q(self.cfg.curriculum_pour_target_arm_q) + + grid_size = int(self.cfg.curriculum_randomized_reset_ik_grid_size) + nominal_source = torch.as_tensor(self.cfg.cup_reset_pos, device=self.device) + if self.cfg.curriculum_randomized_source_radius_range is None: + source_range = torch.as_tensor( + self.cfg.curriculum_randomized_source_position_range, + device=self.device, + dtype=torch.float32, + ) + x_offsets = torch.linspace(-source_range[0], source_range[0], grid_size, device=self.device) + y_offsets = torch.linspace(-source_range[1], source_range[1], grid_size, device=self.device) + offset_x, offset_y = torch.meshgrid(x_offsets, y_offsets, indexing="ij") + correlation = float(self.cfg.curriculum_randomized_source_xy_correlation) + if float(source_range[0]) > 0.0: + reachable_diagonal_y = offset_x / source_range[0] * source_range[1] + else: + reachable_diagonal_y = torch.zeros_like(offset_y) + offset_y = correlation * reachable_diagonal_y + (1.0 - correlation) * offset_y + offsets = torch.stack((offset_x.flatten(), offset_y.flatten()), dim=-1) + base_source_positions = nominal_source.repeat(offsets.shape[0], 1) + base_source_positions[:, :2] += offsets + radial_facing = False + else: + base_source_positions = polar_workspace_cells( + nominal_source, + radius_range=self.cfg.curriculum_randomized_source_radius_range, + azimuth_half_range=float(self.cfg.curriculum_randomized_source_azimuth_range), + grid_size=grid_size, + ) + radial_facing = True + + samples_per_source = int(self.cfg.curriculum_randomized_reset_ik_samples_per_source) + pair_count = samples_per_source // 2 + if self.cfg.curriculum_randomized_reset_tcp_offset_lower is None: + pair_index = torch.arange(pair_count, device=self.device, dtype=torch.float32) + 0.5 + pair_directions = ( + 2.0 + * torch.stack( + ( + torch.frac(pair_index * 0.754877666), + torch.frac(pair_index * 0.569840296), + torch.frac(pair_index * 0.438447187), + ), + dim=-1, + ) + - 1.0 + ) + paired_jitter = pair_directions * torch.as_tensor( + self.cfg.curriculum_randomized_reset_tcp_jitter, + device=self.device, + ) + jitter_parts = [paired_jitter, -paired_jitter] + if samples_per_source % 2: + jitter_parts.insert(0, torch.zeros((1, 3), device=self.device)) + tcp_jitter_samples = torch.cat(jitter_parts, dim=0) + else: + tcp_jitter_samples = asymmetric_reset_offset_samples( + self.cfg.curriculum_randomized_reset_tcp_offset_lower, + self.cfg.curriculum_randomized_reset_tcp_offset_upper, + samples_per_source, + device=self.device, + dtype=torch.float32, + ) + zero_jitter_slot = int(torch.argmin(torch.linalg.vector_norm(tcp_jitter_samples, dim=-1)).item()) + + # Pair upright source yaw with the symmetric TCP samples. Keeping both signs in each source + # cell avoids a directional reset bias while retaining the existing compact bank size. + yaw_range = float(self.cfg.curriculum_randomized_source_yaw_range) + pair_yaws = ( + (torch.arange(pair_count, device=self.device, dtype=torch.float32) + 1.0) / max(pair_count, 1) * yaw_range + ) + yaw_parts = [pair_yaws, -pair_yaws] + if samples_per_source % 2: + yaw_parts.insert(0, torch.zeros(1, device=self.device)) + source_yaw_samples = torch.cat(yaw_parts, dim=0) + + source_cell_count = base_source_positions.shape[0] + nominal_source_cell = int( + torch.argmin(torch.linalg.vector_norm(base_source_positions - nominal_source, dim=-1)).item() + ) + base_source_positions = base_source_positions.repeat_interleave(samples_per_source, dim=0) + base_tcp_jitter = tcp_jitter_samples.repeat(source_cell_count, 1) + reset_rotation_vectors = reset_rotation_vector_samples( + self.cfg.curriculum_randomized_reset_tcp_rotation_angle_range, + samples_per_source, + device=self.device, + dtype=torch.float32, + ) + rotation_sample_ids = balanced_cyclic_permutations( + torch.arange(samples_per_source, device=self.device), + source_cell_count, + ) + base_tcp_rotation_vectors = reset_rotation_vectors[rotation_sample_ids].reshape(-1, 3) + # Rotate the yaw-to-jitter pairing in each source cell. Every cell retains identical yaw + # and jitter/rotation marginals, while each fixed slot sees every perturbation globally. + base_source_yaw_jitter = balanced_cyclic_permutations(source_yaw_samples, source_cell_count).reshape(-1) + rows_per_extent = base_source_positions.shape[0] + extent_levels = tuple(float(extent) for extent in self.cfg.curriculum_randomization_extent_levels) + level_count = len(extent_levels) + + # Scale the same balanced source-cell/reset-offset design at every extent instead of + # filtering one full-range bank. Equal-size level censuses preserve all marginals while + # expanding the physical domain smoothly; the final extent is bit-for-bit full amplitude. + base_source_offsets = base_source_positions - nominal_source + source_positions = nominal_source + scale_randomization_rows_by_extent( + base_source_offsets, extent_levels + ).reshape(-1, 3) + tcp_jitter = scale_randomization_rows_by_extent(base_tcp_jitter, extent_levels).reshape(-1, 3) + tcp_rotation_vectors = scale_randomization_rows_by_extent( + base_tcp_rotation_vectors, + extent_levels, + ).reshape(-1, 3) + source_yaw_jitter = scale_randomization_rows_by_extent( + base_source_yaw_jitter, + extent_levels, + ).reshape(-1) + if radial_facing: + source_yaws = torch.atan2(source_positions[:, 1], source_positions[:, 0]) + source_yaw_jitter + else: + source_yaws = source_yaw_jitter + randomized_bank_size = rows_per_extent * level_count + zero_bank_row = nominal_source_cell * samples_per_source + zero_jitter_slot + + # Stages two and three reuse safe varied arm starts around the nominal upright source. + # Append dedicated solve-only rows so the randomized nominal-source rows keep all yaw values; + # these rows are sliced out before constructing the stage-four extent banks. + source_positions = torch.cat((source_positions, nominal_source.repeat(samples_per_source, 1)), dim=0) + tcp_jitter = torch.cat((tcp_jitter, tcp_jitter_samples), dim=0) + tcp_rotation_vectors = torch.cat( + (tcp_rotation_vectors, torch.zeros((samples_per_source, 3), device=self.device)), + dim=0, + ) + source_yaws = torch.cat((source_yaws, torch.zeros(samples_per_source, device=self.device)), dim=0) + bank_size = source_positions.shape[0] + source_quaternions = torch.zeros((bank_size, 4), device=self.device) + source_quaternions[:, 2] = torch.sin(0.5 * source_yaws) + source_quaternions[:, 3] = torch.cos(0.5 * source_yaws) + + grasp_offset_c = torch.zeros((bank_size, 3), device=self.device) + grasp_offset_c[:, 2] = float(self.cfg.cup_grasp_height) + standoff_c = torch.as_tensor( + self.cfg.curriculum_randomized_reset_tcp_standoff, + device=self.device, + ).expand(bank_size, -1) + # Standoff and jitter are cup-local. Rotating both with source yaw keeps every reset on the + # same side-approach ray instead of sweeping the horizontal fingers across a cup corner. + tcp_offset_c = grasp_offset_c + standoff_c + tcp_jitter + tcp_positions = source_positions + math_utils.quat_apply(source_quaternions, tcp_offset_c) + + target_positions_w = tcp_positions + aligned_target_rotations_w = math_utils.quat_mul( + source_quaternions, + self._desired_grasp_tcp_quat_c[:1].expand(bank_size, -1), + ) + rotation_angles = torch.linalg.vector_norm(tcp_rotation_vectors, dim=-1) + rotation_axes = tcp_rotation_vectors / rotation_angles.clamp_min(1.0e-9).unsqueeze(-1) + reset_rotation_quaternions = math_utils.quat_from_angle_axis(rotation_angles, rotation_axes) + target_rotations_w = math_utils.quat_mul( + aligned_target_rotations_w, + reset_rotation_quaternions, + ).contiguous() + + target_name = "reset_tcp" + + def ik_objectives() -> list: + return [ + NewtonIKPoseObjectiveCfg( + body_name=self.cfg.tcp_body_name, + name=target_name, + body_offset_pos=self.cfg.tcp_offset_pos, + body_offset_rot=self.cfg.tcp_offset_rot, + # Grasp capture tolerates only millimetres of cross-track error. Weight + # position strongly enough that accepted IK cost cannot hide a bad insertion. + position_weight=100.0, + rotation_weight=5.0, + ), + NewtonIKJointLimitObjectiveCfg(weight=1.0), + ] + + # The horizontal arm admits several kinematic branches. Sample only the randomized reset + # pose so each row starts on an interior branch, then use a separate single-seed solver to + # track every subsequent waypoint continuously without branch jumps. + reset_solver = NewtonIKSolver( + NewtonIKSolverCfg( + optimizer="lm", + jacobian_mode="analytic", + sampler="gauss", + n_seeds=64, + noise_std=0.75, + iterations=int(self.cfg.curriculum_randomized_reset_ik_iterations), + lambda_initial=0.1, + ), + model=model, + num_envs=bank_size, + device=str(model.device), + objectives=ik_objectives(), + link_resolver=lambda body_name: hand_id, + ) + pose_objective = reset_solver.objectives_by_name[target_name] + pose_objective.position_objective.set_target_positions( + wp.from_torch(target_positions_w.contiguous(), dtype=wp.vec3) + ) + pose_objective.rotation_objective.set_target_rotations( + wp.from_torch(target_rotations_w.contiguous(), dtype=wp.vec4) + ) + + seed = wp.to_torch(model.joint_q).to(device=self.device, dtype=torch.float32).repeat(bank_size, 1) + seed[:, arm_coordinate_ids] = torch.as_tensor(self.cfg.arm_home, device=self.device) + seed[:, finger_coordinate_ids] = float(self.cfg.gripper_open_pos) + reset_solver.solve(wp.from_torch(seed.contiguous(), dtype=wp.float32)) + expanded_q = wp.to_torch(reset_solver.joint_q).reshape( + bank_size, + reset_solver.cfg.n_seeds, + -1, + ) + expanded_costs = wp.to_torch(reset_solver.costs).reshape(bank_size, reset_solver.cfg.n_seeds) + expanded_arm_q = expanded_q[:, :, arm_coordinate_ids] + arm_limits = self._joint_pos_limits_t[0, self._arm_joint_ids] + expanded_margin = torch.minimum( + expanded_arm_q - arm_limits[:, 0], + arm_limits[:, 1] - expanded_arm_q, + ).amin(dim=-1) + candidate_valid = ( + torch.isfinite(expanded_arm_q).all(dim=-1) + & torch.isfinite(expanded_costs) + & (expanded_costs <= float(self.cfg.curriculum_randomized_reset_ik_max_cost)) + & (expanded_margin >= float(self.cfg.curriculum_randomized_reset_ik_joint_margin)) + ) + candidate_rows = torch.arange(bank_size, device=self.device) + seed_count = int(reset_solver.cfg.n_seeds) + reset_has_candidate = candidate_valid.any(dim=-1) + + # A high-clearance reset solution can still run into a joint limit during the straight + # lateral insertion. Preserve all feasible reset branches through pre-grasp and grasp, + # then choose the trajectory with the largest minimum clearance over all three poses. + # Invalid candidates are replaced with a finite row-local seed before launching IK so one + # bad sample cannot poison the batched solver; their validity mask remains false. + trajectory_valid = candidate_valid.clone() + trajectory_margin = torch.where( + trajectory_valid, + expanded_margin, + torch.full_like(expanded_margin, -torch.inf), + ) + valid_fallback_indices = torch.argmax(trajectory_valid.to(dtype=torch.int32), dim=-1) + margin_fallback_indices = torch.argmax( + torch.nan_to_num(expanded_margin, nan=-torch.inf), + dim=-1, + ) + fallback_indices = torch.where(reset_has_candidate, valid_fallback_indices, margin_fallback_indices) + fallback_q = expanded_q[candidate_rows, fallback_indices].unsqueeze(1) + trajectory_q = torch.where(trajectory_valid.unsqueeze(-1), expanded_q, fallback_q).clone() + + continuation_solver = NewtonIKSolver( + NewtonIKSolverCfg( + optimizer="lm", + jacobian_mode="analytic", + sampler="none", + n_seeds=1, + iterations=int(self.cfg.curriculum_randomized_reset_ik_iterations), + lambda_initial=0.1, + ), + model=model, + num_envs=bank_size * seed_count, + device=str(model.device), + objectives=ik_objectives(), + link_resolver=lambda body_name: hand_id, + ) + continuation_objective = continuation_solver.objectives_by_name[target_name] + + def solve_candidate_waypoint( + target_pose: torch.Tensor, + initial_guess: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + nonlocal trajectory_margin, trajectory_valid + expanded_target = target_pose.repeat_interleave(seed_count, dim=0) + continuation_objective.position_objective.set_target_positions( + wp.from_torch(expanded_target[:, :3].contiguous(), dtype=wp.vec3) + ) + continuation_objective.rotation_objective.set_target_rotations( + wp.from_torch(expanded_target[:, 3:7].contiguous(), dtype=wp.vec4) + ) + solved_full = ( + wp.to_torch( + continuation_solver.solve( + wp.from_torch(initial_guess.reshape(bank_size * seed_count, -1).contiguous(), dtype=wp.float32) + ) + ) + .reshape(bank_size, seed_count, -1) + .clone() + ) + solved_arm = solved_full[:, :, arm_coordinate_ids] + solved_costs = wp.to_torch(continuation_solver.costs).reshape(bank_size, seed_count).clone() + solved_margin = torch.minimum( + solved_arm - arm_limits[:, 0], + arm_limits[:, 1] - solved_arm, + ).amin(dim=-1) + waypoint_valid = ( + torch.isfinite(solved_arm).all(dim=-1) + & torch.isfinite(solved_costs) + & (solved_costs <= float(self.cfg.curriculum_randomized_reset_ik_max_cost)) + & (solved_margin >= float(self.cfg.curriculum_randomized_reset_ik_joint_margin)) + ) + prior_valid = trajectory_valid.clone() + trajectory_valid &= waypoint_valid + trajectory_margin = torch.where( + trajectory_valid, + torch.minimum(trajectory_margin, solved_margin), + torch.full_like(trajectory_margin, -torch.inf), + ) + has_candidate = trajectory_valid.any(dim=-1) + valid_fallback_indices = torch.argmax(trajectory_valid.to(dtype=torch.int32), dim=-1) + prior_fallback_indices = torch.argmax(prior_valid.to(dtype=torch.int32), dim=-1) + valid_fallback = solved_full[candidate_rows, valid_fallback_indices] + prior_fallback = initial_guess[candidate_rows, prior_fallback_indices] + fallback = torch.where(has_candidate.unsqueeze(-1), valid_fallback, prior_fallback).unsqueeze(1) + sanitized = torch.where(trajectory_valid.unsqueeze(-1), solved_full, fallback) + return sanitized, solved_arm, solved_costs, solved_margin + + # Route every reset through a centered lateral pre-grasp. The open fingers then insert + # along hand +Z, parallel to the table, without sweeping across the glass. + pregrasp_tcp_positions = source_positions + math_utils.quat_apply( + source_quaternions, + grasp_offset_c + standoff_c, + ) + pregrasp_target_positions_w = pregrasp_tcp_positions + pregrasp_target_pose = torch.cat( + (pregrasp_target_positions_w, aligned_target_rotations_w), + dim=-1, + ) + pregrasp_candidates, pregrasp_arm_candidates, pregrasp_cost_candidates, pregrasp_margin_candidates = ( + solve_candidate_waypoint(pregrasp_target_pose, trajectory_q) + ) + + # Split the horizontal insertion into two straight Cartesian segments. Interpolating the + # joint coordinates directly from a 12 cm standoff bowed the physical TCP below the glass + # on some randomized IK branches; the midpoint keeps the open fingers on the side-approach + # ray before the grasp gate is allowed to close. + midgrasp_tcp_positions = source_positions + math_utils.quat_apply( + source_quaternions, + grasp_offset_c + 0.5 * standoff_c, + ) + midgrasp_target_pose = torch.cat( + (midgrasp_tcp_positions, aligned_target_rotations_w), + dim=-1, + ) + midgrasp_candidates, midgrasp_arm_candidates, midgrasp_cost_candidates, midgrasp_margin_candidates = ( + solve_candidate_waypoint(midgrasp_target_pose, pregrasp_candidates) + ) + + # Solve the paired grasp waypoint from the centered midpoint. Keeping every pose in the + # same prevalidated bank avoids reset-time IK and preserves stationary action semantics. + grasp_offset_c[:, 2] -= float(self.cfg.curriculum_grasp_descent_overshoot) + grasp_tcp_positions = source_positions + math_utils.quat_apply(source_quaternions, grasp_offset_c) + grasp_target_positions_w = grasp_tcp_positions + grasp_target_pose = torch.cat( + (grasp_target_positions_w, aligned_target_rotations_w), + dim=-1, + ) + grasp_candidates, grasp_arm_candidates, grasp_cost_candidates, grasp_margin_candidates = ( + solve_candidate_waypoint(grasp_target_pose, midgrasp_candidates) + ) + + # Keep all branches alive through the upright lift as well. This source-relative waypoint + # is where a few edge-of-workspace branches that are safe at grasp first reach a limit. + source_delta = source_positions - nominal_source + carry_target_pose = nominal_carry_tcp_pose.repeat(bank_size, 1) + carry_position_range = torch.as_tensor( + (*self.cfg.curriculum_randomized_carry_position_range, 0.0), + device=self.device, + ) + carry_target_pose[:, :3] += torch.clamp(source_delta, min=-carry_position_range, max=carry_position_range) + # Canonicalize yaw during the upright lift/carry. Keeping the radial grasp orientation all + # the way to a fixed receiver-side pour forced broad-angle starts through folded wrist + # branches. The cup remains upright while this collision-screened waypoint smoothly + # unwinds source yaw before transport. + carry_target_pose[:, 3:7] = self._desired_grasp_tcp_quat_c[:1] + carry_candidates, carry_arm_candidates, carry_cost_candidates, carry_margin_candidates = ( + solve_candidate_waypoint(carry_target_pose, grasp_candidates) + ) + + # Keep receiver geometry fixed across the arm-start variants of one source cell. Two + # low-discrepancy coordinates cover the receiver region while the geometric projection + # below keeps it safely behind the source. Centering both sequences on the nominal source + # cell makes the zero-amplitude frontier exactly reproduce the authored task. + cell_index = torch.arange(source_cell_count, device=self.device, dtype=torch.float32) + center_cell = float(nominal_source_cell) + base_unit_samples = torch.stack( + ( + torch.remainder((cell_index - center_cell) * 0.754877666 + 0.5, 1.0), + torch.remainder((cell_index - center_cell) * 0.569840296 + 0.5, 1.0), + ), + dim=-1, + ).repeat_interleave(samples_per_source, dim=0) + source_outer_half_x = self.cfg.source_cup_inner_width / 2.0 + self.cfg.source_cup_wall_thickness + source_outer_half_y = self.cfg.source_cup_inner_depth / 2.0 + self.cfg.source_cup_wall_thickness + target_outer_half_y = self.cfg.target_cup_inner_depth / 2.0 + self.cfg.target_cup_wall_thickness + minimum_y_separation = ( + source_outer_half_x * torch.abs(torch.sin(source_yaws[:randomized_bank_size])) + + source_outer_half_y * torch.abs(torch.cos(source_yaws[:randomized_bank_size])) + + target_outer_half_y + + float(self.cfg.curriculum_randomized_cup_clearance) + ) + target_range = torch.as_tensor( + self.cfg.curriculum_randomized_target_position_range, + device=self.device, + ) + target_xy_parts = [] + for level, extent in enumerate(extent_levels): + rows = slice(level * rows_per_extent, (level + 1) * rows_per_extent) + target_xy_parts.append( + target_xy_behind_source( + source_positions[rows, :2], + target_center=self.cfg.curriculum_randomized_target_center_xy, + target_half_range=target_range * extent, + minimum_y_separation=minimum_y_separation[rows], + unit_samples=base_unit_samples, + ) + ) + target_xy = torch.cat(target_xy_parts, dim=0) + target_positions = torch.as_tensor(self.cfg.target_cup_reset_pos, device=self.device).repeat( + source_positions.shape[0], 1 + ) + target_positions[:randomized_bank_size, :2] = target_xy + + # Broad receiver positions need more than the single authored IK seed used by the narrow + # task. Solve a compact set of branches for both upright-pour and deep-tilt endpoints, then + # choose the shortest jointly valid pair. Source-side branch selection below still accounts + # for the final carry-to-pour transition. + receiver_seed_count = 16 + solver = NewtonIKSolver( + NewtonIKSolverCfg( + optimizer="lm", + jacobian_mode="analytic", + sampler="gauss", + n_seeds=receiver_seed_count, + noise_std=0.5, + iterations=int(self.cfg.curriculum_randomized_reset_ik_iterations), + lambda_initial=0.1, + ), + model=model, + num_envs=bank_size, + device=str(model.device), + objectives=ik_objectives(), + link_resolver=lambda body_name: hand_id, + ) + pose_objective = solver.objectives_by_name[target_name] + + def solve_reference_candidates( + target_pose: torch.Tensor, + initial_guess: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + pose_objective.position_objective.set_target_positions( + wp.from_torch(target_pose[:, :3].contiguous(), dtype=wp.vec3) + ) + pose_objective.rotation_objective.set_target_rotations( + wp.from_torch(target_pose[:, 3:7].contiguous(), dtype=wp.vec4) + ) + solver.solve(wp.from_torch(initial_guess.contiguous(), dtype=wp.float32)) + raw_solved_full = ( + wp.to_torch(solver.joint_q).reshape(bank_size, int(solver.cfg.n_seeds), model.joint_coord_count).clone() + ) + solved_arm = raw_solved_full[:, :, arm_coordinate_ids] + solved_costs = wp.to_torch(solver.costs).reshape(bank_size, receiver_seed_count).clone() + solved_margin = torch.minimum( + solved_arm - arm_limits[:, 0], + arm_limits[:, 1] - solved_arm, + ).amin(dim=-1) + solved_valid = ( + torch.isfinite(solved_arm).all(dim=-1) + & torch.isfinite(solved_costs) + & (solved_costs <= float(self.cfg.curriculum_randomized_reset_ik_max_cost)) + & (solved_margin >= float(self.cfg.curriculum_randomized_reset_ik_joint_margin)) + ) + return raw_solved_full, solved_arm, solved_costs, solved_margin, solved_valid + + nominal_target = torch.as_tensor(self.cfg.target_cup_reset_pos, device=self.device) + target_delta = target_positions - nominal_target + + pour_target_pose = nominal_pour_tcp_pose.repeat(bank_size, 1) + pour_target_pose[:, :3] += target_delta + clearance_by_row = torch.zeros(bank_size, device=self.device) + clearance_by_row[:randomized_bank_size] = torch.as_tensor( + extent_levels, + device=self.device, + ).repeat_interleave(rows_per_extent) * float(self.cfg.curriculum_randomized_pour_clearance) + pour_target_pose[:, 2] += clearance_by_row + pour_target_pose[:, 3:7] = self._desired_grasp_tcp_quat_c[:1] + pour_seed = seed.clone() + pour_seed[:, arm_coordinate_ids] = torch.as_tensor(self.cfg.curriculum_pour_arm_q, device=self.device) + pour_full_candidates, pour_arm_candidates, pour_cost_candidates, pour_margin_candidates, pour_valid = ( + solve_reference_candidates( + pour_target_pose, + pour_seed, + ) + ) + + tilt_target_pose = nominal_tilt_tcp_pose.repeat(bank_size, 1) + tilt_target_pose[:, :3] += target_delta + tilt_target_pose[:, 2] += clearance_by_row + tilt_seed = seed.clone() + tilt_seed[:, arm_coordinate_ids] = torch.as_tensor( + self.cfg.curriculum_pour_target_arm_q, + device=self.device, + ) + tilt_full_candidates, tilt_arm_candidates, tilt_cost_candidates, tilt_margin_candidates, tilt_valid = ( + solve_reference_candidates( + tilt_target_pose, + tilt_seed, + ) + ) + + receiver_pair_valid = pour_valid.unsqueeze(-1) & tilt_valid.unsqueeze(-2) + receiver_pair_transition = ( + (tilt_arm_candidates.unsqueeze(-3) - pour_arm_candidates.unsqueeze(-2)).square().sum(dim=-1) + ) + pour_nominal = torch.as_tensor(self.cfg.curriculum_pour_arm_q, device=self.device) + tilt_nominal = torch.as_tensor(self.cfg.curriculum_pour_target_arm_q, device=self.device) + # Prefer the authored elbow/wrist branch whenever several receiver solutions have similar + # path length. Without this tie-break, millimetric target changes can flip the complete + # joint reference by more than a radian even though the task-space pose is continuous. + nominal_branch_weight = 0.5 + receiver_pair_transition += nominal_branch_weight * ( + (pour_arm_candidates - pour_nominal).square().sum(dim=-1).unsqueeze(-1) + + (tilt_arm_candidates - tilt_nominal).square().sum(dim=-1).unsqueeze(-2) + ) + receiver_pair_score = torch.where( + receiver_pair_valid, + receiver_pair_transition, + torch.full_like(receiver_pair_transition, torch.inf), + ) + receiver_valid = receiver_pair_valid.flatten(start_dim=1).any(dim=-1) + source_trajectory_valid = trajectory_valid.any(dim=-1) + source_transition = ( + (pregrasp_arm_candidates - expanded_arm_q).square().sum(dim=-1) + + (midgrasp_arm_candidates - pregrasp_arm_candidates).square().sum(dim=-1) + + (grasp_arm_candidates - midgrasp_arm_candidates).square().sum(dim=-1) + + (carry_arm_candidates - grasp_arm_candidates).square().sum(dim=-1) + ) + source_nominal_deviation = (expanded_arm_q - self._nominal_reference_waypoints_t[0]).square().sum(dim=-1) + ( + carry_arm_candidates - self._nominal_reference_waypoints_t[4] + ).square().sum(dim=-1) + + # Select complete reset->carry->pour->tilt paths jointly. Committing to the lowest-cost + # receiver pair first can discard a different receiver branch that is much closer to the + # source-relative carry branch, especially near a radial workspace boundary. Retain four + # tilt alternatives per pour branch, rank the resulting source/pour/tilt paths, and send + # the best complete candidates through the exact self-collision screen below. + tilt_alternatives_per_pour = min(4, receiver_seed_count) + fallback_tilt_score = receiver_pair_score.clone() + fallback_tilt_score[:, :, 0] = torch.inf + fallback_scores_by_pour, fallback_tilt_indices_by_pour = torch.topk( + fallback_tilt_score, + k=tilt_alternatives_per_pour - 1, + dim=-1, + largest=False, + sorted=True, + ) + receiver_scores_by_pour = torch.cat( + (receiver_pair_score[:, :, :1], fallback_scores_by_pour), + dim=-1, + ) + zero_tilt_indices = torch.zeros_like(fallback_tilt_indices_by_pour[:, :, :1]) + receiver_tilt_indices_by_pour = torch.cat( + (zero_tilt_indices, fallback_tilt_indices_by_pour), + dim=-1, + ) + carry_to_pour_transition = ( + (pour_arm_candidates.unsqueeze(1) - carry_arm_candidates.unsqueeze(2)).square().sum(dim=-1) + ) + complete_path_score = ( + source_transition[:, :, None, None] + + nominal_branch_weight * source_nominal_deviation[:, :, None, None] + + carry_to_pour_transition[:, :, :, None] + + receiver_scores_by_pour[:, None, :, :] + ) + complete_path_valid = trajectory_valid[:, :, None, None] & torch.isfinite( + receiver_scores_by_pour[:, None, :, :] + ) + open_reset_branch = expanded_arm_q[:, :, 5] <= float(self.cfg.curriculum_randomized_reset_joint6_max) + row_has_open_reset_branch = (complete_path_valid & open_reset_branch[:, :, None, None]).any( + dim=(1, 2, 3), + keepdim=True, + ) + complete_path_valid &= ~row_has_open_reset_branch | open_reset_branch[:, :, None, None] + complete_path_score = torch.where( + complete_path_valid, + complete_path_score, + torch.full_like(complete_path_score, torch.inf), + ) + # Broad radial starts expose multiple valid Franka IK branches. Reserve part of the exact + # collision screen for the deterministic, unperturbed source seed; otherwise receiver-path + # combinations can fill the global top-k and silently remove the only branch continuous + # with the nominal grasp. The remaining slots retain the shortest complete paths globally. + collision_candidate_count = min(64, complete_path_score[0].numel()) + receiver_paths_per_source = receiver_seed_count * tilt_alternatives_per_pour + deterministic_source_candidate_count = min(8, receiver_paths_per_source, collision_candidate_count) + flat_complete_path_score = complete_path_score.flatten(start_dim=1) + flat_complete_path_valid = complete_path_valid.flatten(start_dim=1) + deterministic_source_score = complete_path_score[:, 0].flatten(start_dim=1).clone() + deterministic_source_score[:, 0] = torch.inf + deterministic_source_fallback_indices = torch.topk( + deterministic_source_score, + k=deterministic_source_candidate_count - 1, + dim=-1, + largest=False, + sorted=True, + ).indices + deterministic_source_path_indices = torch.cat( + ( + torch.zeros((bank_size, 1), device=self.device, dtype=torch.long), + deterministic_source_fallback_indices, + ), + dim=-1, + ) + global_path_indices = torch.topk( + flat_complete_path_score, + k=collision_candidate_count - deterministic_source_candidate_count, + dim=-1, + largest=False, + sorted=True, + ).indices + collision_path_indices = torch.cat((deterministic_source_path_indices, global_path_indices), dim=-1) + collision_source_indices = torch.div( + collision_path_indices, + receiver_paths_per_source, + rounding_mode="floor", + ) + collision_receiver_indices = collision_path_indices.remainder(receiver_paths_per_source) + collision_pour_indices = torch.div( + collision_receiver_indices, + tilt_alternatives_per_pour, + rounding_mode="floor", + ) + collision_tilt_slots = collision_receiver_indices.remainder(tilt_alternatives_per_pour) + collision_tilt_indices = receiver_tilt_indices_by_pour[ + candidate_rows.unsqueeze(-1), + collision_pour_indices, + collision_tilt_slots, + ] + ranked_endpoint_valid = torch.gather(flat_complete_path_valid, 1, collision_path_indices) + + def gather_candidates(values: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: + return torch.gather( + values, + 1, + indices.unsqueeze(-1).expand(-1, -1, values.shape[-1]), + ).contiguous() + + def gather_source_waypoint(values: torch.Tensor, finger_position: float) -> torch.Tensor: + gathered = gather_candidates(values, collision_source_indices) + gathered = torch.where(ranked_endpoint_valid.unsqueeze(-1), gathered, seed.unsqueeze(1)) + gathered[:, :, finger_coordinate_ids] = finger_position + return gathered + + pour_collision_candidates = gather_candidates(pour_full_candidates, collision_pour_indices) + pour_collision_candidates = torch.where( + ranked_endpoint_valid.unsqueeze(-1), + pour_collision_candidates, + pour_seed.unsqueeze(1), + ) + pour_collision_candidates[:, :, finger_coordinate_ids] = float(self.cfg.gripper_preload_pos) + tilt_collision_candidates = gather_candidates(tilt_full_candidates, collision_tilt_indices) + tilt_collision_candidates = torch.where( + ranked_endpoint_valid.unsqueeze(-1), + tilt_collision_candidates, + tilt_seed.unsqueeze(1), + ) + tilt_collision_candidates[:, :, finger_coordinate_ids] = float(self.cfg.gripper_preload_pos) + + source_collision_candidates = ( + gather_source_waypoint(expanded_q, float(self.cfg.gripper_open_pos)), + gather_source_waypoint(pregrasp_candidates, float(self.cfg.gripper_open_pos)), + gather_source_waypoint(midgrasp_candidates, float(self.cfg.gripper_open_pos)), + gather_source_waypoint(grasp_candidates, float(self.cfg.gripper_open_pos)), + gather_source_waypoint(carry_candidates, float(self.cfg.gripper_preload_pos)), + ) + collision_free = self._collision_free_ik_candidates( + prototype_builder, + (*source_collision_candidates, pour_collision_candidates, tilt_collision_candidates), + ) + ranked_valid = ranked_endpoint_valid & collision_free + row_has_collision_free_path = ranked_valid.any(dim=-1) + ranked_complete_path_score = torch.gather(flat_complete_path_score, 1, collision_path_indices) + + # Resolve the redundant source-side IK branch against one exact nominal reach path. The + # nominal row is solved in the same batched census and screened through the same complete + # collision path, so this is a physical branch anchor rather than a hard-coded posture. + # Lexicographic selection (nearest source branch, then shortest complete path) prevents a + # millimetric geometry increment from changing elbow/wrist branches at curriculum promotion. + nominal_reach_row = randomized_bank_size + zero_jitter_slot + nominal_seed_valid = ranked_valid[nominal_reach_row] & (collision_source_indices[nominal_reach_row] == 0) + if not bool(nominal_seed_valid.any()): + raise RuntimeError( + "The randomized curriculum requires a collision-free deterministic zero-jitter nominal reach path." + ) + source_preliminary_score = source_transition + nominal_branch_weight * source_nominal_deviation + screened_source_score = torch.gather(source_preliminary_score, 1, collision_source_indices) + nominal_reach_score = torch.where( + nominal_seed_valid, + screened_source_score[nominal_reach_row], + torch.full_like(screened_source_score[nominal_reach_row], torch.inf), + ) + nominal_reach_slot = torch.argmin(nominal_reach_score) + source_collision_arm_paths = torch.stack( + tuple(candidate[:, :, arm_coordinate_ids] for candidate in source_collision_candidates), + dim=2, + ) + nominal_source_arm_path = source_collision_arm_paths[nominal_reach_row, nominal_reach_slot] + source_branch_distance = (source_collision_arm_paths - nominal_source_arm_path).square().sum(dim=(-1, -2)) + valid_source_branch_distance = torch.where( + ranked_valid, + source_branch_distance, + torch.full_like(source_branch_distance, torch.inf), + ) + closest_source_branch = ( + valid_source_branch_distance <= valid_source_branch_distance.amin(dim=-1, keepdim=True) + 1.0e-6 + ) + deterministic_receiver_pair = (collision_pour_indices == 0) & (collision_tilt_indices == 0) + deterministic_receiver_available = (ranked_valid & closest_source_branch & deterministic_receiver_pair).any( + dim=-1, keepdim=True + ) + preferred_receiver_pair = ~deterministic_receiver_available | deterministic_receiver_pair + collision_selection_score = torch.where( + ranked_valid & closest_source_branch & preferred_receiver_pair, + ranked_complete_path_score, + torch.full_like(ranked_complete_path_score, torch.inf), + ) + # Some Cartesian samples near the requested workspace boundary admit endpoint IK but no + # executable self-collision-free branch. Keep finite fallback data for diagnostics, then + # exclude those rows from the reset pools below instead of weakening physical collision + # constraints or biasing the controller with an impossible reference. + collision_selection_score = torch.where( + row_has_collision_free_path.unsqueeze(-1), + collision_selection_score, + ranked_complete_path_score, + ) + best_collision_slots = torch.argmin(collision_selection_score, dim=-1) + best_indices = collision_source_indices[candidate_rows, best_collision_slots] + best_pour_indices = collision_pour_indices[candidate_rows, best_collision_slots] + best_tilt_indices = collision_tilt_indices[candidate_rows, best_collision_slots] + collision_free_count = ranked_valid.sum(dim=-1) + arm_q = expanded_arm_q[candidate_rows, best_indices].clone() + costs = expanded_costs[candidate_rows, best_indices].clone() + margin = expanded_margin[candidate_rows, best_indices].clone() + pregrasp_arm_q = pregrasp_arm_candidates[candidate_rows, best_indices].clone() + pregrasp_costs = pregrasp_cost_candidates[candidate_rows, best_indices].clone() + pregrasp_margin = pregrasp_margin_candidates[candidate_rows, best_indices].clone() + midgrasp_arm_q = midgrasp_arm_candidates[candidate_rows, best_indices].clone() + midgrasp_costs = midgrasp_cost_candidates[candidate_rows, best_indices].clone() + midgrasp_margin = midgrasp_margin_candidates[candidate_rows, best_indices].clone() + grasp_arm_q = grasp_arm_candidates[candidate_rows, best_indices].clone() + grasp_costs = grasp_cost_candidates[candidate_rows, best_indices].clone() + grasp_margin = grasp_margin_candidates[candidate_rows, best_indices].clone() + carry_arm_q = carry_arm_candidates[candidate_rows, best_indices].clone() + carry_costs = carry_cost_candidates[candidate_rows, best_indices].clone() + carry_margin = carry_margin_candidates[candidate_rows, best_indices].clone() + pour_arm_q = pour_arm_candidates[candidate_rows, best_pour_indices] + pour_costs = pour_cost_candidates[candidate_rows, best_pour_indices] + pour_margin = pour_margin_candidates[candidate_rows, best_pour_indices] + tilt_arm_q = tilt_arm_candidates[candidate_rows, best_tilt_indices] + tilt_costs = tilt_cost_candidates[candidate_rows, best_tilt_indices] + tilt_margin = tilt_margin_candidates[candidate_rows, best_tilt_indices] + + # Raw census rows that fail the complete-path screen remain useful for coverage diagnostics, + # but must never leave nonfinite joint data in the bank. Store an explicit invalid sentinel + # in their scalar metadata and a finite nominal trajectory in their joint arrays; sampling + # pools below contain only rows with a positive collision-free candidate count. + selected_valid = row_has_collision_free_path.unsqueeze(-1) + nominal_waypoints = self._nominal_reference_waypoints_t + arm_q = torch.where(selected_valid, arm_q, nominal_waypoints[0]) + pregrasp_arm_q = torch.where(selected_valid, pregrasp_arm_q, nominal_waypoints[1]) + midgrasp_arm_q = torch.where(selected_valid, midgrasp_arm_q, nominal_waypoints[2]) + grasp_arm_q = torch.where(selected_valid, grasp_arm_q, nominal_waypoints[3]) + carry_arm_q = torch.where(selected_valid, carry_arm_q, nominal_waypoints[4]) + pour_arm_q = torch.where(selected_valid, pour_arm_q, nominal_waypoints[5]) + tilt_arm_q = torch.where(selected_valid, tilt_arm_q, nominal_waypoints[6]) + invalid_cost = torch.full_like(costs, torch.inf) + invalid_margin = torch.full_like(margin, -torch.inf) + costs = torch.where(row_has_collision_free_path, costs, invalid_cost) + pregrasp_costs = torch.where(row_has_collision_free_path, pregrasp_costs, invalid_cost) + midgrasp_costs = torch.where(row_has_collision_free_path, midgrasp_costs, invalid_cost) + grasp_costs = torch.where(row_has_collision_free_path, grasp_costs, invalid_cost) + carry_costs = torch.where(row_has_collision_free_path, carry_costs, invalid_cost) + pour_costs = torch.where(row_has_collision_free_path, pour_costs, invalid_cost) + tilt_costs = torch.where(row_has_collision_free_path, tilt_costs, invalid_cost) + margin = torch.where(row_has_collision_free_path, margin, invalid_margin) + pregrasp_margin = torch.where(row_has_collision_free_path, pregrasp_margin, invalid_margin) + midgrasp_margin = torch.where(row_has_collision_free_path, midgrasp_margin, invalid_margin) + grasp_margin = torch.where(row_has_collision_free_path, grasp_margin, invalid_margin) + carry_margin = torch.where(row_has_collision_free_path, carry_margin, invalid_margin) + pour_margin = torch.where(row_has_collision_free_path, pour_margin, invalid_margin) + tilt_margin = torch.where(row_has_collision_free_path, tilt_margin, invalid_margin) + + # Level zero is an exact behavioral continuation of the mastered full task. Its first four + # waypoints must come from the dedicated nominal-source reach bank; the authored waypoint + # tensor intentionally repeats ``arm_home`` there and is only a placeholder until those + # insertion poses are solved. Store the exact zero-jitter nominal reach path in one + # canonical zero-bank row, then preserve the known-safe authored carry/pour/tilt branch. + # Arm-start diversity is introduced continuously by the later nonzero offset extents. + arm_q[zero_bank_row] = arm_q[nominal_reach_row] + pregrasp_arm_q[zero_bank_row] = pregrasp_arm_q[nominal_reach_row] + midgrasp_arm_q[zero_bank_row] = midgrasp_arm_q[nominal_reach_row] + grasp_arm_q[zero_bank_row] = grasp_arm_q[nominal_reach_row] + carry_arm_q[zero_bank_row] = nominal_waypoints[4] + pour_arm_q[zero_bank_row] = nominal_waypoints[5] + tilt_arm_q[zero_bank_row] = nominal_waypoints[6] + costs[zero_bank_row] = costs[nominal_reach_row] + pregrasp_costs[zero_bank_row] = pregrasp_costs[nominal_reach_row] + midgrasp_costs[zero_bank_row] = midgrasp_costs[nominal_reach_row] + grasp_costs[zero_bank_row] = grasp_costs[nominal_reach_row] + carry_costs[zero_bank_row] = carry_costs[nominal_reach_row] + pour_costs[zero_bank_row] = pour_costs[nominal_reach_row] + tilt_costs[zero_bank_row] = tilt_costs[nominal_reach_row] + margin[zero_bank_row] = margin[nominal_reach_row] + pregrasp_margin[zero_bank_row] = pregrasp_margin[nominal_reach_row] + midgrasp_margin[zero_bank_row] = midgrasp_margin[nominal_reach_row] + grasp_margin[zero_bank_row] = grasp_margin[nominal_reach_row] + carry_margin[zero_bank_row] = carry_margin[nominal_reach_row] + pour_margin[zero_bank_row] = pour_margin[nominal_reach_row] + tilt_margin[zero_bank_row] = tilt_margin[nominal_reach_row] + collision_free_count[zero_bank_row] = collision_free_count[nominal_reach_row] + row_has_collision_free_path[zero_bank_row] = True + + randomized_rows = slice(0, randomized_bank_size) + reach_rows = slice(randomized_bank_size, bank_size) + self._randomized_source_pos_bank_t = source_positions[randomized_rows] + self._randomized_source_yaw_bank_t = source_yaws[randomized_rows] + self._randomized_source_quat_bank_t = source_quaternions[randomized_rows] + self._randomized_target_pos_bank_t = target_positions[randomized_rows] + self._randomized_tcp_pos_bank_t = tcp_positions[randomized_rows] + self._randomized_tcp_quat_bank_t = target_rotations_w[randomized_rows] + self._randomized_tcp_rotation_vector_bank_t = tcp_rotation_vectors[randomized_rows] + self._randomized_arm_q_bank_t = arm_q[randomized_rows] + self._randomized_pregrasp_arm_q_bank_t = pregrasp_arm_q[randomized_rows] + self._randomized_midgrasp_arm_q_bank_t = midgrasp_arm_q[randomized_rows] + self._randomized_grasp_arm_q_bank_t = grasp_arm_q[randomized_rows] + self._randomized_carry_arm_q_bank_t = carry_arm_q[randomized_rows] + self._randomized_pour_arm_q_bank_t = pour_arm_q[randomized_rows] + self._randomized_tilt_arm_q_bank_t = tilt_arm_q[randomized_rows] + self._randomized_reset_ik_cost_t = costs[randomized_rows] + self._randomized_reset_ik_margin_t = margin[randomized_rows] + self._randomized_collision_free_candidate_count_t = collision_free_count[randomized_rows] + self._randomized_pregrasp_ik_cost_t = pregrasp_costs[randomized_rows] + self._randomized_pregrasp_ik_margin_t = pregrasp_margin[randomized_rows] + self._randomized_midgrasp_ik_cost_t = midgrasp_costs[randomized_rows] + self._randomized_midgrasp_ik_margin_t = midgrasp_margin[randomized_rows] + self._randomized_grasp_ik_cost_t = grasp_costs[randomized_rows] + self._randomized_grasp_ik_margin_t = grasp_margin[randomized_rows] + self._randomized_carry_ik_cost_t = carry_costs[randomized_rows] + self._randomized_carry_ik_margin_t = carry_margin[randomized_rows] + self._randomized_pour_ik_cost_t = pour_costs[randomized_rows] + self._randomized_pour_ik_margin_t = pour_margin[randomized_rows] + self._randomized_tilt_ik_cost_t = tilt_costs[randomized_rows] + self._randomized_tilt_ik_margin_t = tilt_margin[randomized_rows] + extent_bank_indices = torch.arange(randomized_bank_size, device=self.device).reshape( + level_count, + rows_per_extent, + ) + randomized_collision_free = ( + row_has_collision_free_path[randomized_rows] + & (arm_q[randomized_rows, 5] <= float(self.cfg.curriculum_randomized_reset_joint6_max)) + ).reshape(level_count, rows_per_extent) + randomized_reference_path = torch.stack( + ( + arm_q[randomized_rows], + pregrasp_arm_q[randomized_rows], + midgrasp_arm_q[randomized_rows], + grasp_arm_q[randomized_rows], + carry_arm_q[randomized_rows], + pour_arm_q[randomized_rows], + tilt_arm_q[randomized_rows], + ), + dim=1, + ) + reference_branch_delta = torch.abs(randomized_reference_path - randomized_reference_path[zero_bank_row]).amax( + dim=(-1, -2) + ) + extent_by_row = reference_branch_delta.new_tensor(extent_levels).repeat_interleave(rows_per_extent) + # Early frontiers should expand task geometry, not introduce unrelated redundant IK + # branches. Permit joint displacement proportional to physical extent, then remove this + # continuity filter once the policy has mastered half of the randomization amplitude. + branch_continuity_limit = 0.04 + 3.0 * extent_by_row + branch_continuity_limit = torch.where( + extent_by_row >= 0.5, + torch.full_like(branch_continuity_limit, torch.inf), + branch_continuity_limit, + ) + randomized_collision_free &= (reference_branch_delta <= branch_continuity_limit).reshape( + level_count, rows_per_extent + ) + self._randomized_reference_branch_delta_t = reference_branch_delta + self._randomized_reference_branch_limit_t = branch_continuity_limit + source_cell_ids = torch.arange(rows_per_extent, device=self.device) // samples_per_source + # The exact zero-amplitude anchor is one canonical task, not 539 duplicate randomization + # rows. Later levels retain equal source-cell weighting while expanding arm-start diversity + # with the physical extent. Exposing every redundant IK/reset variant at the first one- + # percent geometry step creates a discrete task jump even though the cup poses barely move. + index_pools = [torch.as_tensor((zero_bank_row,), device=self.device)] + weight_pools = [torch.ones(1, device=self.device)] + source_cell_counts = [1] + minimum_variant_counts = [1] + minimum_variants = int(self.cfg.curriculum_randomized_min_reset_variants_per_source) + for extent, level_indices, level_valid in zip( + extent_levels[1:], + extent_bank_indices[1:].unbind(dim=0), + randomized_collision_free[1:].unbind(dim=0), + strict=True, + ): + variants_per_cell = torch.bincount(source_cell_ids[level_valid], minlength=source_cell_count) + level_valid = level_valid & (variants_per_cell >= minimum_variants)[source_cell_ids] + + # Keep the safest, most continuous paths first and grow toward the complete validated + # arm-start marginal. The final extent selects all sample slots, while early levels + # isolate learning the new Cartesian geometry from a simultaneous redundant-IK jump. + variant_limit = max(1, math.ceil(extent * samples_per_source)) + score_by_cell = torch.where( + level_valid, + reference_branch_delta[level_indices], + torch.full_like(reference_branch_delta[level_indices], torch.inf), + ).reshape(source_cell_count, samples_per_source) + ordered_slots = torch.argsort(score_by_cell, dim=1, stable=True) + selected_by_cell = torch.zeros_like(score_by_cell, dtype=torch.bool) + selected_by_cell.scatter_(1, ordered_slots[:, :variant_limit], True) + level_valid &= selected_by_cell.reshape(-1) + + index_pool = level_indices[level_valid] + _, inverse_cell_ids, rows_per_cell = torch.unique( + source_cell_ids[level_valid], + sorted=True, + return_inverse=True, + return_counts=True, + ) + # Give every feasible source XY cell equal probability, then sample its validated + # yaw/jitter branches uniformly. Flat row sampling would overrepresent central cells, + # where more IK seeds survive the posture and self-collision screens. + weight_pool = rows_per_cell[inverse_cell_ids].to(dtype=torch.float32).reciprocal() + index_pools.append(index_pool) + weight_pools.append(weight_pool) + source_cell_counts.append(int(rows_per_cell.numel())) + minimum_variant_counts.append(int(rows_per_cell.min().item()) if rows_per_cell.numel() else 0) + self._randomized_extent_index_pools = tuple(index_pools) + self._randomized_extent_index_weights = tuple(weight_pools) + self._randomized_extent_source_cell_counts = tuple(source_cell_counts) + self._randomized_extent_minimum_variant_counts = tuple(minimum_variant_counts) + if any(pool.numel() == 0 for pool in self._randomized_extent_index_pools): + raise RuntimeError("Every randomization extent must contain a collision-free reset pose.") + minimum_source_cells = math.ceil( + grid_size * grid_size * float(self.cfg.curriculum_randomized_min_source_cell_fraction) + ) + if any(count < minimum_source_cells for count in self._randomized_extent_source_cell_counts[1:]): + source_valid_by_level = source_trajectory_valid[randomized_rows].reshape(level_count, rows_per_extent) + receiver_valid_by_level = receiver_valid[randomized_rows].reshape(level_count, rows_per_extent) + + def covered_cell_counts(valid_by_level: torch.Tensor) -> tuple[int, ...]: + return tuple( + int(torch.unique(source_cell_ids[level_valid]).numel()) + for level_valid in valid_by_level.unbind(dim=0) + ) + + raise RuntimeError( + "Randomized reset posture screening retained too little source workspace coverage: " + f"required at least {minimum_source_cells}/{grid_size * grid_size} source XY cells per extent, " + f"got {self._randomized_extent_source_cell_counts}; source-waypoint IK covered " + f"{covered_cell_counts(source_valid_by_level)}, receiver-waypoint IK covered " + f"{covered_cell_counts(receiver_valid_by_level)}." + ) + required_minimum_variant_counts = tuple( + 1 if extent == 0.0 else min(minimum_variants, max(1, math.ceil(extent * samples_per_source))) + for extent in extent_levels + ) + if any( + actual < required + for actual, required in zip( + self._randomized_extent_minimum_variant_counts, + required_minimum_variant_counts, + strict=True, + ) + ): + + def variant_histograms(valid_rows: torch.Tensor) -> tuple[tuple[int, ...], ...]: + return tuple( + tuple( + int((torch.bincount(source_cell_ids[level_valid], minlength=source_cell_count) == count).sum()) + for count in range(samples_per_source + 1) + ) + for level_valid in valid_rows.unbind(dim=0) + ) + + source_valid_by_level = source_trajectory_valid[randomized_rows].reshape(level_count, rows_per_extent) + receiver_valid_by_level = receiver_valid[randomized_rows].reshape(level_count, rows_per_extent) + slot_survival_counts = tuple( + tuple(int(count) for count in level_valid.reshape(source_cell_count, samples_per_source).sum(dim=0)) + for level_valid in randomized_collision_free.unbind(dim=0) + ) + raise RuntimeError( + "Randomized reset posture screening retained too little arm-start diversity: " + f"required per-level minima {required_minimum_variant_counts}, " + f"got {self._randomized_extent_minimum_variant_counts}. " + "Per-level source-cell histograms indexed by retained-variant count are " + f"source-IK={variant_histograms(source_valid_by_level)}, " + f"receiver-IK={variant_histograms(receiver_valid_by_level)}, " + f"collision-free={variant_histograms(randomized_collision_free)}. " + f"Collision-free sample-slot survival counts are {slot_survival_counts}." + ) + if self.cfg.curriculum_randomized_source_radius_range is not None: + final_level_offset = (level_count - 1) * rows_per_extent + final_source_cells = torch.unique( + (self._randomized_extent_index_pools[-1] - final_level_offset) // samples_per_source, + sorted=True, + ) + final_radial_rings = torch.unique( + torch.div(final_source_cells, grid_size, rounding_mode="floor"), + sorted=True, + ) + final_azimuth_slots = torch.unique(final_source_cells.remainder(grid_size), sorted=True) + required_radial_rings = torch.arange(grid_size, device=self.device) + center_azimuth_slot = grid_size // 2 + has_both_azimuth_sides = bool( + (final_azimuth_slots[0] < center_azimuth_slot) & (final_azimuth_slots[-1] > center_azimuth_slot) + ) + has_broad_azimuth_span = bool(final_azimuth_slots[-1] - final_azimuth_slots[0] >= center_azimuth_slot) + if ( + not torch.equal(final_radial_rings, required_radial_rings) + or not has_both_azimuth_sides + or not has_broad_azimuth_span + ): + raise RuntimeError( + "Randomized reset posture screening did not retain the configured polar-workspace " + "coverage: required every radial ring and broad coverage on both azimuth sides, got radial " + f"rings {final_radial_rings.tolist()} and azimuth slots {final_azimuth_slots.tolist()}." + ) + if source_positions[reach_rows].shape[0] != samples_per_source: + raise RuntimeError( + f"Expected {samples_per_source} dedicated nominal-source reach poses, " + f"found {source_positions[reach_rows].shape[0]}." + ) + reach_indices = torch.arange(randomized_bank_size, bank_size, device=self.device) + reach_indices = reach_indices[row_has_collision_free_path[reach_rows]] + if reach_indices.numel() < minimum_variants: + raise RuntimeError( + "Nominal-source reach posture screening retained too little arm-start diversity: " + f"required at least {minimum_variants} variants, got {reach_indices.numel()}." + ) + self._reach_tcp_pos_bank_t = tcp_positions[reach_indices] + self._reach_source_yaw_bank_t = source_yaws[reach_indices] + self._reach_arm_q_bank_t = arm_q[reach_indices] + self._reach_pregrasp_arm_q_bank_t = pregrasp_arm_q[reach_indices] + self._reach_midgrasp_arm_q_bank_t = midgrasp_arm_q[reach_indices] + self._reach_grasp_arm_q_bank_t = grasp_arm_q[reach_indices] + self._reach_reset_ik_cost_t = costs[reach_indices] + self._reach_reset_ik_margin_t = margin[reach_indices] + self._reach_collision_free_candidate_count_t = collision_free_count[reach_indices] + # Newton IK and the zero-copy Warp/Torch views above run asynchronously. The temporary + # solver and prototype are released when this method returns, so complete every gather + # before their backing allocations can be reclaimed. This is a one-time startup barrier. + wp.synchronize_device(model.device) + + def _build_independent_reset_fallbacks(self) -> None: + """Precompute guaranteed-clearance fallbacks for independently mixed reset rows.""" + bank_size = self._randomized_source_pos_bank_t.shape[0] + arm_fallback = torch.full((bank_size,), -1, device=self.device, dtype=torch.long) + target_fallback = torch.full_like(arm_fallback, -1) + minimum_arm_distance = float(self.cfg.curriculum_independent_arm_min_tcp_distance) + + for level, pool in enumerate(self._randomized_extent_index_pools): + source_position = self._randomized_source_pos_bank_t[pool] + if self.cfg.curriculum_independent_arm_fraction_levels[level] > 0.0: + arm_distance = torch.cdist(source_position, self._randomized_tcp_pos_bank_t[pool]) + farthest_arm_slot = torch.argmax(arm_distance, dim=1) + farthest_arm_distance = arm_distance.gather(1, farthest_arm_slot.unsqueeze(-1)).squeeze(-1) + if bool(torch.any(farthest_arm_distance < minimum_arm_distance)): + raise RuntimeError( + f"Randomization level {level} has no independent arm reset with the required " + f"{minimum_arm_distance:.3f} m TCP/source clearance." + ) + arm_fallback[pool] = pool[farthest_arm_slot] + else: + arm_fallback[pool] = pool + + if self.cfg.curriculum_independent_target_fraction_levels[level] > 0.0: + candidate_targets = pool.unsqueeze(0).expand(pool.numel(), -1) + target_clearance = self._independent_target_clearance( + pool, + candidate_targets, + ) + farthest_target_slot = torch.argmax(target_clearance, dim=1) + farthest_target_clearance = target_clearance.gather( + 1, + farthest_target_slot.unsqueeze(-1), + ).squeeze(-1) + if bool(torch.any(farthest_target_clearance < 0.0)): + raise RuntimeError( + f"Randomization level {level} has no independent receiver reset with the configured " + "rectangular cup clearance." + ) + target_fallback[pool] = pool[farthest_target_slot] + else: + target_fallback[pool] = pool + + self._independent_arm_fallback_index_t = arm_fallback + self._independent_target_fallback_index_t = target_fallback + + def _independent_target_clearance( + self, + source_indices: torch.Tensor, + target_indices: torch.Tensor, + ) -> torch.Tensor: + """Return conservative rectangular source/receiver separation margins [m].""" + if source_indices.ndim != 1 or target_indices.ndim != 2 or target_indices.shape[0] != source_indices.shape[0]: + raise ValueError("Independent reset indices must have shapes (N,) and (N, K).") + source_position = self._randomized_source_pos_bank_t[source_indices] + source_yaw = self._randomized_source_yaw_bank_t[source_indices] + target_position = self._randomized_target_pos_bank_t[target_indices] + source_half_x = 0.5 * float(self.cfg.source_cup_inner_width) + float(self.cfg.source_cup_wall_thickness) + source_half_y = 0.5 * float(self.cfg.source_cup_inner_depth) + float(self.cfg.source_cup_wall_thickness) + target_half_x = 0.5 * float(self.cfg.target_cup_inner_width) + float(self.cfg.target_cup_wall_thickness) + target_half_y = 0.5 * float(self.cfg.target_cup_inner_depth) + float(self.cfg.target_cup_wall_thickness) + clearance = float(self.cfg.curriculum_randomized_cup_clearance) + source_aabb_half_x = source_half_x * torch.abs(torch.cos(source_yaw)) + source_half_y * torch.abs( + torch.sin(source_yaw) + ) + source_aabb_half_y = source_half_x * torch.abs(torch.sin(source_yaw)) + source_half_y * torch.abs( + torch.cos(source_yaw) + ) + position_delta = torch.abs(target_position[:, :, :2] - source_position[:, None, :2]) + clearance_x = position_delta[:, :, 0] - (source_aabb_half_x[:, None] + target_half_x + clearance) + clearance_y = position_delta[:, :, 1] - (source_aabb_half_y[:, None] + target_half_y + clearance) + # A positive margin on either world axis is sufficient to separate the two rectangles. + return torch.maximum(clearance_x, clearance_y) + + # ----------------------------------------------------------- poses / obs + def _pose_w_to_e(self, pose_w: torch.Tensor) -> torch.Tensor: + """Convert a public world-frame pose view to a finite environment-frame pose.""" + pos = torch.nan_to_num(pose_w[:, :3], nan=0.0, posinf=0.0, neginf=0.0) - self.env_origins + raw_quat = pose_w[:, 3:7] + quat = torch.nan_to_num(raw_quat, nan=0.0, posinf=0.0, neginf=0.0) + norm = torch.linalg.norm(quat, dim=-1, keepdim=True) + ident = torch.zeros_like(raw_quat) + ident[:, 3] = 1.0 + valid = torch.isfinite(raw_quat).all(dim=-1, keepdim=True) & (norm > 1.0e-6) + quat = torch.where(valid, quat / torch.clamp(norm, min=1.0e-6), ident) + return torch.cat((pos, quat), dim=-1) + + def ee_pose_e(self) -> torch.Tensor: + """End-effector (panda_hand) pose in the env frame: ``(num_envs, 7)`` pos + xyzw quat.""" + return self._pose_w_to_e(self._robot.data.body_link_pose_w.torch[:, self._tcp_body_idx]) + + def cup_pose_e(self) -> torch.Tensor: + """Cup body pose in the env frame: ``(num_envs, 7)`` pos + xyzw quat.""" + return self._pose_w_to_e(self._source_cup.data.root_link_pose_w.torch) + + def cup_velocity_w(self) -> torch.Tensor: + """Source-cup linear and angular velocity in the world frame [m/s, rad/s].""" + return self._source_cup.data.root_link_vel_w.torch + + def target_pose_e(self) -> torch.Tensor: + """Receiving-cup pose in the env frame: ``(num_envs, 7)`` pos + xyzw quat.""" + return self._pose_w_to_e(self._target_cup.data.root_link_pose_w.torch) + + def tcp_pose_e(self) -> torch.Tensor: + """Tool-centre pose in the robot-root/environment frame.""" + body_pose_w = self._robot.data.body_link_pose_w.torch[:, self._tcp_body_idx] + root_pose_w = self._robot.data.root_link_pose_w.torch + pos, quat = math_utils.subtract_frame_transforms( + root_pose_w[:, :3], root_pose_w[:, 3:7], body_pose_w[:, :3], body_pose_w[:, 3:7] + ) + pos, quat = math_utils.combine_frame_transforms(pos, quat, self._tcp_offset_pos, self._tcp_offset_quat) + return torch.cat((torch.nan_to_num(pos), torch.nan_to_num(quat)), dim=-1) + + def tcp_pos_e(self) -> torch.Tensor: + return self.tcp_pose_e()[:, :3] + + def grasp_approach_error(self) -> tuple[torch.Tensor, torch.Tensor]: + """Return signed axial and unsigned cross-track TCP errors from the grasp point [m].""" + cup_pose = self.cup_pose_e() + error_e = self.tcp_pos_e() - self.cup_grasp_point_e() + error_c = math_utils.quat_apply_inverse(cup_pose[:, 3:7], error_e) + axial = torch.sum(error_c * self._grasp_approach_axis_c, dim=-1) + cross_track = torch.linalg.vector_norm( + error_c - axial.unsqueeze(-1) * self._grasp_approach_axis_c, + dim=-1, + ) + return axial, cross_track + + def cup_grasp_point_e(self) -> torch.Tensor: + """World-facing grasp point at the middle of the source cup walls, in env coordinates.""" + pose = self.cup_pose_e() + offset = torch.zeros((self.num_envs, 3), device=self.device) + offset[:, 2] = float(self.cfg.cup_grasp_height) + return pose[:, :3] + math_utils.quat_apply(pose[:, 3:7], offset) + + def gripper_width(self) -> torch.Tensor: + """Distance represented by the two symmetric Panda finger joint positions [m].""" + finger_pos = self._robot.data.joint_pos.torch[:, self._finger_joint_ids] + width = finger_pos.sum(dim=-1) + valid = torch.isfinite(finger_pos).all(dim=-1) & torch.isfinite(width) + return torch.where(valid, width, torch.full_like(width, float(self.gripper_open_width))) + + def finger_joint_pos(self) -> torch.Tensor: + """Individual policy-controlled finger joint positions [m].""" + return self._robot.data.joint_pos.torch[:, self._finger_joint_ids] + + def finger_joint_vel(self) -> torch.Tensor: + """Individual policy-controlled finger joint velocities [m/s].""" + return self._robot.data.joint_vel.torch[:, self._finger_joint_ids] + + def desired_grasp_tcp_quat_c(self) -> torch.Tensor: + """Desired TCP orientation in the source-cup frame as canonical XYZW quaternions.""" + return self._desired_grasp_tcp_quat_c + + def arm_joint_pos(self) -> torch.Tensor: + """Current policy-controlled arm joint positions [rad].""" + return self._robot.data.joint_pos.torch[:, self._arm_joint_ids] + + @property + def gripper_open_width(self) -> float: + return 2.0 * float(self.cfg.gripper_open_pos) + + @property + def gripper_grasp_width(self) -> float: + return 2.0 * float(self.cfg.cup_grasp_box_half[1]) + + @property + def cup_reset_height(self) -> float: + return float(self.cfg.cup_reset_pos[2]) + + @property + def num_particles(self) -> int: + return self._num_particles + + def set_curriculum_stage( + self, + env_ids: list[int] | torch.Tensor | slice, + stage: int, + ) -> None: + """Assign a curriculum stage and success threshold to selected environments.""" + if stage < 0 or stage >= len(self.cfg.curriculum_stage_names): + raise ValueError(f"Curriculum stage {stage} is out of range.") + self.curriculum_stage[env_ids] = stage + self.pour_target_frac[env_ids] = float(self.cfg.curriculum_target_frac[stage]) + + def set_curriculum_randomization_level( + self, + env_ids: list[int] | torch.Tensor | slice, + level: int, + ) -> None: + """Assign one prevalidated source-randomization extent to selected environments.""" + if level < 0 or level >= len(self._randomized_extent_index_pools): + raise ValueError(f"Curriculum randomization level {level} is out of range.") + self.curriculum_randomization_level[env_ids] = level + + def particle_pos_e(self) -> torch.Tensor: + """Per-env MPM particle positions in env coordinates, shape ``(N, P, 3)``.""" + return self._media.data.particle_pos_w.torch - self.env_origins[:, None, :] + + def particle_vel_e(self) -> torch.Tensor: + """Per-env MPM particle velocities in environment axes, shape ``(N, P, 3)``.""" + # Environments differ only by translation, so world and environment velocity axes coincide. + return self._media.data.particle_vel_w.torch + + def _points_inside_cup( + self, points_e: torch.Tensor, pose_e: torch.Tensor, lo: torch.Tensor, hi: torch.Tensor + ) -> torch.Tensor: + rel = points_e - pose_e[:, None, :3] + quat = pose_e[:, None, 3:7].expand(-1, points_e.shape[1], -1) + local = math_utils.quat_apply_inverse(quat, rel) + margin = float(self.cfg.particle_count_margin) + return ((local >= lo - margin) & (local <= hi + margin)).all(dim=-1) + + def _particle_region_masks(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Boolean ``(source, target, spilled)`` masks, cached within one manager step.""" + step = int(getattr(self, "common_step_counter", -1)) + if self._particle_region_cache is not None and self._particle_region_cache_step == step: + return self._particle_region_cache + points = self.particle_pos_e() + source = self._points_inside_cup(points, self.cup_pose_e(), self._source_inner_lo_t, self._source_inner_hi_t) + target_region = self._points_inside_cup( + points, + self.target_pose_e(), + self._target_inner_lo_t, + self._target_inner_hi_t, + ) + # Geometric overlap is not delivery: nesting the source cup inside the receiver must not + # score particles that remain physically contained by the source cup. + target = _delivered_particle_mask(source, target_region) + spill_height = float(self.cfg.spill_table_height) + float(self.cfg.particle_count_margin) + spilled = _spilled_particle_mask(points, source, target, max_height=spill_height) + self._particle_region_cache = (source, target, spilled) + self._particle_region_cache_step = step + return source, target, spilled + + def particles_in_target_mask(self) -> torch.Tensor: + """Particles inside the target cup and no longer inside the source cup.""" + return self._particle_region_masks()[1] + + def particle_region_masks(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return cached source, source-exclusive target, and irreversible-spill masks.""" + return self._particle_region_masks() + + def update_held_delivery_tracker(self, held_pour: torch.Tensor) -> None: + """Record particles whose target-entry edge occurs during a held pour. + + The tracker is idempotent within one manager step because success termination and reward + evaluation consume the same state. An unheld entry does not permanently disqualify a + particle: after it leaves the target, a later valid held re-entry can still qualify. + + Args: + held_pour: Per-environment mask for a preloaded, lifted source grasp. + """ + if held_pour.shape != (self.num_envs,): + raise ValueError(f"held_pour must have shape ({self.num_envs},), got {tuple(held_pour.shape)}.") + step = int(self.common_step_counter) + if self._held_delivery_tracker_step == step: + return + in_target = self.particles_in_target_mask() + target_entry = in_target & ~self._target_entry_seen + self._held_delivered |= target_entry & held_pour.unsqueeze(-1) + self._target_entry_seen.copy_(in_target) + self._held_delivery_tracker_step = step + + def held_delivered_mask(self) -> torch.Tensor: + """Particles with at least one target-entry edge during a held pour.""" + return self._held_delivered + + def current_held_delivered_mask(self) -> torch.Tensor: + """Validly delivered particles that remain inside the receiving cup.""" + return self.particles_in_target_mask() & self._held_delivered + + def particles_spilled_mask(self) -> torch.Tensor: + """Per-particle irreversible-spill membership used by one-time penalties.""" + return self._particle_region_masks()[2] + + def count_in_source(self) -> torch.Tensor: + return self._particle_region_masks()[0].sum(dim=1).float() + + def count_in_target(self) -> torch.Tensor: + return self._particle_region_masks()[1].sum(dim=1).float() + + def count_spilled(self) -> torch.Tensor: + return self._particle_region_masks()[2].sum(dim=1).float() + + def spilled_fraction(self) -> torch.Tensor: + return self.count_spilled() / max(self.num_particles, 1) + + def state_finite(self) -> torch.Tensor: + """Per-env instability guard over robot, source cup, and MPM media state.""" + cup_velocity = self._source_cup.data.root_link_vel_w.torch + return _state_finite( + self._robot.data.joint_pos.torch, + self._robot.data.joint_vel.torch, + self._robot.data.body_link_pose_w.torch[:, self._tcp_body_idx], + self._source_cup.data.root_link_pose_w.torch, + cup_velocity[:, :3], + cup_velocity[:, 3:], + self._media.data.particle_pos_w.torch, + ) + + def rigid_state_in_bounds(self) -> torch.Tensor: + """Return whether finite rigid state remains within task-safe observation bounds.""" + cup_velocity = self._source_cup.data.root_link_vel_w.torch + hand_pose_w = self._robot.data.body_link_pose_w.torch[:, self._tcp_body_idx] + tcp_position_w, tcp_quaternion_w = math_utils.combine_frame_transforms( + hand_pose_w[:, :3], + hand_pose_w[:, 3:7], + self._tcp_offset_pos, + self._tcp_offset_quat, + ) + tcp_pose_w = torch.cat((tcp_position_w, tcp_quaternion_w), dim=-1) + return _rigid_state_in_bounds( + self._robot.data.joint_pos.torch, + self._robot.data.joint_vel.torch, + self._joint_pos_limits_t, + tcp_pose_w, + self._source_cup.data.root_link_pose_w.torch, + cup_velocity[:, :3], + cup_velocity[:, 3:], + self.env_origins, + self._particle_workspace_lower_t, + self._particle_workspace_upper_t, + joint_position_margin=self.cfg.state_bound_joint_position_margin, + max_joint_velocity=self.cfg.state_bound_max_joint_velocity, + max_cup_linear_velocity=self.cfg.state_bound_max_cup_linear_velocity, + max_cup_angular_velocity=self.cfg.state_bound_max_cup_angular_velocity, + ) + + def particles_in_workspace(self) -> torch.Tensor: + """Return a per-environment mask for media inside the configured local workspace.""" + return _particles_in_workspace( + self.particle_pos_e(), + self._particle_workspace_lower_t, + self._particle_workspace_upper_t, + ) + + @staticmethod + def _select_first_safe_candidate( + candidates: torch.Tensor, + safe: torch.Tensor, + fallback: torch.Tensor, + ) -> torch.Tensor: + """Select the first safe sampled index, or a prevalidated fallback.""" + first_safe = torch.argmax(safe.to(dtype=torch.int32), dim=1) + selected = candidates[torch.arange(candidates.shape[0], device=candidates.device), first_safe] + return torch.where(safe.any(dim=1), selected, fallback) + + def _sample_independent_reset_indices( + self, + source_indices: torch.Tensor, + levels: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Sample collision-screened arm and receiver rows independently from the source row.""" + attempts = int(self.cfg.curriculum_independent_sample_attempts) + candidate_columns = [ + sample_index_pools( + self._randomized_extent_index_pools, + levels, + weights=self._randomized_extent_index_weights, + ) + for _ in range(attempts) + ] + candidates = torch.stack(candidate_columns, dim=1) + rows = torch.arange(source_indices.shape[0], device=self.device) + source_position = self._randomized_source_pos_bank_t[source_indices] + + candidate_tcp = self._randomized_tcp_pos_bank_t[candidates] + tcp_distance = torch.linalg.vector_norm(candidate_tcp - source_position[:, None, :], dim=-1) + arm_safe = tcp_distance >= float(self.cfg.curriculum_independent_arm_min_tcp_distance) + independent_arm = self._select_first_safe_candidate( + candidates, + arm_safe, + self._independent_arm_fallback_index_t[source_indices], + ) + + target_clearance = self._independent_target_clearance(source_indices, candidates) + target_safe = target_clearance >= 0.0 + independent_target = self._select_first_safe_candidate( + candidates, + target_safe, + self._independent_target_fallback_index_t[source_indices], + ) + + arm_fraction = source_position.new_tensor(self.cfg.curriculum_independent_arm_fraction_levels)[levels] + target_fraction = source_position.new_tensor(self.cfg.curriculum_independent_target_fraction_levels)[levels] + arm_indices = torch.where(torch.rand_like(arm_fraction) < arm_fraction, independent_arm, source_indices) + target_indices = torch.where( + torch.rand_like(target_fraction) < target_fraction, + independent_target, + source_indices, + ) + # Keep this local variable explicit: it makes shape errors in future pool changes fail near + # sampling rather than later during indexed assignment. + if arm_indices.shape != rows.shape or target_indices.shape != rows.shape: + raise RuntimeError("Independent reset sampling returned an invalid index shape.") + return arm_indices, target_indices + + # ----------------------------------------------------------- reset + def _reset_from_dataset(self, env_ids: torch.Tensor, world_mask: torch.Tensor) -> None: + """Restore exact dataset rows and clear all per-world solver history.""" + rows = self.reset_dataset_row_id[env_ids] + if bool(torch.any((rows < 0) | (rows >= self._reset_dataset_states["category"].numel()))): + raise RuntimeError("The reset-dataset curriculum must assign every environment a valid row.") + states = self._reset_dataset_states + arm_q = states["arm_joint_position"][rows] + arm_qd = states["arm_joint_velocity"][rows] + finger_q = states["finger_joint_position"][rows] + finger_qd = states["finger_joint_velocity"][rows] + finger_target = states["finger_joint_target"][rows] + + self._robot.write_joint_position_to_sim_index( + position=arm_q, + joint_ids=self._arm_joint_ids, + env_ids=env_ids, + ) + self._robot.write_joint_velocity_to_sim_index( + velocity=arm_qd, + joint_ids=self._arm_joint_ids, + env_ids=env_ids, + ) + self._robot.set_joint_position_target_index( + target=arm_q, + joint_ids=self._arm_joint_ids, + env_ids=env_ids, + ) + self._robot.write_joint_position_to_sim_index( + position=finger_q, + joint_ids=self._finger_joint_ids, + env_ids=env_ids, + ) + self._robot.write_joint_velocity_to_sim_index( + velocity=finger_qd, + joint_ids=self._finger_joint_ids, + env_ids=env_ids, + ) + self._robot.set_joint_position_target_index( + target=finger_target, + joint_ids=self._finger_joint_ids, + env_ids=env_ids, + ) + self.action_manager.get_term("gripper_action").set_reset_position( + finger_target[:, :1], + env_ids=env_ids, + ) + # Consume the public FK invalidation before rigid proxies and observations access bodies. + _ = self._robot.data.body_link_pose_w + + source_pose = states["source_root_pose"][rows].clone() + source_pose[:, :3] += self.env_origins[env_ids] + target_pose = states["target_root_pose"][rows].clone() + target_pose[:, :3] += self.env_origins[env_ids] + self._source_cup.write_root_pose_to_sim_index(root_pose=source_pose, env_ids=env_ids) + self._source_cup.write_root_velocity_to_sim_index( + root_velocity=states["source_root_velocity"][rows], + env_ids=env_ids, + ) + self._target_cup.write_root_pose_to_sim_index(root_pose=target_pose, env_ids=env_ids) + self._target_cup.write_root_velocity_to_sim_index( + root_velocity=states["target_root_velocity"][rows], + env_ids=env_ids, + ) + + layout_ids = states["particle_layout_id"][rows].long() + local_position = self._reset_dataset_particle_local_position[layout_ids] + local_velocity = self._reset_dataset_particle_local_velocity[layout_ids] + particle_count = local_position.shape[1] + source_quat = source_pose[:, None, 3:7].expand(-1, particle_count, -1) + particle_position = math_utils.quat_apply(source_quat, local_position) + source_pose[:, None, :3] + particle_velocity = math_utils.quat_apply(source_quat, local_velocity) + self._media.write_particle_pos_to_sim_index(particle_position, env_ids=env_ids) + self._media.write_particle_velocity_to_sim_index(particle_velocity, env_ids=env_ids) + + # Public particle writers restore both Newton state buffers. This masked solver reset then + # clears MPM stress/deformation and every private contact/collider history for the worlds. + NewtonManager.reset_solver_state( + world_mask=None if self.num_envs == 1 else wp.from_torch(world_mask, dtype=wp.bool), + flags=newton.StateFlags.BODY | newton.StateFlags.PARTICLE, + ) + self._last_source_bank_index[env_ids] = -1 + self._last_arm_bank_index[env_ids] = -1 + self._last_target_bank_index[env_ids] = -1 + self._particle_region_cache = None + self._particle_region_cache_step = -1 + self.episode_succeeded[env_ids] = False + self.ep_max_target_frac[env_ids] = 0.0 + self._success_dwell_count[env_ids] = 0 + self._lost_grasp_dwell_count[env_ids] = 0 + # A grasping cache row is an offline-validated demonstrated grasp. Seed the latch from the + # row category so opening the hand immediately after reset cannot evade dropped-cup + # termination before the first runtime grasp observation. Non-grasping rows remain + # unlatched and may freely approach the cup. + self._lifted_grasp_seen[env_ids] = states["category"][rows] == GRASPING_CATEGORY + self._target_entry_seen[env_ids] = False + self._held_delivered[env_ids] = False + self._held_delivery_tracker_step = -1 + + def reset_pour_scene(self, env_ids: torch.Tensor) -> None: + """Reset the arm, source cup, and particles through their public asset APIs.""" + if not isinstance(env_ids, torch.Tensor): + env_ids = torch.as_tensor(list(env_ids), device=self.device, dtype=torch.long) + env_ids = env_ids.to(device=self.device, dtype=torch.long) + if env_ids.numel() == 0: + return + world_mask = boolean_selection_mask(self.num_envs, env_ids) + if self._uses_reset_dataset: + self._reset_from_dataset(env_ids, world_mask) + return + n = env_ids.numel() + stage = self.curriculum_stage[env_ids] + arm_q = self._curriculum_arm_q_t[stage].clone() + cup_pos_e = torch.as_tensor(self.cfg.cup_reset_pos, device=self.device).repeat(n, 1) + target_pos_e = torch.as_tensor(self.cfg.target_cup_reset_pos, device=self.device).repeat(n, 1) + source_quat = self._curriculum_cup_quat_t[stage].clone() + target_quat = self._curriculum_cup_quat_t[stage].clone() + self._last_source_bank_index[env_ids] = -1 + self._last_arm_bank_index[env_ids] = -1 + self._last_target_bank_index[env_ids] = -1 + grasp_rows = torch.nonzero(stage == self._grasp_stage_index, as_tuple=False).flatten() + if grasp_rows.numel() > 0: + arm_q[grasp_rows] = self._reach_grasp_arm_q_bank_t[0] + approach_rows = torch.nonzero( + (stage >= self._approach_stage_index) & (stage < self._full_stage_index), + as_tuple=False, + ).flatten() + if approach_rows.numel() > 0: + approach_indices = stage[approach_rows] - self._approach_stage_index + arm_q[approach_rows] = self._curriculum_approach_arm_q_t[approach_indices] + full_rows = torch.nonzero(stage == self._full_stage_index, as_tuple=False).flatten() + if full_rows.numel() > 0: + arm_q[full_rows] = self._reach_pregrasp_arm_q_bank_t[0] + randomized_rows = torch.nonzero(stage == self._randomized_stage_index, as_tuple=False).flatten() + if randomized_rows.numel() > 0: + randomization_levels = self.curriculum_randomization_level[env_ids[randomized_rows]] + source_indices = sample_index_pools( + self._randomized_extent_index_pools, + randomization_levels, + weights=self._randomized_extent_index_weights, + ) + arm_indices, target_indices = self._sample_independent_reset_indices( + source_indices, + randomization_levels, + ) + arm_q[randomized_rows] = self._randomized_arm_q_bank_t[arm_indices] + cup_pos_e[randomized_rows] = self._randomized_source_pos_bank_t[source_indices] + source_quat[randomized_rows] = self._randomized_source_quat_bank_t[source_indices] + target_pos_e[randomized_rows] = self._randomized_target_pos_bank_t[target_indices] + randomized_env_ids = env_ids[randomized_rows] + self._last_source_bank_index[randomized_env_ids] = source_indices + self._last_arm_bank_index[randomized_env_ids] = arm_indices + self._last_target_bank_index[randomized_env_ids] = target_indices + + zero_arm_velocity = torch.zeros_like(arm_q) + self._robot.write_joint_position_to_sim_index( + position=arm_q, + joint_ids=self._arm_joint_ids, + env_ids=env_ids, + ) + self._robot.write_joint_velocity_to_sim_index( + velocity=zero_arm_velocity, + joint_ids=self._arm_joint_ids, + env_ids=env_ids, + ) + self._robot.set_joint_position_target_index( + target=arm_q, + joint_ids=self._arm_joint_ids, + env_ids=env_ids, + ) + # The operator-only absolute joint controller uses the reset pose as its zero-command + # origin. The training controller is a standard relative action and has no such method. + arm_action = self.action_manager.get_term("arm_action") + if hasattr(arm_action, "set_action_offset"): + arm_action.set_action_offset(arm_q, env_ids=env_ids) + finger_position = self._curriculum_finger_pos_t[stage].unsqueeze(-1).expand(-1, len(FINGER_JOINTS)).clone() + finger_drive_target = finger_position.clone() + self._robot.write_joint_position_to_sim_index( + position=finger_position, + joint_ids=self._finger_joint_ids, + env_ids=env_ids, + ) + self._robot.write_joint_velocity_to_sim_index( + velocity=torch.zeros_like(finger_position), + joint_ids=self._finger_joint_ids, + env_ids=env_ids, + ) + self._robot.set_joint_position_target_index( + target=finger_drive_target, + joint_ids=self._finger_joint_ids, + env_ids=env_ids, + ) + gripper_target = torch.where( + (stage < self._grasp_stage_index).unsqueeze(-1), + torch.full((n, 1), float(self.cfg.gripper_preload_pos), device=self.device), + torch.full((n, 1), float(self.cfg.gripper_open_pos), device=self.device), + ) + self.action_manager.get_term("gripper_action").set_reset_position( + gripper_target, + env_ids=env_ids, + ) + + # Public root/joint writers invalidate FK. Reading a public body-pose view consumes + # the accumulated articulation mask, making all dirtied body poses authoritative + # before solver caches and source-cup proxy transforms are refreshed. Priming the + # robot view also prevents its next observation from issuing a redundant FK launch. + _ = self._robot.data.body_link_pose_w + + tcp_pose_e = self.tcp_pose_e()[env_ids] + lifted_stage = stage < self._grasp_stage_index + held_source_quat = math_utils.quat_mul( + tcp_pose_e[:, 3:7], + math_utils.quat_conjugate(self._desired_grasp_tcp_quat_c[env_ids]), + ) + source_quat = torch.where(lifted_stage.unsqueeze(-1), held_source_quat, source_quat) + grasp_offset = torch.zeros((n, 3), device=self.device) + grasp_offset[:, 2] = float(self.cfg.cup_grasp_height) + tcp_cup_pos_e = tcp_pose_e[:, :3] - math_utils.quat_apply(source_quat, grasp_offset) + # Lifted stages follow their solved TCP. Full and randomized stages use authored or + # IK-paired table positions. + cup_pos_e = torch.where(lifted_stage.unsqueeze(-1), tcp_cup_pos_e, cup_pos_e) + cup_world = cup_pos_e + self.env_origins[env_ids] + cup_pose = torch.cat((cup_world, source_quat), dim=-1) + self._source_cup.write_root_pose_to_sim_index(root_pose=cup_pose, env_ids=env_ids) + self._source_cup.write_root_velocity_to_sim_index( + root_velocity=cup_pose.new_zeros((n, 6)), + env_ids=env_ids, + ) + target_world = target_pos_e + self.env_origins[env_ids] + target_pose = torch.cat((target_world, target_quat), dim=-1) + self._target_cup.write_root_pose_to_sim_index(root_pose=target_pose, env_ids=env_ids) + self._target_cup.write_root_velocity_to_sim_index( + root_velocity=target_pose.new_zeros((n, 6)), + env_ids=env_ids, + ) + + new_p = self._sample_cup_media(cup_world, source_quat) + self._media.write_particle_pos_to_sim_index(new_p, env_ids=env_ids) + self._media.write_particle_velocity_to_sim_index(torch.zeros_like(new_p), env_ids=env_ids) + # Particle writers restore q/qd in both state buffers. Reset constitutive and collider + # history for the same worlds so a captured replay cannot retain the previous episode. + NewtonManager.reset_solver_state( + world_mask=None if self.num_envs == 1 else wp.from_torch(world_mask, dtype=wp.bool), + flags=newton.StateFlags.BODY | newton.StateFlags.PARTICLE, + ) + self._particle_region_cache = None + self._particle_region_cache_step = -1 + self.episode_succeeded[env_ids] = False + self.ep_max_target_frac[env_ids] = 0.0 + self._success_dwell_count[env_ids] = 0 + self._lost_grasp_dwell_count[env_ids] = 0 + self._lifted_grasp_seen[env_ids] = False + self._target_entry_seen[env_ids] = False + self._held_delivered[env_ids] = False + # Selective resets can occur after this step's termination/reward pass. Invalidating the + # scalar cache is cheap and keeps direct reset/replay workflows correct as well. + self._held_delivery_tracker_step = -1 + + def _sample_cup_media(self, cup_pos: torch.Tensor, cup_quat: torch.Tensor) -> torch.Tensor: + """Transform the local media lattice into selected cup poses on the simulation device.""" + particle_count = self._media_local_points_t.shape[0] + local_points = self._media_local_points_t.unsqueeze(0).expand(cup_pos.shape[0], -1, -1) + quaternions = cup_quat.unsqueeze(1).expand(-1, particle_count, -1) + return math_utils.quat_apply(quaternions, local_points) + cup_pos.unsqueeze(1) + + +class FrankaPourResetDatasetValidationEnv(FrankaPourEnv): + """Offline replay environment that accepts schema-valid candidate datasets.""" + + def _validate_loaded_reset_dataset(self, payload: dict) -> None: + """Validate candidate structure without requiring output provenance yet.""" + validate_reset_dataset( + payload, + expected_task_contract=build_franka_pour_reset_task_contract(self), + ) + + +class FrankaPourResetSamplerEnv(FrankaPourEnv): + """One-world offline scene that skips procedural reset banks and RL managers.""" + + def load_managers(self) -> None: + """Resolve only scene data consumed by the reset-dataset generator.""" + dev = self.device + self._robot = self.scene["robot"] + self._source_cup = self.scene["source_cup"] + self._target_cup = self.scene["target_cup"] + self._media: MPMObject = self.scene["media"] + self._arm_joint_ids, _ = self._robot.find_joints(ARM_JOINTS, preserve_order=True) + self._finger_joint_ids, _ = self._robot.find_joints(FINGER_JOINTS, preserve_order=True) + self._joint_pos_limits_t = self._robot.data.joint_pos_limits.torch.clone() + self.env_origins = self.scene.env_origins.to(device=dev, dtype=torch.float32) + self._num_particles = int(self._media.particles_per_object) + self._media_local_points_t = torch.as_tensor(self._media_local_points, device=dev, dtype=torch.float32) + + # ManagerBasedRLEnv.close() deletes these attributes unconditionally. No manager is + # constructed because this scene is never reset or stepped through the Gym API. + self.command_manager = None + self.reward_manager = None + self.termination_manager = None + self.curriculum_manager = None + self.recorder_manager = None + self.action_manager = None + self.observation_manager = None diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/pour_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/pour_env_cfg.py new file mode 100644 index 000000000000..ed571bcaa87c --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/pour_env_cfg.py @@ -0,0 +1,2139 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Franka grasp-a-cup-of-MPM-media-and-pour, on the stable Isaac-Lift-Cube-Franka foundation. + +Scene assets are borrowed from the lift task (standard Franka + the SeattleLab table USD) for a +stable, familiar base. On top we add a coupled Newton solver with **proxy coupling**: + +* an MJWarp ``arm`` entry owns the robot, the dynamic source cup, and the fixed receiver, and +* an implicit ``media`` entry owns the MPM particles. + +The source cup is a real dynamic rigid body resting on the table: the Franka grasps it with its fingers +through Newton-generated friction contacts resolved by MJWarp, and a Newton proxy mapping exposes +both cups' ``COLLIDE_PARTICLES`` cavity meshes to the MPM solver as auto-pose-synced colliders. +This replaces the earlier welded-kinematic-cup design. + +The source cup carries two co-located shapes on the same body: a solid grasp box (``COLLIDE_SHAPES``, +arm-entry-only) the fingers can actually grip, and a hollow cavity mesh (``COLLIDE_PARTICLES``) the +proxy bridges to MPM. Both learning variants use relative joint commands; the reset-dataset variant +uses a binary symmetric gripper. +""" + +from __future__ import annotations + +import math +from copy import deepcopy + +from isaaclab_newton.assets import MPMObjectCfg +from isaaclab_newton.physics import ( + MJWarpSolverCfg, + MPMSolverCfg, + NewtonCfg, + NewtonCollisionPipelineCfg, +) +from isaaclab_newton.sim.schemas import MujocoJointCfg +from isaaclab_newton.sim.spawners.mpm import MPMParticleMaterialCfg + +import isaaclab.sim as sim_utils +from isaaclab.assets import AssetBaseCfg, RigidObjectCfg +from isaaclab.envs import ManagerBasedRLEnvCfg +from isaaclab.managers import CurriculumTermCfg as CurrTerm +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 RewardTermCfg as RewTerm +from isaaclab.managers import SceneEntityCfg +from isaaclab.managers import TerminationTermCfg as DoneTerm +from isaaclab.scene import InteractiveSceneCfg +from isaaclab.sim.schemas import MassCfg, UsdPhysicsCollisionCfg, UsdPhysicsRigidBodyCfg +from isaaclab.sim.spawners.from_files.from_files_cfg import GroundPlaneCfg, UsdFileCfg +from isaaclab.sim.spawners.materials import RigidBodyMaterialBaseCfg +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR +from isaaclab.utils.configclass import configclass +from isaaclab.visualizers import VisualizerCfg + +from isaaclab_contrib.coupling import CouplerEntryCfg, CouplerProxyCfg, CouplerProxyMappingCfg + +from isaaclab_tasks.utils.adaptive_reset_sampler import AdaptiveResetSamplerCfg + +from isaaclab_assets.robots.franka import FRANKA_PANDA_CFG + +from . import mdp +from .cube_bowl_mesh import cube_bowl_inner_bounds +from .cube_bowl_spawner_cfg import CubeBowlSpawnerCfg +from .cup_media import build_media_object_cfg, cup_cavity_lattice + +RIGID_ENTRY = "arm" +MPM_ENTRY = "media" +FRANKA_POUR_ROBOT_USD_PATH = "omniverse://isaac-dev.ov.nvidia.com/Isaac/IsaacLab/Robots/FrankaEmika/franka_panda.usda" +FRANKA_POUR_ARM_COLLISION_PROXIES = frozenset( + { + "link0_c", + "link1_c", + "link2_c", + "link3_c", + "link4_c", + "link5_c0", + "link5_c1", + "link5_c2", + "link6_c", + "link7_c", + } +) +SPILL_FLOOR_LABEL_PATTERN = r".*/SpillFloor$" +GRASP_APPROACH_STAGE_NAMES = ( + "approach_1", + "approach_2", + "approach_3", + "approach_4", + "approach_5", + "approach_6", +) +CURRICULUM_STAGE_NAMES = ( + "drain", + "deep_tilt", + "tilt", + "pour", + "near_carry", + "mid_carry", + "carry", + "grasp", + *GRASP_APPROACH_STAGE_NAMES, + "full", + "randomized", +) + + +def spawn_franka_with_arm_collisions( + prim_path: str, + cfg: UsdFileCfg, + translation: tuple[float, float, float] | None = None, + orientation: tuple[float, float, float, float] | None = None, + **kwargs, +): + """Spawn the canonical Franka and activate only its dedicated arm collision proxies.""" + from pxr import Usd, UsdPhysics # noqa: PLC0415 + + robot_prim = sim_utils.spawn_from_usd(prim_path, cfg, translation, orientation, **kwargs) + for root_prim in sim_utils.find_matching_prims(prim_path, stage=robot_prim.GetStage()): + self_collision = root_prim.GetAttribute("newton:selfCollisionEnabled") + if not self_collision or not self_collision.Set(True): + raise RuntimeError(f"Franka asset at {root_prim.GetPath()} has no writable Newton self-collision flag.") + proxy_roots = { + prim.GetName(): prim.GetParent().GetPath() + for prim in Usd.PrimRange(root_prim, Usd.TraverseInstanceProxies()) + if prim.GetName() in FRANKA_POUR_ARM_COLLISION_PROXIES and prim.GetParent().GetName() == prim.GetName() + } + if proxy_roots.keys() != FRANKA_POUR_ARM_COLLISION_PROXIES: + missing = sorted(FRANKA_POUR_ARM_COLLISION_PROXIES.difference(proxy_roots)) + raise RuntimeError(f"Franka asset at {root_prim.GetPath()} is missing collision proxies: {missing}.") + for proxy_root in proxy_roots.values(): + root_prim.GetStage().OverridePrim(proxy_root).SetInstanceable(False) + for proxy_root in proxy_roots.values(): + collision_prim = root_prim.GetStage().GetPrimAtPath(proxy_root.AppendChild(proxy_root.name)) + UsdPhysics.CollisionAPI(collision_prim).GetCollisionEnabledAttr().Set(True) + return robot_prim + + +PANDA_ARM_JOINT_LIMITS = ( + (-2.8973, 2.8973), + (-1.7628, 1.7628), + (-2.8973, 2.8973), + (-3.0718, -0.0698), + (-2.8973, 2.8973), + (-0.0175, 3.7525), + (-2.8973, 2.8973), +) + + +def _mpm_solver_cfg(cfg: FrankaPourEnvCfg) -> MPMSolverCfg: + """Return the task's unique implicit-MPM solver config.""" + entries = [entry for entry in cfg.sim.physics.solver_cfg.entries if entry.name == MPM_ENTRY] + if len(entries) != 1: + raise ValueError(f"Expected exactly one {MPM_ENTRY!r} solver entry, found {len(entries)}.") + return entries[0].solver_cfg + + +def _resolve_mpm_cell_cap(cfg: FrankaPourEnvCfg) -> int: + """Resolve the total MPM active-cell capacity without mutating ``cfg``. + + Sparse training reserves an aligned hard upper bound per independent world so Newton can + capture topology rebuilds. Fixed and dense grids retain their configured capacity unless an + explicit total override is provided. + + Returns: + The total capacity to assign to the MPM solver entry. + """ + solver_cfg = _mpm_solver_cfg(cfg) + override = cfg.mpm_cell_cap_override + if override is not None: + capacity = int(override) + elif solver_cfg.grid_type == "sparse": + alignment = int(cfg.mpm_cell_capacity_alignment) + if alignment <= 0: + raise ValueError(f"Franka Pour MPM cell-capacity alignment must be positive, got {alignment}.") + particle_count = int(cup_cavity_lattice(cfg)[0].shape[0]) + per_world = ((particle_count + alignment - 1) // alignment) * alignment + capacity = per_world * int(cfg.scene.num_envs) + else: + capacity = int(solver_cfg.max_active_cell_count) + + if capacity <= 0: + raise ValueError(f"Franka Pour MPM capacity must be positive, got {capacity}.") + return capacity + + +@configclass +class PourSceneCfg(InteractiveSceneCfg): + """Lift-task scene assets plus resolved cups and MPM media.""" + + # SeattleLab table (top at env z=0), exactly as the Isaac-Lift-Cube-Franka scene. + table = AssetBaseCfg( + prim_path="{ENV_REGEX_NS}/Table", + init_state=AssetBaseCfg.InitialStateCfg(pos=[0.5, 0, 0], rot=[0, 0, 0.707, 0.707]), # xyzw, matches Lift + spawn=UsdFileCfg(usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Mounts/SeattleLabTable/table_instanceable.usd"), + ) + plane = AssetBaseCfg( + prim_path="/World/GroundPlane", + init_state=AssetBaseCfg.InitialStateCfg(pos=[0, 0, -1.05]), + spawn=GroundPlaneCfg(), + ) + light = AssetBaseCfg( + prim_path="/World/light", spawn=sim_utils.DomeLightCfg(color=(0.75, 0.75, 0.75), intensity=3000.0) + ) + robot = FRANKA_PANDA_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + robot.spawn.usd_path = FRANKA_POUR_ROBOT_USD_PATH + robot.spawn.func = spawn_franka_with_arm_collisions + # The task-specific spawner overrides the source USD's Newton-native value; retain the matching + # backend-independent intent in the articulation configuration as well. + robot.spawn.articulation_props.enabled_self_collisions = True + # Resolve the implicit-actuator parameters from the requested USD rather than overwriting them + # with Isaac Lab's generic Franka gains and limits. + robot.actuators = { + name: actuator_cfg.replace( + effort_limit_sim=None, + velocity_limit_sim=None, + stiffness=None, + damping=None, + armature=None, + ) + for name, actuator_cfg in robot.actuators.items() + } + robot.spawn.joint_drive_props = [MujocoJointCfg(actuatorgravcomp=True)] + # Built by :meth:`FrankaPourEnvCfg.finalize` from the final override values. + source_cup: RigidObjectCfg | None = None + target_cup: RigidObjectCfg | None = None + media: MPMObjectCfg | None = None + + +@configclass +class ActionsCfg: + """Relative arm-joint increments and one continuous symmetric-gripper command.""" + + arm_action = mdp.RelativeJointPositionActionCfg( + asset_name="robot", + joint_names=[f"panda_joint{i}" for i in range(1, 8)], + preserve_order=True, + # Eight-hundredths of a radian per policy step gives useful reach authority without + # bypassing the articulation position drives or encoding a demonstrated trajectory. + scale=0.08, + use_zero_offset=True, + ) + gripper_action = mdp.CurriculumGripperPositionActionCfg( + asset_name="robot", + joint_names=["panda_finger.*"], + # Zero action holds the contact-safe preload. Negative actions close farther and positive + # actions continuously open the fingers, so the policy—not a phase interlock—owns grasping. + scale=0.016, + alpha=0.2, + close_position=0.021, + neutral_position=0.04, + open_position=0.04, + default_position=0.024, + limit_to_preload=False, + force_open_before_phase_stage=-1, + # The lift gate still requires persistent bilateral deflection and actual cup motion. A + # 5 cm/s finger-settling threshold avoids spending most of a five-second attempt waiting + # for sub-millimetre drive oscillations to decay. + contact_max_velocity=0.05, + ) + + +@configclass +class ObservationsCfg: + @configclass + class PolicyCfg(ObsGroup): + """Sensor-compatible robot, gripper, and cup geometry available to the actor.""" + + arm_q = ObsTerm( + func=mdp.joint_pos_rel, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=["panda_joint.*"])}, + scale=0.3, + ) + arm_qd = ObsTerm( + func=mdp.joint_vel_rel, + params={"asset_cfg": SceneEntityCfg("robot", joint_names=["panda_joint.*"])}, + scale=0.05, + ) + time_remaining = ObsTerm(func=mdp.time_remaining_obs) + pour_target_fraction = ObsTerm(func=mdp.pour_target_fraction_obs) + tcp_pose = ObsTerm(func=mdp.tcp_pose_obs) + cup_pose = ObsTerm(func=mdp.cup_pose_obs) + target_pose = ObsTerm(func=mdp.target_pose_obs) + tcp_to_grasp_position_c = ObsTerm(func=mdp.tcp_to_grasp_position_c_obs, scale=10.0) + grasp_to_tcp_quat = ObsTerm(func=mdp.grasp_to_tcp_quat_obs) + target_position_c = ObsTerm(func=mdp.target_position_c_obs, scale=5.0) + finger_position = ObsTerm(func=mdp.finger_position_obs, scale=25.0) + finger_velocity = ObsTerm(func=mdp.finger_velocity_obs, scale=5.0) + gripper_target = ObsTerm(func=mdp.gripper_target_obs, scale=25.0) + gripper_contact = ObsTerm(func=mdp.gripper_contact_obs, scale=250.0) + last_action = ObsTerm(func=mdp.last_action, scale=0.2) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = True + + @configclass + class PrivilegedCfg(ObsGroup): + """Exact simulation state available only to the asymmetric critic.""" + + success_dwell = ObsTerm(func=mdp.success_dwell_obs) + lost_grasp_dwell = ObsTerm(func=mdp.lost_grasp_dwell_obs) + cup_velocity = ObsTerm(func=mdp.cup_velocity_obs, scale=0.1) + particle_fractions = ObsTerm(func=mdp.particle_fractions_obs) + particle_transfer = ObsTerm(func=mdp.particle_transfer_obs) + held_delivery_history = ObsTerm(func=mdp.held_delivery_history_obs) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = True + + policy: PolicyCfg = PolicyCfg() + privileged: PrivilegedCfg = PrivilegedCfg() + + +@configclass +class RewardsCfg: + # One discounted hierarchical physical potential spans every backward-curriculum reset. It is + # policy-invariant at PPO's gamma, so holding or wiggling cannot improve discounted return. + task_progress = RewTerm( + func=mdp.PourTaskProgress, + weight=5.0, + params={ + "target_height": 0.12, + "reach_std": 0.07, + "grasp_reach_std": 0.015, + "grasp_preload_position": 0.024, + "lift_height": 0.06, + "align_std": 0.12, + "source_offset_xy": (0.0, 0.05), + "target_tilt": math.radians(140.0), + "pour_direction_xy": (0.0, -1.0), + "source_mouth_height": 0.099, + "alignment_radius": 0.15, + # Tilt is an exploration bootstrap only for the supplied-grasp stages. Full-task + # policies are optimized by actual held particle transfer rather than a prescribed pose. + "active_through_stage": 6, + "min_lift_height": 0.05, + "max_tcp_distance": 0.018, + "max_gripper_width_error": 0.006, + "max_gripper_command": 0.024, + # Must match PPO gamma for policy-invariant discounted potential shaping. + "discount_factor": 0.99, + }, + ) + # The full-task reset starts outside the narrow reach kernel above. This broad physical-pose + # potential preserves a Cartesian gradient after premature closure and rewards the cup-relative + # side-grasp orientation without prescribing a trajectory or exposing a phase variable. + approach_progress = RewTerm( + func=mdp.ApproachProgress, + weight=8.0, + params={ + "position_std": 0.20, + "orientation_std": 0.75, + "open_hand_fraction": 0.35, + "active_from_stage": 8, + "discount_factor": 0.99, + }, + ) + # Once approach reaches the contact neighborhood, distinguish a loaded side grasp from empty + # closure and retain a monotonic signal until the glass has cleared the table. + grasp_lift_progress = RewTerm( + func=mdp.GraspLiftProgress, + weight=10.0, + params={ + "target_height": 0.10, + "grasp_reach_std": 0.025, + "grasp_preload_position": 0.024, + "grasp_fraction": 0.40, + "active_from_stage": 4, + "discount_factor": 0.99, + }, + ) + # Signed held-delivery progress is capped at the active success threshold. Particles leaving + # the receiver repay their credit, and an unsuccessful episode repays any credit still held. + delivered = RewTerm( + func=mdp.HeldDeliveryProgress, + weight=30.0, + params={ + "min_lift_height": 0.05, + "max_tcp_distance": 0.018, + "max_gripper_width_error": 0.006, + "max_gripper_command": 0.024, + }, + ) + success = RewTerm(func=mdp.pour_success_bonus, weight=25.0) + # Airborne transfer is excluded; each particle is penalized once after reaching the table + # outside both cups. Termination bounds the failure at just over ten percent. + spill = RewTerm(func=mdp.NewlySpilledParticles, weight=-30.0) + # Count overlapping failures and an unsuccessful deadline once. This keeps a transient dump + # that misses the stable-success predicate strictly worse than completing the task. + failure = RewTerm(func=mdp.terminal_failure, weight=-35.0) + action_rate = RewTerm(func=mdp.action_rate_l2, weight=-1.0e-4) + # Penalize unnecessarily large relative arm increments and finger commands without prescribing + # a reference trajectory. + action_magnitude = RewTerm(func=mdp.action_l2, weight=-0.05) + + +@configclass +class ResetDatasetRewardsCfg: + """Stage-independent OmniReset-style rewards for reset-dataset training.""" + + # Task rewards: generic reach and goal-set distance plus the strict particle success state. + # The broad kernel remains informative across full-workspace reaching while the reset dataset + # supplies close-contact precision without adding a task-specific grasp trajectory. + reach = RewTerm(func=mdp.tcp_cup_distance_tanh, weight=0.1, params={"std": 0.3}) + # Particle distances span only a few decimetres; this preserves useful transfer contrast + # while retaining the same task-independent tanh form used by OmniReset. + goal_distance = RewTerm(func=mdp.media_target_distance_tanh, weight=0.1, params={"std": 0.2}) + # A target occupancy of 30% is an immediate terminal success. Dividing the terminal pulse by + # the policy step keeps its integrated contribution equal to this unit weight. + success = RewTerm(func=mdp.pour_success_bonus, weight=1.0) + + # Smoothness is one semantic reward group, kept as separate standard terms for diagnostics. + action_magnitude = RewTerm(func=mdp.action_l2, weight=-1.0e-4) + action_rate = RewTerm(func=mdp.action_rate_l2, weight=-1.0e-3) + joint_velocity = RewTerm( + func=mdp.finite_joint_velocity_l2, + weight=-1.0e-2, + params={ + "asset_cfg": SceneEntityCfg("robot", joint_names=["panda_joint.*"]), + "max_velocity": 20.0, + }, + ) + + # Ordinary fixed-horizon completion is neutral. ``terminal_failure`` is already normalized by + # the policy step, so this weight produces one exact -1 abnormal-state pulse per episode. + failure = RewTerm(func=mdp.terminal_failure, weight=-1.0, params={"include_time_out": False}) + + +@configclass +class TerminationsCfg: + failure = DoneTerm(func=mdp.nonfinite_failure) + extreme_rigid_state = DoneTerm(func=mdp.extreme_rigid_state) + lost_grasp = DoneTerm( + func=mdp.lost_lifted_grasp, + params={ + "dwell_time_s": 0.05, + "max_tcp_distance": 0.018, + "max_gripper_width_error": 0.006, + "max_gripper_command": 0.024, + }, + ) + spill = DoneTerm(func=mdp.excessive_spill) + particle_out_of_bounds = DoneTerm(func=mdp.particle_out_of_bounds) + # Success follows every failure predicate, then the custom timeout excludes same-step success. + success = DoneTerm( + func=mdp.stable_pour_success, + params={ + "dwell_time_s": 0.15, + "min_lift_height": 0.05, + "max_tcp_distance": 0.018, + "max_gripper_width_error": 0.006, + "max_gripper_command": 0.024, + }, + ) + time_out = DoneTerm(func=mdp.unsuccessful_time_out, time_out=True) + + +@configclass +class EventsCfg: + reset_scene = EventTerm(func=mdp.reset_pour_scene, mode="reset") + + +@configclass +class CurriculumCfg: + stage = CurrTerm(func=mdp.PourCurriculum) + + +@configclass +class ResetDatasetCurriculumCfg: + """Adaptive curriculum over a validated reset-state dataset.""" + + reset_dataset = CurrTerm(func=mdp.PourResetDatasetCurriculum) + + +@configclass +class FrankaPourEnvCfg(ManagerBasedRLEnvCfg): + """Franka grasping a dynamic cup of MPM media on the lift foundation, proxy-coupled solver.""" + + scene: PourSceneCfg = PourSceneCfg(num_envs=2, env_spacing=2.5, replicate_physics=True) + observations: ObservationsCfg = ObservationsCfg() + actions: ActionsCfg = ActionsCfg() + rewards: RewardsCfg = RewardsCfg() + terminations: TerminationsCfg = TerminationsCfg() + events: EventsCfg = EventsCfg() + curriculum: CurriculumCfg = CurriculumCfg() + + # ---- Franka layout / reset (horizontal gripper behind the glass, ready to grasp) ---- + # The fingers approach parallel to the table instead of descending over the rim. This + # configuration seeds the far-side Newton-IK reset bank; the policy must still acquire physical + # contact, lift, carry, and pour through direct actions. + arm_home: tuple[float, float, float, float, float, float, float] = ( + -1.07505691, + 0.76868522, + 0.53213346, + -2.93226814, + 2.50670838, + 1.40047050, + 0.21146376, + ) + # Task-space metadata is independent of the policy action representation. SpaceMouse teleop + # uses the same frame for its input-only IK adapter, while PPO commands joint positions. + tcp_body_name: str = "panda_hand" + tcp_offset_pos: tuple[float, float, float] = (0.0, 0.0, 0.107) + tcp_offset_rot: tuple[float, float, float, float] = (0.0, 0.0, 0.0, 1.0) + gripper_open_pos: float = 0.04 # finger position the cup is grasped from (fingers start open) + # A target inside the 0.028 m geometric contact position proves active squeeze/preload instead + # of open fingers passively compressed by cup contact. + # The zero-action target retained the tall, media-loaded glass in validation without the + # instability caused by higher friction. The continuous command may close 3 mm farther or open + # to the full 40 mm finger position. + # The close target is derived from this value and ``gripper_close_offset`` during finalization. + gripper_preload_pos: float = 0.024 + gripper_close_offset: float = 0.003 + + # ---- source cup (dynamic, grasped by the fingers) ---- + # A 56 x 56 x 119 mm hollow glass leaves 24 mm of clearance in the Panda's 80 mm opening. Its + # visible outer wall and solid grasp proxy have exactly the same extents. The taller profile + # leaves 36 mm of wall above the horizontal grasp TCP, so the fingers engage the side without + # pinching or entering the rim while the hand remains clear of the table. + source_cup_inner_width: float = 0.042 + source_cup_inner_depth: float = 0.042 + source_cup_cavity_depth: float = 0.110 + source_cup_wall_thickness: float = 0.007 + source_cup_bottom_thickness: float = 0.009 + source_cup_friction: float = 0.9 + cup_mass: float = 0.05 + # Match the visible glass's exact 56 x 56 x 119 mm outer envelope. Direct contact instrumentation + # showed that a 60 mm TCP made the horizontal hand collide with the table in randomized poses; + # 83 mm preserves a side grasp while providing physical clearance for the complete approach. + cup_grasp_box_half: tuple[float, float, float] = (0.028, 0.028, 0.0595) + cup_grasp_height: float = 0.083 + # Cup-local TCP orientation for a true side grasp. A +90 degree rotation about cup +Y maps + # Panda tool +Z to cup +X, leaving both the finger length and jaw axis parallel to the table. + cup_grasp_tcp_quat_c: tuple[float, float, float, float] = ( + 0.0, + math.sqrt(0.5), + 0.0, + math.sqrt(0.5), + ) + # A/B grasp testing: mu=1 let the media-loaded cup roll out during the carry, while mu=2 kept + # bilateral contact through the full carry/tilt path without raising tangential stiffness. + cup_grasp_box_friction: float = 2.0 + # The Newton default (ke=2.5e3 N/m) permits several millimeters of penetration under the + # finger position drive. That is enough for the fingers to cross the narrow grasp proxy and + # lose the contact manifold entirely. Match Newton's rigid-robot contact recipe instead. + # A/B sweep: 50 kN/m retained the cup but allowed ~4 mm corner penetration at mu=2; + # 100 kN/m reduced typical penetration to ~2 mm and stayed stable; 200 kN/m ejected the cup. + grasp_contact_ke: float = 1.0e5 + grasp_contact_kd: float = 5.0e2 + grasp_contact_kf: float = 1.0e3 + # Cup reset pose, in the env frame. The cup rests on the table (top z=0) directly under the home + # gripper, opening up. z is the cup base height (the body local origin sits at the outer base). + cup_reset_pos: tuple[float, float, float] = (0.5, 0.0, 0.0) + + # ---- receiving cup (fixed, represented once and proxied between solvers) ---- + # The receiver is wider than the source during the initial learning curriculum. It remains a + # proper hollow cup while making early particle-delivery experience substantially less sparse. + target_cup_inner_width: float = 0.140 + target_cup_inner_depth: float = 0.140 + target_cup_cavity_depth: float = 0.065 + target_cup_wall_thickness: float = 0.009 + target_cup_bottom_thickness: float = 0.009 + target_cup_friction: float = 0.8 + # Keep enough initial clearance for the Panda finger collision meshes. Moving this wide receiver + # to y=-0.12 made its rigid rim touch the pre-grasp hand and explosively eject media at step 0. + target_cup_reset_pos: tuple[float, float, float] = (0.5, -0.18, 0.0) + # The grasped source origin moves about 5 cm toward environment -y during the deep +x tilt. + # Starting behind the receiver keeps the draining mouth centered throughout that motion. + pour_source_offset_xy: tuple[float, float] = (0.0, 0.05) + collider_margin: float = 0.002 + + # Full-task threshold kept as a standalone compatibility knob. Earlier curriculum stages use + # the values below; :attr:`curriculum_target_frac` combines both sources. + # A 30% transfer is 74 of the 245 particles. The validated side-pour motion reaches about 41%, + # leaving enough margin that the success predicate measures manipulation rather than the + # lower tail of large-batch MPM/contact variation. + pour_target_frac: float = 0.30 + particle_count_margin: float = 0.003 + # Particle point samples resting on the z=0 MPM spill plane settle within the containment + # margin above it. Only points in that contact band and outside both cups are true spills. + spill_table_height: float = 0.0 + max_spill_fraction: float = 0.10 + # A transfer must remain above its stage threshold for this duration before successful + # termination. This rejects transient particle crossings and aligns reward with curriculum. + # Nine consecutive control steps reject transient particle crossings while leaving enough of + # the finite horizon for the terminal event after a late, valid randomized grasp. + success_dwell_time_s: float = 0.15 + lost_grasp_dwell_time_s: float = 0.05 + """Continuous post-lift grasp loss required before failure [s].""" + success_min_lift_height: float = 0.05 + success_max_tcp_distance: float = 0.018 + # A contact-free hand reaches a 64 mm measured gap at the bounded 24 mm command, exactly 8 mm + # wider than this 56 mm cup. Requiring <=6 mm distinguishes real bilateral cup contact while + # retaining roughly 3 mm of measured true-grasp variation. + success_max_gripper_width_error: float = 0.006 + # ``None`` derives the largest command that still guarantees the configured drive deflection + # at geometric cup contact. This keeps continuous near-preload exploration eligible while the + # physical bilateral-contact predicate rejects a genuinely opening or empty hand. + success_max_gripper_command: float | None = None + + def _resolved_success_max_gripper_command(self) -> float: + if self.success_max_gripper_command is not None: + return float(self.success_max_gripper_command) + contact_limit = float(self.cup_grasp_box_half[1]) - float(self.actions.gripper_action.contact_min_deflection) + return max(float(self.gripper_preload_pos), contact_limit) + + # Reset extreme but finite rigid state before it can enter actor observation normalization. + state_bound_joint_position_margin: float = 0.05 + state_bound_max_joint_velocity: float = 20.0 + state_bound_max_cup_linear_velocity: float = 10.0 + state_bound_max_cup_angular_velocity: float = 50.0 + # Keep finite escaped particles from expanding the sparse NanoVDB hierarchy throughout + # an episode. Bounds are in each environment's local frame and comfortably contain both cups, + # every curriculum reset, and the robot workspace. + # The reset-dataset generator covers the central 90% of the Franka's full 360-degree reachable + # workspace. Keep the sparse-grid safety envelope symmetric behind the base so valid rear + # grasps and their cup-contained media are not mistaken for numerical escapes. + particle_workspace_lower_bound: tuple[float, float, float] = (-1.0, -1.0, -0.5) + particle_workspace_upper_bound: tuple[float, float, float] = (1.5, 1.0, 1.5) + + # ---- success-driven backward curriculum ---- + # The first reset starts with a grasped cup nearly drained over the receiver, then moves backward + # through a partial tilt, an upright pour, a source-side carry, an open-finger grasp, and + # progressively longer open-finger approaches. The final stage adds independently mixed arm, + # source, and receiver resets. This provides dense direct-control experience without an + # automatic trajectory. + # Reset IK is solved once into a bank; asynchronous resets only select prevalidated rows. + curriculum_stage_names: tuple[str, ...] = CURRICULUM_STAGE_NAMES + curriculum_pour_arm_q: tuple[float, float, float, float, float, float, float] = ( + -1.47599292, + 0.33629909, + 0.99845403, + -2.69460344, + 2.62228370, + 1.93315995, + 0.66680431, + ) + # Collision-screened joint-space points on the authored upright-to-deep-pour segment. They are + # reset states only: zero action holds the pose, and the policy must command all further motion. + curriculum_tilt_arm_q: tuple[float, float, float, float, float, float, float] = ( + -1.77459520, + 0.79981079, + 1.18383252, + -2.64059955, + 2.53800059, + 2.19650501, + 1.88206341, + ) + curriculum_deep_tilt_arm_q: tuple[float, float, float, float, float, float, float] = ( + -1.86647283, + 0.94242977, + 1.24087206, + -2.62398297, + 2.51206732, + 2.27753425, + 2.25598929, + ) + curriculum_drain_arm_q: tuple[float, float, float, float, float, float, float] = ( + -1.91241164, + 1.01373926, + 1.26939183, + -2.61567468, + 2.49910069, + 2.31804888, + 2.44295223, + ) + # Roll the side-grasped glass 140 degrees about the horizontal approach axis. The 119 mm glass + # retains its granular media at 120 degrees; this deeper but still natural wrist roll drains it + # without the instability observed at 150 degrees. + curriculum_pour_target_arm_q: tuple[float, float, float, float, float, float, float] = ( + -1.93538105, + 1.04939401, + 1.28365171, + -2.61152053, + 2.49261737, + 2.33830619, + 2.53643370, + ) + curriculum_carry_arm_q: tuple[float, float, float, float, float, float, float] = ( + -1.16845703, + 0.55803788, + 0.95656616, + -2.75139022, + 2.87593412, + 1.73866940, + 0.36629686, + ) + # Move backward from the receiver-side pour pose toward the source-side carry pose in two + # collision-screened joint-space increments. Each reset still derives the held cup pose from + # forward kinematics, so the arm, cup, media, and grasp remain exactly co-located. + curriculum_transport_reset_fractions: tuple[float, float] = (1.0 / 3.0, 2.0 / 3.0) + # The deterministic validation motion reaches ~41% on the full task. First-time particle + # delivery remains rewarded above every success threshold. + curriculum_early_target_frac: tuple[float, ...] = ( + 0.05, + 0.08, + 0.10, + 0.15, + 0.15, + 0.18, + 0.20, + 0.30, + 0.30, + 0.30, + 0.30, + 0.30, + 0.30, + 0.30, + ) + curriculum_randomized_pour_target_frac: float = 0.30 + # Conservative Cartesian half-extents containing the polar workspace below. This field stays + # available for downstream rectangular-reset configurations and particle-workspace validation; + # the Franka Pour preset uses the explicit polar domain when ``source_radius_range`` is set. + curriculum_randomized_source_position_range: tuple[float, float] = (0.30, 0.45) + # Cover the useful front-table reach sector instead of the former thin diagonal strip. At full + # extent the source samples a candidate side-grasp sector from 40--78 cm radius and +/-35 + # degrees about the fixed Franka base. A startup sweep retained complete collision-free paths + # across every radial ring and both angular boundaries with margin beyond the configured IK + # thresholds. Lower curriculum levels interpolate these cells from the authored pose. + curriculum_randomized_source_radius_range: tuple[float, float] | None = (0.40, 0.78) + curriculum_randomized_source_azimuth_range: float = math.radians(35.0) + # Retained for rectangular downstream presets. Zero means a full two-dimensional rectangle; + # it is ignored by the polar Franka Pour preset. + curriculum_randomized_source_xy_correlation: float = 0.0 + # After grasping anywhere in the polar source sector, pull the glass into the central +/-3 cm + # carry corridor while lifting. This preserves the full reach problem without forcing an + # edge-of-workspace upright carry against a joint limit. + curriculum_randomized_carry_position_range: tuple[float, float] = (0.03, 0.03) + # In the polar preset the cup's grasp face points approximately toward the robot. This optional + # value adds a local yaw perturbation around that radial-facing direction so the policy cannot + # rely on the source grasp face pointing toward the robot at the high-randomization frontier. + curriculum_randomized_source_yaw_range: float = math.radians(30.0) + # Keep the authored -Y pour direction while moving the receiver throughout the reachable table + # region. The reset mixer enforces rectangular separation after progressively breaking the + # source/receiver bank pairing. + curriculum_randomized_target_center_xy: tuple[float, float] = (0.50, -0.18) + curriculum_randomized_target_position_range: tuple[float, float] = (0.15, 0.44) + curriculum_randomized_cup_clearance: float = 0.04 + # The randomized stage starts 12 cm behind the glass along cup -X. The runtime rotates this + # cup-local offset and jitter by the source yaw, preserving the horizontal approach geometry. + curriculum_randomized_reset_tcp_standoff: tuple[float, float, float] = (-0.12, 0.0, 0.0) + curriculum_randomized_reset_tcp_jitter: tuple[float, float, float] = (0.02, 0.03, 0.0) + # Optional asymmetric offset box applied on top of the centered pre-grasp standoff. At full + # extent it varies approach depth from 10--28 cm, lateral displacement by +/-16 cm, and height + # by 25 cm. Every pose is solved and collision-screened before entering the reset bank. ``None`` + # preserves the legacy symmetric-jitter design for downstream presets. + curriculum_randomized_reset_tcp_offset_lower: tuple[float, float, float] | None = (-0.16, -0.16, 0.0) + curriculum_randomized_reset_tcp_offset_upper: tuple[float, float, float] | None = (0.02, 0.16, 0.25) + # Reset-only orientation error relative to the cup-aligned grasp frame. Fibonacci-sphere axes + # avoid a preferred rotation direction, while a positive lower angle guarantees that the final + # curriculum never begins perfectly grasp-aligned, so the direct policy must recover from the + # observed pose error instead of receiving an aligned reach problem. + curriculum_randomized_reset_tcp_rotation_angle_range: tuple[float, float] = ( + math.radians(20.0), + math.radians(60.0), + ) + curriculum_randomized_reset_tcp_min_grasp_distance: float = 0.09 + # Move backward along samples already exercised by the complete-path collision sweep. Values + # are interpolation weights from the 6 cm midpoint toward the exact grasp, so decreasing + # weights increase reset distance before ``full`` introduces the aligned 12 cm pre-grasp. + curriculum_grasp_approach_fractions: tuple[float, ...] = (0.75, 0.50, 0.375, 0.25, 0.125, 0.0) + # Keep the held source proxy above the receiver rim throughout the independently solved + # pour-to-tilt joint interpolation. Without this reserve the source corner and two collider + # margins overlap even though both endpoint IK poses are valid. + curriculum_randomized_pour_clearance: float = 0.010 + # Center the closed-finger waypoint on the authored side-grasp point. A lower waypoint made + # the TCP settle about 8 mm below the grasp point and prevented otherwise valid captures. + curriculum_grasp_descent_overshoot: float = 0.0 + # An odd source grid includes the authored nominal pose and both configured XY extrema. Newton + # IK solves this bank once at startup; asynchronous resets only gather prevalidated rows. + curriculum_randomized_reset_ik_grid_size: int = 7 + curriculum_randomized_reset_ik_samples_per_source: int = 11 + # Require the feasible, posture-screened bank to retain broad workspace coverage. Runtime + # sampling is uniform over these source XY cells rather than over their unequal IK row counts. + # The polar grid is a reachability census, not a promise that every Cartesian cell is feasible + # for the complete grasp-to-deep-tilt trajectory. Retain at least one fifth of its cells while + # separately requiring every radius, both angular sides, and broad angular span. + curriculum_randomized_min_source_cell_fraction: float = 0.2 + # Require more than one safe arm start to be available in every retained source cell. The + # curriculum deliberately exposes only the closest of these paths at small extents, then grows + # arm-start diversity with the extent and exposes every surviving path at the final level. + # Screening the full reserve here prevents a broad Cartesian grid from silently collapsing to + # one repeated Franka posture when the high-randomization levels are reached. + curriculum_randomized_min_reset_variants_per_source: int = 2 + curriculum_randomized_reset_ik_iterations: int = 160 + curriculum_randomized_reset_ik_max_cost: float = 1.0e-3 + curriculum_randomized_reset_ik_joint_margin: float = 0.015 + # Backward-compatible upper safeguard for the folded panda_joint6 branch. The default lies just + # inside the URDF limit, so joint-limit and complete-path collision screening perform the + # effective filtering without imposing a narrow workspace-specific posture heuristic. + curriculum_randomized_reset_joint6_max: float = 3.75 + # Introduce the final stage through nested, normalized extents across source pose, receiver pose, + # reset-TCP offset, and the number of exposed arm-start variants. The exact zero-amplitude anchor + # is behaviorally identical to the mastered full task; small early increments prevent an + # IK/reset-bank discontinuity from masquerading as exploration difficulty. The last level + # contains the complete prevalidated randomization bank. + # Together with the fifteen preceding backward-reset stages, these nine frontiers form a + # twenty-four-level curriculum. Small initial extents prevent the first randomized reset from + # simultaneously changing reach, orientation, source placement, and receiver placement beyond + # the support of the mastered nominal policy. + curriculum_randomization_extent_levels: tuple[float, ...] = ( + 0.0, + 0.05, + 0.10, + 0.20, + 0.35, + 0.50, + 0.70, + 0.85, + 1.0, + ) + # Preserve paired collision-screened paths during the first geometry-only frontiers, then + # gradually break arm/source and receiver/source correlations. The final frontier samples all + # three independently, subject to conservative reset clearances. + curriculum_independent_arm_fraction_levels: tuple[float, ...] = ( + 0.0, + 0.0, + 0.0, + 0.0, + 0.10, + 0.25, + 0.50, + 0.75, + 1.0, + ) + curriculum_independent_target_fraction_levels: tuple[float, ...] = ( + 0.0, + 0.0, + 0.0, + 0.10, + 0.25, + 0.50, + 0.70, + 0.85, + 1.0, + ) + curriculum_independent_sample_attempts: int = 8 + curriculum_independent_arm_min_tcp_distance: float = 0.12 + # Stateful manager progress is not part of RSL-RL checkpoints. Set both start controls to the + # last logged values when resuming within the randomized stage. + curriculum_randomization_start_level: int = 0 + curriculum_success_threshold: float = 0.8 + # Standard Franka manipulation tasks train directly on broad reset distributions. Use a lower + # threshold only to expose the next nested randomization frontier; final task mastery still + # requires ``curriculum_success_threshold`` at full amplitude. + curriculum_randomization_promotion_threshold: float = 0.65 + # Use the most recent 4,096 frontier episodes for the success estimate. This is a statistical + # window, not sufficient policy-training exposure when thousands of environments reset at once. + curriculum_min_resets_per_stage: int = 4096 + # Require eight whole vectorized-environment reset cohorts before promotion. At the documented + # 512-environment scale this equals the success window above; at 3,000 environments it prevents + # a frontier from being promoted after only one or two PPO updates. The entry replay mixture + # decays over this same exposure horizon so the predecessor skill remains represented while the + # new prerequisite is consolidated. Set to zero only to retain the absolute-window behavior. + curriculum_min_reset_cohorts_per_stage: float = 8.0 + # Retain the immediately preceding nested task after promotion so the policy does not forget + # already-solved behavior while learning the newly introduced prerequisite. + curriculum_previous_stage_replay_fraction: float = 0.1 + # At a newly introduced frontier, begin with a balanced predecessor mixture and decay toward + # the retention fraction above over one evidence window. This avoids replacing nearly the + # entire rollout distribution at the exact moment a new prerequisite is introduced. + curriculum_frontier_entry_replay_fraction: float = 0.5 + # Set this to the last logged ``Curriculum/stage/stage`` value when resuming training. + curriculum_start_stage: int = 0 + curriculum_freeze: bool = False + + # ---- media (granular sand inside the cup) ---- + # Preserve the former particle volume in the taller 110 mm cavity. This keeps the + # particle count, material mass, sparse-grid capacity, and RL transfer threshold unchanged. + media_fill_frac: float = 0.17181818181818181 + # Normal rollouts peak near 2.2 m/s. This generous clamp prevents a numerically launched + # particle from crossing many NanoVDB upper regions within one manager step, before the + # workspace termination can selectively reset its environment. + particle_max_velocity: float = 10.0 + media_material: MPMParticleMaterialCfg = MPMParticleMaterialCfg( + density=1500.0, + friction=0.7, + yield_pressure=1.0e12, + ) + + # ---- MPM ---- + voxel_size: float = 0.01 + particles_per_cell: float = 2.0 + mpm_iterations: int = 24 + # A rebuildable multi-world grid stores guard/topology voxels in addition to each particle's + # occupied cell. Two aligned blocks per world prevent high world indices from exhausting the + # global captured reserve after tilted fills spread, while retaining a sparse local grid. + mpm_cell_capacity_alignment: int = 512 + mpm_cell_cap_override: int | None = None + # Advance the complete coupled system once per 120 Hz simulation tick. Per-entry refinements + # below avoid repeating the outer collision and proxy-coupling pipeline. + physics_substeps: int = 1 + # Four rigid refinements preserve the former number of arm solves per simulation tick without + # multiplying the coupled proxy exchange or MPM work. + rigid_entry_substeps: int = 4 + # Refine only the implicit-MPM entry twice inside each coupled tick. This improves thin-wall + # collision stability without duplicating the outer collision/coupling pipeline; rigid and MPM + # accuracy remain independently configurable through their per-entry substeps. + mpm_entry_substeps: int = 2 + proxy_iterations: int = 1 + # This scales only the virtual cup inertia inside the destination MPM solve. The rigid solver + # retains the authored cup mass and receives the harvested MPM reaction wrench. A stiff proxy + # prevents split-step cup yielding from letting resting media creep through its floor while + # retaining one inexpensive two-way coupling pass. + proxy_mass_scale: float = 100.0 + # Newton supports captured sparse rebuilds with one local MPM world per replicated environment. + use_cuda_graph: bool = True + + @property + def curriculum_target_frac(self) -> tuple[float, ...]: + """Per-stage delivered-particle success fractions.""" + return ( + *self.curriculum_early_target_frac, + float(self.pour_target_frac), + float(self.curriculum_randomized_pour_target_frac), + ) + + @curriculum_target_frac.setter + def curriculum_target_frac(self, values: tuple[float, ...]) -> None: + if len(values) != len(CURRICULUM_STAGE_NAMES): + raise ValueError(f"curriculum_target_frac must contain {len(CURRICULUM_STAGE_NAMES)} values.") + self.curriculum_early_target_frac = tuple(values[:14]) + self.pour_target_frac = float(values[14]) + self.curriculum_randomized_pour_target_frac = float(values[15]) + + def _curriculum_transport_arm_configs(self) -> tuple[tuple[float, ...], ...]: + """Return joint configurations between the receiver-side pour and source-side carry poses.""" + return tuple( + tuple( + (1.0 - fraction) * pour_q + fraction * carry_q + for pour_q, carry_q in zip( + self.curriculum_pour_arm_q, + self.curriculum_carry_arm_q, + strict=True, + ) + ) + for fraction in self.curriculum_transport_reset_fractions + ) + + def _configure_reward_cfg(self, max_gripper_command: float, *, initialize_stage_gates: bool) -> None: + """Propagate task controls into the reverse-curriculum reward terms.""" + progress_params = self.rewards.task_progress.params + progress_params["grasp_preload_position"] = self.gripper_preload_pos + progress_params["source_offset_xy"] = self.pour_source_offset_xy + progress_params["source_mouth_height"] = self.source_cup_bottom_thickness + self.source_cup_cavity_depth + progress_params["min_lift_height"] = self.success_min_lift_height + progress_params["max_tcp_distance"] = self.success_max_tcp_distance + progress_params["max_gripper_width_error"] = self.success_max_gripper_width_error + progress_params["max_gripper_command"] = max_gripper_command + grasp_lift_params = self.rewards.grasp_lift_progress.params + grasp_lift_params["grasp_preload_position"] = self.gripper_preload_pos + if initialize_stage_gates: + progress_params["active_through_stage"] = self.curriculum_stage_names.index("carry") + self.rewards.approach_progress.params["active_from_stage"] = self.curriculum_stage_names.index("approach_1") + grasp_lift_params["active_from_stage"] = self.curriculum_stage_names.index("near_carry") + delivered_params = self.rewards.delivered.params + delivered_params["min_lift_height"] = self.success_min_lift_height + delivered_params["max_tcp_distance"] = self.success_max_tcp_distance + delivered_params["max_gripper_width_error"] = self.success_max_gripper_width_error + delivered_params["max_gripper_command"] = max_gripper_command + + def __post_init__(self): + self.actions.gripper_action.close_position = max(0.0, self.gripper_preload_pos - self.gripper_close_offset) + self.actions.gripper_action.open_position = self.gripper_open_pos + if self.actions.gripper_action.limit_to_preload: + self.actions.gripper_action.neutral_position = self.gripper_preload_pos + else: + self.actions.gripper_action.neutral_position = self.gripper_open_pos + self.actions.gripper_action.default_position = self.gripper_preload_pos + if not self.actions.gripper_action.use_incremental_target: + self.actions.gripper_action.scale = self.gripper_open_pos - self.gripper_preload_pos + max_gripper_command = self._resolved_success_max_gripper_command() + self._configure_reward_cfg(max_gripper_command, initialize_stage_gates=True) + self.terminations.lost_grasp.params["dwell_time_s"] = self.lost_grasp_dwell_time_s + self.terminations.lost_grasp.params["max_gripper_command"] = max_gripper_command + self.terminations.success.params["max_gripper_command"] = max_gripper_command + self.decimation = 2 + # Recycle failed attempts promptly after the expected manipulation sequence. + self.episode_length_s = 5.0 + # The deadline is part of the task: an attempt that has not poured within five seconds is + # a failed finite-horizon episode and must not be value-bootstrapped by RL wrappers. + self.is_finite_horizon = True + self.sim.dt = 1.0 / 120.0 + self.sim.render_interval = self.decimation + self.sim.use_newton_actuators = False + self.viewer.eye = (1.4, 1.4, 0.9) + self.viewer.lookat = (0.5, 0.0, 0.1) + self.viewer.origin_type = "env" + self.viewer.env_index = 0 + + self._validate_curriculum_cfg() + self._validate_particle_workspace_cfg() + self._apply_robot_cfg() + + self.sim.physics = NewtonCfg( + solver_cfg=CouplerProxyCfg( + scene_cfg=self.scene, + entries=[ + CouplerEntryCfg( + name=RIGID_ENTRY, + # Proxy coupling keeps the MPM stable, so the arm integrator can be the faster + # "implicitfast" (unlike base coupling, which needed "euler"). The cup is a + # dynamic rigid body owned by this entry; Newton generates its contacts and + # the proxy bridges its cavity mesh to the MPM solver. + solver_cfg=MJWarpSolverCfg( + use_mujoco_contacts=False, integrator="implicitfast", njmax=510, nconmax=400 + ), + bodies=[ + SceneEntityCfg("robot"), + SceneEntityCfg("source_cup"), + SceneEntityCfg("target_cup"), + ], + include_static_shapes=True, + substeps=self.rigid_entry_substeps, + ), + CouplerEntryCfg( + name=MPM_ENTRY, + solver_cfg=MPMSolverCfg( + voxel_size=self.voxel_size, + grid_type="sparse", + grid_padding=0, + max_active_cell_count=-1, + strain_basis="P0", + transfer_scheme="apic", + max_iterations=self.mpm_iterations, + warmstart_mode="none", + # PIC27 bounds collider work by particle samples. + velocity_basis="Q1", + collider_basis="pic27", + # "forward": the moving cup carries its media ("backward" drains it). + collider_velocity_mode="forward", + # Keep the task's validated nonlinear solve while sparse topology is + # rebuilt eagerly around the physically separated environments. + solver="jacobi", + separate_worlds=True, + ), + all_particles=True, + bodies=[SPILL_FLOOR_LABEL_PATTERN], + include_static_shapes=False, + include_child_joints=False, + substeps=self.mpm_entry_substeps, + in_place=True, + ), + ], + proxies=[ + CouplerProxyMappingCfg( + source=RIGID_ENTRY, + destination=MPM_ENTRY, + bodies=[SceneEntityCfg("source_cup"), SceneEntityCfg("target_cup")], + mass_scale=self.proxy_mass_scale, + mode="lagged", + # Implicit MPM resolves its proxy colliders internally; the shared outer + # pipeline is only needed for rigid MJWarp contacts. + collision_pipeline=lambda _model: None, + ) + ], + iterations=self.proxy_iterations, + ), + # Rigid contacts use Newton's outer pipeline. Implicit MPM handles particle/shape + # collisions internally, so allocating outer soft contacts would waste O(P*S) work. + collision_cfg=NewtonCollisionPipelineCfg(soft_contact_max=0), + num_substeps=self.physics_substeps, + use_cuda_graph=self.use_cuda_graph, + ) + + def _validate_gripper_action_cfg(self) -> None: + """Validate the reset and action targets against the Panda finger range.""" + gripper_action = self.actions.gripper_action + if not isinstance(gripper_action.use_incremental_target, bool): + raise TypeError("Gripper use_incremental_target must be a bool.") + if gripper_action.binary_threshold is not None: + if ( + isinstance(gripper_action.binary_threshold, bool) + or not math.isfinite(gripper_action.binary_threshold) + or not -1.0 < gripper_action.binary_threshold < 1.0 + ): + raise ValueError("Gripper binary_threshold must be finite and lie strictly between -1 and 1.") + if gripper_action.use_incremental_target: + raise ValueError("Binary and incremental gripper targets are mutually exclusive.") + if ( + not math.isfinite(self.gripper_open_pos) + or not 0.0 < self.gripper_open_pos <= 0.04 + or not math.isfinite(gripper_action.close_position) + or not 0.0 <= gripper_action.close_position < self.gripper_open_pos + or not math.isfinite(gripper_action.scale) + or gripper_action.scale <= 0.0 + or not math.isfinite(gripper_action.alpha) + or not 0.0 < gripper_action.alpha <= 1.0 + ): + raise ValueError( + "Gripper action positions must fit the Panda finger range [0, 0.04] with positive scale and a " + "moving-average weight in (0, 1]." + ) + if not math.isclose(gripper_action.open_position, self.gripper_open_pos, rel_tol=0.0, abs_tol=1.0e-9): + raise ValueError("The gripper action open position must match gripper_open_pos.") + if ( + not math.isfinite(self.gripper_preload_pos) + or not gripper_action.close_position <= self.gripper_preload_pos < self.cup_grasp_box_half[1] + ): + raise ValueError("gripper_preload_pos must lie between the closed and geometric contact positions.") + if ( + not math.isfinite(self.gripper_close_offset) + or not 0.0 <= self.gripper_close_offset <= self.gripper_preload_pos + ): + raise ValueError("gripper_close_offset must lie in [0, gripper_preload_pos].") + contact_command_limit = self.cup_grasp_box_half[1] - gripper_action.contact_min_deflection + max_gripper_command = self._resolved_success_max_gripper_command() + if ( + not math.isfinite(max_gripper_command) + or not self.gripper_preload_pos <= max_gripper_command <= contact_command_limit + ): + raise ValueError( + "success_max_gripper_command must lie between the preload target and the largest command that " + "retains contact_min_deflection at the cup." + ) + max_action_position = self.gripper_preload_pos if gripper_action.limit_to_preload else self.gripper_open_pos + if not math.isclose( + gripper_action.neutral_position, + max_action_position, + rel_tol=0.0, + abs_tol=1.0e-9, + ): + raise ValueError("Gripper action maximum does not match its configured operating interval.") + action_span = max_action_position - gripper_action.close_position + if gripper_action.scale > action_span + 1.0e-9: + raise ValueError("Gripper action scale must not exceed its configured operating interval.") + if gripper_action.default_position is not None and not ( + gripper_action.close_position <= gripper_action.default_position <= gripper_action.neutral_position + ): + raise ValueError("Gripper default position must lie within its configured operating interval.") + + def _validate_source_cup_cfg(self) -> None: + """Validate source-cup geometry, grasp proxy, and media fill as one contract.""" + for field_name in ( + "source_cup_inner_width", + "source_cup_inner_depth", + "source_cup_cavity_depth", + "source_cup_wall_thickness", + "source_cup_bottom_thickness", + "cup_mass", + ): + value = getattr(self, field_name) + if not math.isfinite(value) or value <= 0.0: + raise ValueError(f"{field_name} must be finite and positive.") + if not math.isfinite(self.media_fill_frac) or not 0.0 < self.media_fill_frac <= 1.0: + raise ValueError("media_fill_frac must lie in (0, 1].") + + outer_size = ( + self.source_cup_inner_width + 2.0 * self.source_cup_wall_thickness, + self.source_cup_inner_depth + 2.0 * self.source_cup_wall_thickness, + self.source_cup_cavity_depth + self.source_cup_bottom_thickness, + ) + if len(self.cup_grasp_box_half) != 3 or any( + not math.isfinite(value) or value <= 0.0 for value in self.cup_grasp_box_half + ): + raise ValueError("cup_grasp_box_half must contain three finite positive half-extents.") + if any( + not math.isclose(proxy_half, 0.5 * outer, rel_tol=0.0, abs_tol=1.0e-9) + for proxy_half, outer in zip(self.cup_grasp_box_half, outer_size, strict=True) + ): + raise ValueError("cup_grasp_box_half must exactly match the visible source-cup outer envelope.") + if ( + not math.isfinite(self.cup_grasp_height) + or self.cup_grasp_height <= self.source_cup_bottom_thickness + or self.cup_grasp_height >= outer_size[2] + ): + raise ValueError("cup_grasp_height must lie above the source bottom and below its rim.") + if len(self.cup_grasp_tcp_quat_c) != 4 or any(not math.isfinite(value) for value in self.cup_grasp_tcp_quat_c): + raise ValueError("cup_grasp_tcp_quat_c must contain four finite XYZW values.") + quaternion_norm = math.sqrt(sum(value * value for value in self.cup_grasp_tcp_quat_c)) + if not math.isclose(quaternion_norm, 1.0, rel_tol=0.0, abs_tol=1.0e-6): + raise ValueError("cup_grasp_tcp_quat_c must be a unit quaternion.") + + def _validate_curriculum_progress_cfg(self, stage_count: int) -> None: + """Validate success statistics and promotion controls.""" + if not 0.0 < self.curriculum_success_threshold <= 1.0: + raise ValueError("curriculum_success_threshold must lie in (0, 1].") + if not 0.0 < self.curriculum_randomization_promotion_threshold <= self.curriculum_success_threshold: + raise ValueError( + "curriculum_randomization_promotion_threshold must lie in (0, curriculum_success_threshold]." + ) + if self.curriculum_min_resets_per_stage <= 0: + raise ValueError("curriculum_min_resets_per_stage must be positive.") + cohort_count = self.curriculum_min_reset_cohorts_per_stage + if not math.isfinite(cohort_count) or cohort_count < 0.0: + raise ValueError("curriculum_min_reset_cohorts_per_stage must be finite and nonnegative.") + replay_fraction = self.curriculum_previous_stage_replay_fraction + if not math.isfinite(replay_fraction) or replay_fraction < 0.0 or replay_fraction >= 1.0: + raise ValueError("curriculum_previous_stage_replay_fraction must lie in [0, 1).") + entry_replay_fraction = self.curriculum_frontier_entry_replay_fraction + if ( + not math.isfinite(entry_replay_fraction) + or entry_replay_fraction < replay_fraction + or entry_replay_fraction >= 1.0 + ): + raise ValueError( + "curriculum_frontier_entry_replay_fraction must lie in [curriculum_previous_stage_replay_fraction, 1)." + ) + if self.curriculum_start_stage < 0 or self.curriculum_start_stage >= stage_count: + raise ValueError(f"curriculum_start_stage must lie in [0, {stage_count - 1}].") + extent_levels = self.curriculum_randomization_extent_levels + if not extent_levels: + raise ValueError("curriculum_randomization_extent_levels must not be empty.") + if any(not math.isfinite(level) or level < 0.0 or level > 1.0 for level in extent_levels): + raise ValueError("curriculum_randomization_extent_levels must lie in [0, 1].") + if any(previous >= current for previous, current in zip(extent_levels, extent_levels[1:])): + raise ValueError("curriculum_randomization_extent_levels must be strictly increasing.") + if extent_levels[0] != 0.0: + raise ValueError("curriculum_randomization_extent_levels must start at 0.0.") + if not math.isclose(extent_levels[-1], 1.0, rel_tol=0.0, abs_tol=1.0e-9): + raise ValueError("curriculum_randomization_extent_levels must end at 1.0.") + for field_name in ( + "curriculum_independent_arm_fraction_levels", + "curriculum_independent_target_fraction_levels", + ): + fractions = getattr(self, field_name) + if len(fractions) != len(extent_levels): + raise ValueError(f"{field_name} must align with curriculum_randomization_extent_levels.") + if any(not math.isfinite(value) or value < 0.0 or value > 1.0 for value in fractions): + raise ValueError(f"{field_name} must contain values in [0, 1].") + if any(left > right for left, right in zip(fractions, fractions[1:])): + raise ValueError(f"{field_name} must be nondecreasing.") + if not math.isclose(fractions[0], 0.0, rel_tol=0.0, abs_tol=1.0e-9): + raise ValueError(f"{field_name} must start at 0.0.") + if not math.isclose(fractions[-1], 1.0, rel_tol=0.0, abs_tol=1.0e-9): + raise ValueError(f"{field_name} must end at 1.0.") + if self.curriculum_independent_sample_attempts <= 0: + raise ValueError("curriculum_independent_sample_attempts must be positive.") + if ( + not math.isfinite(self.curriculum_independent_arm_min_tcp_distance) + or self.curriculum_independent_arm_min_tcp_distance <= 0.0 + ): + raise ValueError("curriculum_independent_arm_min_tcp_distance must be finite and positive.") + if self.curriculum_randomization_start_level < 0 or self.curriculum_randomization_start_level >= len( + extent_levels + ): + raise ValueError("curriculum_randomization_start_level must index curriculum_randomization_extent_levels.") + self._validate_reward_cfg(stage_count) + + def _validate_solver_cfg(self) -> None: + """Validate the task-level controls copied into the coupled solver tree.""" + for name in ( + "physics_substeps", + "rigid_entry_substeps", + "mpm_entry_substeps", + "mpm_iterations", + "proxy_iterations", + ): + value = getattr(self, name) + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError(f"{name} must be a positive integer.") + for name in ("voxel_size", "proxy_mass_scale"): + value = getattr(self, name) + if not math.isfinite(value) or value <= 0.0: + raise ValueError(f"{name} must be finite and positive.") + if not isinstance(self.use_cuda_graph, bool): + raise TypeError("use_cuda_graph must be a bool.") + + def _validate_reward_cfg(self, stage_count: int) -> None: + """Validate the reverse-curriculum reward parameters.""" + tilt_params = self.rewards.task_progress.params + target_tilt = float(tilt_params["target_tilt"]) + pour_direction_xy = tilt_params["pour_direction_xy"] + source_offset_xy = self.pour_source_offset_xy + source_mouth_height = float(tilt_params["source_mouth_height"]) + alignment_radius = float(tilt_params["alignment_radius"]) + active_through_stage = int(tilt_params["active_through_stage"]) + discount_factor = float(tilt_params["discount_factor"]) + if not math.isfinite(target_tilt) or not 0.0 < target_tilt < math.pi: + raise ValueError("task_progress target_tilt must lie in (0, pi).") + if ( + len(pour_direction_xy) != 2 + or any(not math.isfinite(value) for value in pour_direction_xy) + or math.hypot(float(pour_direction_xy[0]), float(pour_direction_xy[1])) <= 0.0 + ): + raise ValueError("task_progress pour_direction_xy must contain two finite values and be nonzero.") + if len(source_offset_xy) != 2 or any(not math.isfinite(value) for value in source_offset_xy): + raise ValueError("pour_source_offset_xy must contain two finite values.") + if not math.isfinite(source_mouth_height) or source_mouth_height <= 0.0: + raise ValueError("task_progress source_mouth_height must be finite and positive.") + if not math.isfinite(alignment_radius) or alignment_radius <= 0.0: + raise ValueError("task_progress alignment_radius must be finite and positive.") + if active_through_stage < 0 or active_through_stage >= stage_count: + raise ValueError(f"task_progress active_through_stage must lie in [0, {stage_count - 1}].") + if active_through_stage != self.curriculum_stage_names.index("carry"): + raise ValueError("task_progress active_through_stage must select the carry curriculum stage.") + if not math.isfinite(discount_factor) or not 0.0 < discount_factor <= 1.0: + raise ValueError("task_progress discount_factor must lie in (0, 1].") + self._validate_guidance_reward_cfg() + + def _validate_guidance_reward_cfg(self) -> None: + """Validate full-task approach and grasp-lift potential settings.""" + approach_params = self.rewards.approach_progress.params + for parameter in ("position_std", "orientation_std"): + value = float(approach_params[parameter]) + if not math.isfinite(value) or value <= 0.0: + raise ValueError(f"approach_progress {parameter} must be finite and positive.") + open_hand_fraction = float(approach_params["open_hand_fraction"]) + if not math.isfinite(open_hand_fraction) or not 0.0 <= open_hand_fraction <= 1.0: + raise ValueError("approach_progress open_hand_fraction must lie in [0, 1].") + if int(approach_params["active_from_stage"]) != self.curriculum_stage_names.index("approach_1"): + raise ValueError("approach_progress active_from_stage must select the approach_1 curriculum stage.") + + grasp_lift_params = self.rewards.grasp_lift_progress.params + for parameter in ("target_height", "grasp_reach_std"): + value = float(grasp_lift_params[parameter]) + if not math.isfinite(value) or value <= 0.0: + raise ValueError(f"grasp_lift_progress {parameter} must be finite and positive.") + if not math.isfinite(float(grasp_lift_params["grasp_preload_position"])): + raise ValueError("grasp_lift_progress grasp_preload_position must be finite.") + grasp_fraction = float(grasp_lift_params["grasp_fraction"]) + if not math.isfinite(grasp_fraction) or not 0.0 <= grasp_fraction <= 1.0: + raise ValueError("grasp_lift_progress grasp_fraction must lie in [0, 1].") + if int(grasp_lift_params["active_from_stage"]) != self.curriculum_stage_names.index("near_carry"): + raise ValueError("grasp_lift_progress active_from_stage must select the near_carry curriculum stage.") + + for name, params in ( + ("approach_progress", approach_params), + ("grasp_lift_progress", grasp_lift_params), + ): + value = float(params["discount_factor"]) + if not math.isfinite(value) or not 0.0 < value <= 1.0: + raise ValueError(f"{name} discount_factor must lie in (0, 1].") + + def _validate_arm_action_cfg(self) -> None: + """Validate the task's direct relative arm-joint action.""" + arm_action = self.actions.arm_action + if not isinstance(arm_action, mdp.RelativeJointPositionActionCfg): + raise ValueError("Franka Pour requires RelativeJointPositionActionCfg for the arm.") + if arm_action.joint_names != [f"panda_joint{i}" for i in range(1, 8)] or not arm_action.preserve_order: + raise ValueError("Franka Pour requires all seven Panda arm joints in kinematic order.") + if not isinstance(arm_action.scale, float) or not math.isfinite(arm_action.scale) or arm_action.scale <= 0.0: + raise ValueError("Arm relative-action scale must be a finite positive scalar.") + if not arm_action.use_zero_offset: + raise ValueError("Franka Pour relative arm actions require use_zero_offset=True.") + + def _validate_curriculum_cfg(self) -> None: + """Validate the aligned per-stage backward-curriculum settings.""" + self._validate_solver_cfg() + self._validate_source_cup_cfg() + self._validate_gripper_action_cfg() + self._validate_arm_action_cfg() + stage_count = len(self.curriculum_stage_names) + if self.curriculum_stage_names != CURRICULUM_STAGE_NAMES: + raise ValueError(f"curriculum_stage_names must be {CURRICULUM_STAGE_NAMES!r}.") + if len(self.curriculum_target_frac) != stage_count: + raise ValueError( + f"curriculum_target_frac has {len(self.curriculum_target_frac)} values for {stage_count} stages." + ) + if any( + not math.isfinite(fraction) or fraction <= 0.0 or fraction > 1.0 for fraction in self.curriculum_target_frac + ): + raise ValueError("Curriculum target fractions must lie in (0, 1].") + if tuple(sorted(self.curriculum_target_frac)) != self.curriculum_target_frac: + raise ValueError("Curriculum target fractions must be nondecreasing.") + self._validate_transport_reset_cfg() + self._validate_grasp_approach_reset_cfg() + self._validate_curriculum_progress_cfg(stage_count) + if self.cup_grasp_box_half[1] < 0.0 or self.cup_grasp_box_half[1] > self.gripper_open_pos: + raise ValueError("The curriculum contact position must fit within the open gripper.") + self._validate_randomized_reset_cfg() + self._validate_curriculum_arm_configs() + + def _validate_transport_reset_cfg(self) -> None: + """Validate intermediate held-cup reset locations between pour and carry.""" + pour_stage = self.curriculum_stage_names.index("pour") + carry_stage = self.curriculum_stage_names.index("carry") + fractions = self.curriculum_transport_reset_fractions + if len(fractions) != carry_stage - pour_stage - 1: + raise ValueError( + "curriculum_transport_reset_fractions must provide one value for every stage between pour and carry." + ) + if any(not math.isfinite(value) or not 0.0 < value < 1.0 for value in fractions) or any( + left >= right for left, right in zip(fractions, fractions[1:]) + ): + raise ValueError("curriculum_transport_reset_fractions must increase strictly within (0, 1).") + + def _validate_grasp_approach_reset_cfg(self) -> None: + """Validate collision-screened reset samples between exact grasp and pre-grasp.""" + grasp_stage = self.curriculum_stage_names.index("grasp") + full_stage = self.curriculum_stage_names.index("full") + fractions = self.curriculum_grasp_approach_fractions + if len(fractions) != full_stage - grasp_stage - 1: + raise ValueError( + "curriculum_grasp_approach_fractions must provide one value for every stage between grasp and full." + ) + if any(not math.isfinite(value) or value < 0.0 or value >= 1.0 for value in fractions) or any( + left <= right for left, right in zip(fractions, fractions[1:]) + ): + raise ValueError("curriculum_grasp_approach_fractions must decrease strictly within [0, 1).") + if any(not math.isclose(8.0 * value, round(8.0 * value), abs_tol=1.0e-9) for value in fractions): + raise ValueError( + "curriculum_grasp_approach_fractions must select eighth-segment samples covered by collision screening." + ) + + def _validate_randomized_reset_cfg(self) -> None: + """Validate randomized cup poses and their precomputed IK reset bank.""" + self._validate_randomized_workspace_cfg() + self._validate_randomized_reset_offset_bounds_cfg() + self._validate_randomized_reset_selection_cfg() + self._validate_randomized_grasp_approach_cfg() + self._validate_randomized_ik_solver_cfg() + self._validate_randomized_cup_separation_cfg() + + def _validate_randomized_workspace_cfg(self) -> None: + """Validate randomized source, carry, and receiver workspace geometry.""" + for field_name in ( + "curriculum_randomized_source_position_range", + "curriculum_randomized_carry_position_range", + "curriculum_randomized_target_position_range", + ): + values = getattr(self, field_name) + if len(values) != 2 or any(not math.isfinite(value) or value < 0.0 for value in values): + raise ValueError(f"{field_name} must contain two finite nonnegative values.") + if ( + not math.isfinite(self.curriculum_randomized_source_xy_correlation) + or not 0.0 <= self.curriculum_randomized_source_xy_correlation < 1.0 + ): + raise ValueError("curriculum_randomized_source_xy_correlation must lie in [0, 1).") + radius_range = self.curriculum_randomized_source_radius_range + if radius_range is not None: + if self.cup_reset_pos[0] <= 0.0 or not math.isclose( + self.cup_reset_pos[1], + 0.0, + rel_tol=0.0, + abs_tol=1.0e-9, + ): + raise ValueError( + "The polar source workspace requires cup_reset_pos to lie on the positive robot-base X axis." + ) + if ( + len(radius_range) != 2 + or any(not math.isfinite(value) or value <= 0.0 for value in radius_range) + or radius_range[0] >= radius_range[1] + ): + raise ValueError( + "curriculum_randomized_source_radius_range must contain two finite positive values " + "in increasing order." + ) + if ( + not math.isfinite(self.curriculum_randomized_source_azimuth_range) + or not 0.0 < self.curriculum_randomized_source_azimuth_range < math.pi / 2.0 + ): + raise ValueError("curriculum_randomized_source_azimuth_range must lie in (0, pi / 2).") + min_radius, max_radius = radius_range + azimuth = self.curriculum_randomized_source_azimuth_range + minimum_x = min_radius * math.cos(azimuth) + maximum_x = max_radius + maximum_abs_y = max_radius * math.sin(azimuth) + required_half_range = ( + max(abs(minimum_x - self.cup_reset_pos[0]), abs(maximum_x - self.cup_reset_pos[0])), + maximum_abs_y, + ) + if any( + configured + 1.0e-9 < required + for configured, required in zip( + self.curriculum_randomized_source_position_range, + required_half_range, + strict=True, + ) + ): + raise ValueError( + "curriculum_randomized_source_position_range must contain the configured polar workspace; " + f"required at least {required_half_range}." + ) + elif not math.isfinite(self.curriculum_randomized_source_azimuth_range): + raise ValueError("curriculum_randomized_source_azimuth_range must be finite.") + if any( + carry_range > source_range + for carry_range, source_range in zip( + self.curriculum_randomized_carry_position_range, + self.curriculum_randomized_source_position_range, + strict=True, + ) + ): + raise ValueError( + "curriculum_randomized_carry_position_range must not exceed " + "curriculum_randomized_source_position_range." + ) + if len(self.curriculum_randomized_target_center_xy) != 2 or any( + not math.isfinite(value) for value in self.curriculum_randomized_target_center_xy + ): + raise ValueError("curriculum_randomized_target_center_xy must contain two finite values.") + if ( + not math.isfinite(self.curriculum_randomized_source_yaw_range) + or self.curriculum_randomized_source_yaw_range < 0.0 + or self.curriculum_randomized_source_yaw_range > math.pi / 4.0 + ): + raise ValueError("curriculum_randomized_source_yaw_range must lie in [0, pi / 4].") + if not math.isfinite(self.curriculum_randomized_cup_clearance) or self.curriculum_randomized_cup_clearance < 0: + raise ValueError("curriculum_randomized_cup_clearance must be finite and nonnegative.") + + def _validate_randomized_reset_offset_bounds_cfg(self) -> None: + """Validate legacy jitter and optional asymmetric reset-TCP offset bounds.""" + for field_name in ( + "curriculum_randomized_reset_tcp_standoff", + "curriculum_randomized_reset_tcp_jitter", + ): + values = getattr(self, field_name) + if len(values) != 3 or any(not math.isfinite(value) for value in values): + raise ValueError(f"{field_name} must contain three finite values.") + if any(value < 0.0 for value in self.curriculum_randomized_reset_tcp_jitter): + raise ValueError("curriculum_randomized_reset_tcp_jitter must contain three finite nonnegative values.") + offset_lower = self.curriculum_randomized_reset_tcp_offset_lower + offset_upper = self.curriculum_randomized_reset_tcp_offset_upper + if (offset_lower is None) != (offset_upper is None): + raise ValueError( + "curriculum_randomized_reset_tcp_offset_lower and " + "curriculum_randomized_reset_tcp_offset_upper must either both be set or both be None." + ) + if offset_lower is not None and offset_upper is not None: + if len(offset_lower) != 3 or len(offset_upper) != 3: + raise ValueError("Randomized reset TCP offset bounds must each contain three values.") + if any(not math.isfinite(value) for value in (*offset_lower, *offset_upper)): + raise ValueError("Randomized reset TCP offset bounds must be finite.") + if any(lower > 0.0 or upper < 0.0 or lower > upper for lower, upper in zip(offset_lower, offset_upper)): + raise ValueError("Randomized reset TCP offset bounds must be ordered and contain zero.") + if offset_lower[2] < 0.0: + raise ValueError("Randomized reset TCP offsets must not place the initial TCP below its pre-grasp.") + rotation_range = self.curriculum_randomized_reset_tcp_rotation_angle_range + if len(rotation_range) != 2 or any(not math.isfinite(value) for value in rotation_range): + raise ValueError("curriculum_randomized_reset_tcp_rotation_angle_range must contain two finite values.") + rotation_lower, rotation_upper = rotation_range + if rotation_lower < 0.0 or rotation_lower > rotation_upper or rotation_upper > math.pi / 2.0: + raise ValueError("curriculum_randomized_reset_tcp_rotation_angle_range must be ordered within [0, pi / 2].") + + def _validate_randomized_reset_selection_cfg(self) -> None: + """Validate reset-bank filtering thresholds and row-coverage requirements.""" + if ( + not math.isfinite(self.curriculum_randomized_reset_tcp_min_grasp_distance) + or self.curriculum_randomized_reset_tcp_min_grasp_distance <= 0.0 + ): + raise ValueError("curriculum_randomized_reset_tcp_min_grasp_distance must be finite and positive.") + if ( + not math.isfinite(self.curriculum_randomized_reset_joint6_max) + or self.curriculum_randomized_reset_joint6_max <= 0.0 + ): + raise ValueError("curriculum_randomized_reset_joint6_max must be finite and positive.") + if ( + not math.isfinite(self.curriculum_randomized_min_source_cell_fraction) + or not 0.0 < self.curriculum_randomized_min_source_cell_fraction <= 1.0 + ): + raise ValueError("curriculum_randomized_min_source_cell_fraction must lie in (0, 1].") + if ( + self.curriculum_randomized_min_reset_variants_per_source < 1 + or self.curriculum_randomized_min_reset_variants_per_source + > self.curriculum_randomized_reset_ik_samples_per_source + ): + raise ValueError( + "curriculum_randomized_min_reset_variants_per_source must lie between one and " + "curriculum_randomized_reset_ik_samples_per_source." + ) + if ( + not math.isfinite(self.curriculum_randomized_pour_clearance) + or self.curriculum_randomized_pour_clearance < 0.0 + ): + raise ValueError("curriculum_randomized_pour_clearance must be finite and nonnegative.") + if not math.isfinite(self.curriculum_grasp_descent_overshoot) or self.curriculum_grasp_descent_overshoot < 0.0: + raise ValueError("curriculum_grasp_descent_overshoot must be finite and nonnegative.") + + def _validate_randomized_grasp_approach_cfg(self) -> None: + """Validate horizontal grasp orientation, standoff, and minimum clearance.""" + offset_lower = self.curriculum_randomized_reset_tcp_offset_lower + offset_upper = self.curriculum_randomized_reset_tcp_offset_upper + grasp_qx, grasp_qy, grasp_qz, grasp_qw = self.cup_grasp_tcp_quat_c + tool_axis_c = ( + 2.0 * (grasp_qx * grasp_qz + grasp_qy * grasp_qw), + 2.0 * (grasp_qy * grasp_qz - grasp_qx * grasp_qw), + 1.0 - 2.0 * (grasp_qx * grasp_qx + grasp_qy * grasp_qy), + ) + if abs(tool_axis_c[2]) > 1.0e-6: + raise ValueError("cup_grasp_tcp_quat_c must keep Panda tool +Z parallel to the table.") + jaw_axis_z = 2.0 * (grasp_qy * grasp_qz + grasp_qx * grasp_qw) + if abs(jaw_axis_z) > 1.0e-6: + raise ValueError("cup_grasp_tcp_quat_c must keep the Panda jaw axis parallel to the table.") + standoff_norm = math.sqrt(sum(value * value for value in self.curriculum_randomized_reset_tcp_standoff)) + if standoff_norm <= 1.0e-9: + raise ValueError("curriculum_randomized_reset_tcp_standoff must be nonzero.") + standoff_alignment = sum( + tool_axis * standoff / standoff_norm + for tool_axis, standoff in zip( + tool_axis_c, + self.curriculum_randomized_reset_tcp_standoff, + strict=True, + ) + ) + if not math.isclose(standoff_alignment, -1.0, rel_tol=0.0, abs_tol=1.0e-6): + raise ValueError( + "curriculum_randomized_reset_tcp_standoff must be antiparallel to Panda tool +Z " + "from cup_grasp_tcp_quat_c." + ) + if ( + offset_lower is None + and self.cup_grasp_height - self.curriculum_randomized_reset_tcp_jitter[2] <= self.collider_margin + ): + raise ValueError( + "cup_grasp_height and curriculum_randomized_reset_tcp_jitter must keep every reset TCP " + "above the table by more than collider_margin." + ) + if offset_lower is None or offset_upper is None: + minimum_standoff = math.sqrt( + sum( + max(abs(offset) - jitter, 0.0) ** 2 + for offset, jitter in zip( + self.curriculum_randomized_reset_tcp_standoff, + self.curriculum_randomized_reset_tcp_jitter, + strict=True, + ) + ) + ) + else: + closest_offset = tuple( + min(max(-standoff, lower), upper) + for standoff, lower, upper in zip( + self.curriculum_randomized_reset_tcp_standoff, + offset_lower, + offset_upper, + strict=True, + ) + ) + minimum_standoff = math.sqrt( + sum( + (standoff + offset) ** 2 + for standoff, offset in zip( + self.curriculum_randomized_reset_tcp_standoff, + closest_offset, + strict=True, + ) + ) + ) + if minimum_standoff + 1.0e-9 < self.curriculum_randomized_reset_tcp_min_grasp_distance: + raise ValueError( + "curriculum_randomized_reset_tcp_standoff and curriculum_randomized_reset_tcp_jitter " + "cannot guarantee curriculum_randomized_reset_tcp_min_grasp_distance." + ) + + def _validate_randomized_ik_solver_cfg(self) -> None: + """Validate randomized reset-bank IK discretization and solver tolerances.""" + if self.curriculum_randomized_reset_ik_grid_size < 3 or self.curriculum_randomized_reset_ik_grid_size % 2 == 0: + raise ValueError("curriculum_randomized_reset_ik_grid_size must be an odd integer of at least three.") + if ( + self.curriculum_randomized_reset_ik_samples_per_source < 3 + or self.curriculum_randomized_reset_ik_samples_per_source % 2 == 0 + ): + raise ValueError( + "curriculum_randomized_reset_ik_samples_per_source must be an odd integer of at least three." + ) + if self.curriculum_randomized_reset_ik_iterations <= 0: + raise ValueError("curriculum_randomized_reset_ik_iterations must be positive.") + for field_name in ( + "curriculum_randomized_reset_ik_max_cost", + "curriculum_randomized_reset_ik_joint_margin", + ): + value = getattr(self, field_name) + if not math.isfinite(value) or value < 0.0: + raise ValueError(f"{field_name} must be finite and nonnegative.") + + def _validate_randomized_cup_separation_cfg(self) -> None: + """Validate that every randomized source pose admits a separated receiver pose.""" + source_outer_half_x = self.source_cup_inner_width / 2.0 + self.source_cup_wall_thickness + source_outer_half_y = self.source_cup_inner_depth / 2.0 + self.source_cup_wall_thickness + target_outer_half_y = self.target_cup_inner_depth / 2.0 + self.target_cup_wall_thickness + if self.curriculum_randomized_source_radius_range is None: + maximum_projection_yaw = min( + self.curriculum_randomized_source_yaw_range, + math.atan2(source_outer_half_x, source_outer_half_y), + ) + maximum_source_half_y = source_outer_half_x * math.sin( + maximum_projection_yaw + ) + source_outer_half_y * math.cos(maximum_projection_yaw) + else: + # A radial-facing cup can reach any global yaw in the polar sector. The diagonal is the + # conservative support of the square source under arbitrary upright yaw. + maximum_source_half_y = math.hypot(source_outer_half_x, source_outer_half_y) + minimum_separation = maximum_source_half_y + target_outer_half_y + self.curriculum_randomized_cup_clearance + if self.curriculum_randomized_source_radius_range is None: + minimum_source_y = self.cup_reset_pos[1] - self.curriculum_randomized_source_position_range[1] + else: + minimum_source_y = -self.curriculum_randomized_source_radius_range[1] * math.sin( + self.curriculum_randomized_source_azimuth_range + ) + minimum_target_y = ( + self.curriculum_randomized_target_center_xy[1] - self.curriculum_randomized_target_position_range[1] + ) + if minimum_source_y - minimum_separation < minimum_target_y - 1.0e-6: + raise ValueError( + "curriculum_randomized_target_position_range leaves no collision-free target y-position " + "at the minimum randomized source y-position." + ) + + def _validate_curriculum_arm_configs(self) -> None: + """Validate authored arm waypoints against the action limits.""" + arm_configs = ( + self.curriculum_drain_arm_q, + self.curriculum_deep_tilt_arm_q, + self.curriculum_tilt_arm_q, + self.curriculum_pour_arm_q, + *self._curriculum_transport_arm_configs(), + self.curriculum_pour_target_arm_q, + self.curriculum_carry_arm_q, + self.arm_home, + ) + for arm_q in arm_configs: + if len(arm_q) != 7: + raise ValueError("Every curriculum arm configuration must contain seven joint positions.") + for joint_name, position, (lower, upper) in zip( + self.actions.arm_action.joint_names, + arm_q, + PANDA_ARM_JOINT_LIMITS, + strict=True, + ): + if not math.isfinite(position) or position < lower or position > upper: + raise ValueError( + f"Curriculum joint position {joint_name}={position} lies outside [{lower}, {upper}]." + ) + + def _validate_particle_workspace_cfg(self) -> None: + """Validate finite local particle bounds and all configured media reset poses.""" + if not math.isfinite(self.particle_max_velocity) or self.particle_max_velocity <= 0.0: + raise ValueError("particle_max_velocity must be finite and positive.") + if not math.isfinite(self.spill_table_height): + raise ValueError("spill_table_height must be finite.") + if not 0.0 < self.max_spill_fraction < 1.0: + raise ValueError("max_spill_fraction must lie in (0, 1).") + if not math.isfinite(self.success_dwell_time_s) or self.success_dwell_time_s <= 0.0: + raise ValueError("success_dwell_time_s must be finite and positive.") + if not math.isfinite(self.lost_grasp_dwell_time_s) or self.lost_grasp_dwell_time_s <= 0.0: + raise ValueError("lost_grasp_dwell_time_s must be finite and positive.") + for field_name in ( + "success_min_lift_height", + "success_max_tcp_distance", + "success_max_gripper_width_error", + ): + value = getattr(self, field_name) + if not math.isfinite(value) or value <= 0.0: + raise ValueError(f"{field_name} must be finite and positive.") + if not math.isfinite(self.state_bound_joint_position_margin) or self.state_bound_joint_position_margin < 0.0: + raise ValueError("state_bound_joint_position_margin must be finite and nonnegative.") + for field_name in ( + "state_bound_max_joint_velocity", + "state_bound_max_cup_linear_velocity", + "state_bound_max_cup_angular_velocity", + ): + value = getattr(self, field_name) + if not math.isfinite(value) or value <= 0.0: + raise ValueError(f"{field_name} must be finite and positive.") + if not math.isfinite(self.particle_count_margin) or self.particle_count_margin < 0.0: + raise ValueError("particle_count_margin must be finite and nonnegative.") + lower = self.particle_workspace_lower_bound + upper = self.particle_workspace_upper_bound + if len(lower) != 3 or len(upper) != 3: + raise ValueError("particle_workspace bounds must each contain three coordinates.") + if any(not math.isfinite(value) for value in (*lower, *upper)): + raise ValueError("particle_workspace bounds must be finite.") + if any(lo >= hi for lo, hi in zip(lower, upper, strict=True)): + raise ValueError("particle_workspace lower bounds must be smaller than upper bounds.") + if not lower[2] <= self.spill_table_height <= upper[2]: + raise ValueError("spill_table_height must lie inside the particle workspace z bounds.") + + local_points = cup_cavity_lattice(self)[0] + local_lo = local_points.min(axis=0) + local_hi = local_points.max(axis=0) + source_range = (*self.curriculum_randomized_source_position_range, 0.0) + # A radial XY envelope is conservative for every configured upright yaw and keeps this + # validation independent of the reset bank's finite angular samples. + local_xy_radius = max(math.hypot(float(point[0]), float(point[1])) for point in local_points) + source_lo = ( + self.cup_reset_pos[0] - source_range[0] - local_xy_radius, + self.cup_reset_pos[1] - source_range[1] - local_xy_radius, + float(local_lo[2] + self.cup_reset_pos[2]), + ) + source_hi = ( + self.cup_reset_pos[0] + source_range[0] + local_xy_radius, + self.cup_reset_pos[1] + source_range[1] + local_xy_radius, + float(local_hi[2] + self.cup_reset_pos[2]), + ) + target_local_lo, target_local_hi = cube_bowl_inner_bounds( + self.target_cup_inner_width, + self.target_cup_inner_depth, + self.target_cup_cavity_depth, + self.target_cup_bottom_thickness, + ) + target_range = (*self.curriculum_randomized_target_position_range, 0.0) + target_lo = tuple( + float(point + position - extent - self.particle_count_margin) + for point, position, extent in zip( + target_local_lo, + self.target_cup_reset_pos, + target_range, + strict=True, + ) + ) + target_hi = tuple( + float(point + position + extent + self.particle_count_margin) + for point, position, extent in zip( + target_local_hi, + self.target_cup_reset_pos, + target_range, + strict=True, + ) + ) + for region_name, region_lo, region_hi in ( + ("randomized source media", source_lo, source_hi), + ("randomized target cavity", target_lo, target_hi), + ): + if any(value < bound for value, bound in zip(region_lo, lower, strict=True)) or any( + value > bound for value, bound in zip(region_hi, upper, strict=True) + ): + raise ValueError(f"particle_workspace bounds do not contain the {region_name}.") + + def _apply_robot_cfg(self) -> None: + """Apply final task reset positions to the scene robot.""" + self.scene.robot.init_state.joint_pos.update( + dict(zip([f"panda_joint{i}" for i in range(1, 8)], self.arm_home, strict=True)) + ) + self.scene.robot.init_state.joint_pos["panda_finger_joint.*"] = self.gripper_open_pos + + def _apply_solver_cfg_overrides(self) -> None: + """Propagate final top-level controls into the constructed coupled-solver config.""" + coupled_cfg = self.sim.physics.solver_cfg + arm_entries = [entry for entry in coupled_cfg.entries if entry.name == RIGID_ENTRY] + if len(arm_entries) != 1: + raise ValueError(f"Expected exactly one {RIGID_ENTRY!r} solver entry, found {len(arm_entries)}.") + media_entries = [entry for entry in coupled_cfg.entries if entry.name == MPM_ENTRY] + if len(media_entries) != 1: + raise ValueError(f"Expected exactly one {MPM_ENTRY!r} solver entry, found {len(media_entries)}.") + proxies = [ + proxy for proxy in coupled_cfg.proxies if proxy.source == RIGID_ENTRY and proxy.destination == MPM_ENTRY + ] + if len(proxies) != 1: + raise ValueError(f"Expected exactly one {RIGID_ENTRY!r}-to-{MPM_ENTRY!r} proxy, found {len(proxies)}.") + + mpm_solver_cfg = _mpm_solver_cfg(self) + mpm_solver_cfg.voxel_size = self.voxel_size + mpm_solver_cfg.max_iterations = self.mpm_iterations + arm_entries[0].substeps = self.rigid_entry_substeps + media_entries[0].substeps = self.mpm_entry_substeps + self.sim.physics.num_substeps = self.physics_substeps + self.sim.physics.use_cuda_graph = self.use_cuda_graph + coupled_cfg.iterations = self.proxy_iterations + proxies[0].mass_scale = self.proxy_mass_scale + + def finalize(self) -> FrankaPourEnvCfg: + """Return an independent config with all derived scene assets resolved.""" + resolved = deepcopy(self) + resolved.sim.render_interval = resolved.decimation + # Hydra and command-line overrides are applied after ``__post_init__`` constructs the + # nested Newton solver tree. Reapply every public top-level solver control before resolving + # derived configuration. + resolved._apply_solver_cfg_overrides() + # Command-line overrides are applied after ``__post_init__``. Re-resolve the custom action + # bound so its open target cannot diverge from the physical reset configuration. + resolved.actions.gripper_action.close_position = max( + 0.0, resolved.gripper_preload_pos - resolved.gripper_close_offset + ) + resolved.actions.gripper_action.open_position = resolved.gripper_open_pos + if resolved.actions.gripper_action.limit_to_preload: + resolved.actions.gripper_action.neutral_position = resolved.gripper_preload_pos + else: + resolved.actions.gripper_action.neutral_position = resolved.gripper_open_pos + resolved.actions.gripper_action.default_position = resolved.gripper_preload_pos + if not resolved.actions.gripper_action.use_incremental_target: + resolved.actions.gripper_action.scale = resolved.gripper_open_pos - resolved.gripper_preload_pos + max_gripper_command = resolved._resolved_success_max_gripper_command() + resolved._configure_reward_cfg(max_gripper_command, initialize_stage_gates=False) + resolved.terminations.lost_grasp.params["dwell_time_s"] = resolved.lost_grasp_dwell_time_s + resolved.terminations.lost_grasp.params["max_tcp_distance"] = resolved.success_max_tcp_distance + resolved.terminations.lost_grasp.params["max_gripper_width_error"] = resolved.success_max_gripper_width_error + resolved.terminations.lost_grasp.params["max_gripper_command"] = max_gripper_command + resolved._validate_curriculum_cfg() + resolved._validate_particle_workspace_cfg() + resolved._apply_robot_cfg() + if resolved.terminations.success.func is mdp.immediate_pour_success: + resolved.terminations.success.params = {} + else: + resolved.terminations.success.params["dwell_time_s"] = resolved.success_dwell_time_s + resolved.terminations.success.params["min_lift_height"] = resolved.success_min_lift_height + resolved.terminations.success.params["max_tcp_distance"] = resolved.success_max_tcp_distance + resolved.terminations.success.params["max_gripper_width_error"] = resolved.success_max_gripper_width_error + resolved.terminations.success.params["max_gripper_command"] = max_gripper_command + resolved.scene.source_cup = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/SourceCup", + init_state=RigidObjectCfg.InitialStateCfg( + pos=resolved.cup_reset_pos, + rot=(0.0, 0.0, 0.0, 1.0), + ), + spawn=CubeBowlSpawnerCfg( + inner_width=resolved.source_cup_inner_width, + inner_depth=resolved.source_cup_inner_depth, + cavity_depth=resolved.source_cup_cavity_depth, + wall_thickness=resolved.source_cup_wall_thickness, + bottom_thickness=resolved.source_cup_bottom_thickness, + display_color=(0.95, 0.82, 0.16), + grasp_proxy_half_extents=resolved.cup_grasp_box_half, + mass_props=MassCfg(mass=resolved.cup_mass), + rigid_props=UsdPhysicsRigidBodyCfg(rigid_body_enabled=True, kinematic_enabled=False), + collision_props=UsdPhysicsCollisionCfg(collision_enabled=True), + physics_material=RigidBodyMaterialBaseCfg( + static_friction=resolved.cup_grasp_box_friction, + dynamic_friction=resolved.cup_grasp_box_friction, + ), + ), + ) + resolved.scene.target_cup = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/TargetCup", + init_state=RigidObjectCfg.InitialStateCfg( + pos=resolved.target_cup_reset_pos, + rot=(0.0, 0.0, 0.0, 1.0), + ), + spawn=CubeBowlSpawnerCfg( + inner_width=resolved.target_cup_inner_width, + inner_depth=resolved.target_cup_inner_depth, + cavity_depth=resolved.target_cup_cavity_depth, + wall_thickness=resolved.target_cup_wall_thickness, + bottom_thickness=resolved.target_cup_bottom_thickness, + display_color=(0.20, 0.55, 0.90), + grasp_proxy_half_extents=None, + rigid_props=UsdPhysicsRigidBodyCfg(rigid_body_enabled=True, kinematic_enabled=True), + physics_material=RigidBodyMaterialBaseCfg( + static_friction=resolved.target_cup_friction, + dynamic_friction=resolved.target_cup_friction, + ), + ), + ) + resolved.scene.media = build_media_object_cfg( + resolved, + resolved.cup_reset_pos, + (0.0, 0.0, 0.0, 1.0), + ) + mpm_solver_cfg = _mpm_solver_cfg(resolved) + mpm_solver_cfg.max_active_cell_count = _resolve_mpm_cell_cap(resolved) + if mpm_solver_cfg.grid_type == "sparse" and resolved.sim.physics.use_cuda_graph: + # The compact initial fill underestimates hierarchy nodes needed after a particle moves + # into a different NanoVDB region. Reserve the task's workspace-derived headroom per + # independent world; this changes capacity only, not MPM stepping or physics. + world_count = int(resolved.scene.num_envs) + mpm_solver_cfg.max_lower_node_count = max(32, 16 * world_count) + mpm_solver_cfg.max_upper_node_count = max(32, (world_count + 1) // 2) + resolved.sim.physics.solver_cfg.scene_cfg = resolved.scene + return resolved + + +@configclass +class FrankaPourEnvCfg_RESET_DATASET(FrankaPourEnvCfg): + """Validated reset-dataset curriculum with a stage-independent reward.""" + + curriculum: ResetDatasetCurriculumCfg = ResetDatasetCurriculumCfg() + rewards: ResetDatasetRewardsCfg = ResetDatasetRewardsCfg() + curriculum_early_target_frac: tuple[float, ...] = (0.30,) * 14 + # The dataset is generated and collision-validated offline, then restored directly at reset. + # Training samples it through the adaptive curriculum below; frozen evaluation samples either + # every row or the requested highest-objective grasp subset. + reset_dataset_path: str = "datasets/franka_pour/reset_dataset.pt" + reset_dataset_content_sha256: str | None = None + reset_dataset_top_grasp_count: int | None = None + # Each process retains bounded online success evidence. Training begins on the easiest rows, + # calibrates a broad replay distribution near 50% success, and probes only the adjacent harder + # frontier. The state is intentionally rank-local and starts fresh on RSL-RL resume because + # environment state is not currently part of its checkpoints. The reusable sampler config + # remains directly overridable through Hydra. + reset_dataset_sampler: AdaptiveResetSamplerCfg = AdaptiveResetSamplerCfg() + # A spill is irreversible once more than 30% of the media rests on the table outside both + # vessels. At 245 particles the strict threshold terminates on particle 74. + max_spill_fraction: float = 0.30 + + def __post_init__(self): + # Use the standard Isaac Lab relative joint-position action with the position drives + # authored by the selected Franka USD. A 0.015 rad tracking error keeps the stiff proximal + # USD drives within useful torque authority while improving contact-phase precision. + self.actions.arm_action.scale = 0.015 + # OmniReset uses one region-independent binary gripper action. Keep the existing moving- + # average filter to damp IID exploration chatter without accumulating stale open commands. + # Adjust its coefficient to retain the same physical-time response at three times the rate. + self.actions.gripper_action.use_incremental_target = False + self.actions.gripper_action.binary_threshold = 0.0 + self.actions.gripper_action.alpha = 1.0 - (1.0 - self.actions.gripper_action.alpha) ** (1.0 / 3.0) + super().__post_init__() + # Run the policy at 30 Hz over the 120 Hz simulation. Thirteen frames span 0.4 seconds from + # oldest to newest, matching five frames at 10 Hz, while 32 PPO steps span 1.067 seconds. + self.decimation = 4 + self.sim.render_interval = self.decimation + # Seven seconds gives broad reaching rows time to act while turning failed attempts over + # more quickly. At 30 Hz this is exactly 210 policy steps. + self.episode_length_s = 7.0 + self.observations.policy.history_length = 13 + self.is_finite_horizon = False + # The first step with at least ``pour_target_frac`` (30%) of the media in the receiver is + # terminal success. It deliberately has no dwell, retained-grasp, lift, or trajectory-history + # requirement. Failure predicates are evaluated first and retain same-step precedence. + self.terminations.success.func = mdp.immediate_pour_success + self.terminations.success.params = {} + self.terminations.lost_grasp.params["terminate"] = False + self.terminations.spill.params["terminate"] = True + self.terminations.time_out.func = mdp.unsuccessful_time_out + + def _configure_reward_cfg(self, max_gripper_command: float, *, initialize_stage_gates: bool) -> None: + """Keep the general reward independent of curriculum and grasp thresholds.""" + del max_gripper_command, initialize_stage_gates + + def _validate_reward_cfg(self, stage_count: int) -> None: + """Validate the general reward without introducing reset-stage dependencies.""" + for name in ("reach", "goal_distance"): + std = float(getattr(self.rewards, name).params["std"]) + if not math.isfinite(std) or std <= 0.0: + raise ValueError(f"{name} std must be finite and positive.") + max_velocity = float(self.rewards.joint_velocity.params["max_velocity"]) + if not math.isfinite(max_velocity) or max_velocity <= 0.0: + raise ValueError("joint_velocity max_velocity must be finite and positive.") + + def _validate_curriculum_cfg(self) -> None: + super()._validate_curriculum_cfg() + if not isinstance(self.reset_dataset_path, str): + raise TypeError("reset_dataset_path must be a string.") + if not self.reset_dataset_path: + raise ValueError("reset_dataset_path must be nonempty.") + top_grasp_count = self.reset_dataset_top_grasp_count + if top_grasp_count is not None: + if not isinstance(top_grasp_count, int) or isinstance(top_grasp_count, bool) or top_grasp_count <= 0: + raise ValueError("reset_dataset_top_grasp_count must be a positive integer or None.") + if not self.curriculum_freeze: + raise ValueError("reset_dataset_top_grasp_count requires curriculum_freeze=True.") + expected_hash = self.reset_dataset_content_sha256 + if expected_hash is not None and ( + not isinstance(expected_hash, str) + or len(expected_hash) != 64 + or any(character not in "0123456789abcdef" for character in expected_hash) + ): + raise ValueError("reset_dataset_content_sha256 must be a lowercase SHA-256 or None.") + self.reset_dataset_sampler.validate_values() + + +@configclass +class FrankaPourEnvCfg_RESET_DATASET_EVAL(FrankaPourEnvCfg_RESET_DATASET): + """Frozen full-distribution reset-dataset evaluation.""" + + curriculum_freeze: bool = True + + +@configclass +class FrankaPourEnvCfg_RESET_DATASET_PLAY(FrankaPourEnvCfg_RESET_DATASET_EVAL): + """Reset-dataset playback using the captured sparse multi-world MPM configuration.""" + + def __post_init__(self): + super().__post_init__() + self.scene.num_envs = 1 + self.use_cuda_graph = True + self.sim.physics.use_cuda_graph = True + # Retain the base view direction while moving the camera about one metre closer. + self.viewer.eye = (0.9, 0.65, 0.5) + self.sim.default_visualizer_cfg = VisualizerCfg(eye=self.viewer.eye, lookat=self.viewer.lookat) + + +# Compatibility aliases for checkpoints and commands produced while this task was experimental. +@configclass +class FrankaPourEnvCfg_PLAY(FrankaPourEnvCfg): + """Playback using the training task's captured sparse multi-world MPM configuration.""" + + def __post_init__(self): + super().__post_init__() + self.scene.num_envs = 1 + self.use_cuda_graph = True + self.sim.physics.use_cuda_graph = True + self.curriculum_start_stage = len(self.curriculum_stage_names) - 1 + self.curriculum_randomization_start_level = len(self.curriculum_randomization_extent_levels) - 1 + self.curriculum_freeze = True + + +@configclass +class FrankaPourEnvCfg_TELEOP(FrankaPourEnvCfg_PLAY): + """Teleop preset: 1 env, no RL time-out (operator resets manually).""" + + def __post_init__(self): + super().__post_init__() + # SpaceMouse IK emits direct seven-joint targets; keep that operator-only interface + # separate from the policy's relative-joint action representation. + joint_clip = { + joint_name: limits + for joint_name, limits in zip(self.actions.arm_action.joint_names, PANDA_ARM_JOINT_LIMITS, strict=True) + } + self.actions.arm_action = mdp.CurriculumJointPositionActionCfg( + asset_name="robot", + joint_names=[f"panda_joint{i}" for i in range(1, 8)], + scale=0.5, + alpha=0.2, + project_reference_through_stage=-1, + use_default_offset=True, + preserve_order=True, + clip=joint_clip, + ) + self.actions.gripper_action.force_open_before_phase_stage = -1 + self.actions.gripper_action.limit_to_preload = False + self.actions.gripper_action.neutral_position = self.gripper_open_pos + self.actions.gripper_action.default_position = self.gripper_open_pos + self.actions.gripper_action.scale = self.gripper_open_pos - self.actions.gripper_action.close_position + self.scene.num_envs = 1 + self.terminations.time_out = None + self.episode_length_s = 3600.0 + + def finalize(self) -> FrankaPourEnvCfg_TELEOP: + """Preserve the operator preset's full open-to-close gripper range.""" + resolved = super().finalize() + resolved.actions.gripper_action.default_position = resolved.gripper_open_pos + resolved.actions.gripper_action.scale = ( + resolved.gripper_open_pos - resolved.actions.gripper_action.close_position + ) + return resolved + + def _validate_arm_action_cfg(self) -> None: + """Validate the operator-only absolute joint-position action.""" + arm_action = self.actions.arm_action + # ``configclass`` validates the inherited policy action before this preset's post-init + # replaces it with the operator controller. + if isinstance(arm_action, mdp.RelativeJointPositionActionCfg): + FrankaPourEnvCfg._validate_arm_action_cfg(self) + return + if not isinstance(arm_action, mdp.CurriculumJointPositionActionCfg): + raise ValueError("Franka Pour teleoperation requires CurriculumJointPositionActionCfg for the arm.") + if arm_action.joint_names != [f"panda_joint{i}" for i in range(1, 8)] or not arm_action.preserve_order: + raise ValueError("Franka Pour teleoperation requires all seven Panda arm joints in kinematic order.") + expected_clip = { + joint_name: limits + for joint_name, limits in zip(arm_action.joint_names, PANDA_ARM_JOINT_LIMITS, strict=True) + } + if arm_action.clip != expected_clip: + raise ValueError("Franka Pour teleoperation requires Panda joint-limit clipping.") + return + + +# Deprecated compatibility names; use the corresponding ``RESET_DATASET`` configurations. These +# aliases preserve configuration and task lookup, not compatibility with older 7-action policies. +ResetMixtureRewardsCfg = ResetDatasetRewardsCfg +ResetMixtureCurriculumCfg = ResetDatasetCurriculumCfg +FrankaPourEnvCfg_RESET_MIXTURE = FrankaPourEnvCfg_RESET_DATASET +FrankaPourEnvCfg_RESET_MIXTURE_EVAL = FrankaPourEnvCfg_RESET_DATASET_EVAL +FrankaPourEnvCfg_RESET_MIXTURE_PLAY = FrankaPourEnvCfg_RESET_DATASET_PLAY diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/reset_dataset_generator.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/reset_dataset_generator.py new file mode 100644 index 000000000000..b514f68f8649 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/reset_dataset_generator.py @@ -0,0 +1,1898 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Franka Pour adapter for the reusable reset-dataset generation utilities. + +This module owns proposal generation, Newton IK, rigid collision validation, objective scoring, +and validation of the task-specific direct-state contract used by the runtime reset curriculum. +Generic batching, integrity hashing, atomic persistence, and adaptive runtime sampling live in +``isaaclab_tasks.utils`` so another task only needs to provide its proposals and validators. Particles +are represented by the existing deterministic cup-local fill lattice; replay transforms that one +layout by the cached source-cup pose and starts the MPM solver with zero velocity, stress, and +history. +""" + +from __future__ import annotations + +import math +from collections import defaultdict +from collections.abc import Mapping, Sequence +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import torch + +from isaaclab.utils import math as math_utils + +from isaaclab_tasks.utils.reset_dataset import ( + reset_dataset_collect_batches, + reset_dataset_content_digest, + reset_dataset_digest, + reset_dataset_save_atomic, + reset_dataset_validate_header, +) + +if TYPE_CHECKING: + from .pour_env import FrankaPourEnv + + +FRANKA_POUR_RESET_DATASET_FORMAT = "franka_pour_reset_dataset" +FRANKA_POUR_RESET_DATASET_SCHEMA_VERSION = 6 +FRANKA_POUR_RESET_DATASET_TASK_ID = "Isaac-Pour-Franka-Reset-Dataset-v0" +NON_GRASPING_CATEGORY = 0 +GRASPING_CATEGORY = 1 +RESET_DATASET_GRASPING_COUNT = 10_000 +RESET_DATASET_NON_GRASPING_COUNT = 10_000 +RESET_DATASET_NEAR_POUR_COUNT = 1_000 + +_ARM_JOINT_NAMES = tuple(f"panda_joint{index}" for index in range(1, 8)) +_FINGER_JOINT_NAMES = ("panda_finger_joint1", "panda_finger_joint2") +_STATE_KEYS = ( + "arm_joint_position", + "arm_joint_velocity", + "finger_joint_position", + "finger_joint_velocity", + "finger_joint_target", + "source_root_pose", + "source_root_velocity", + "target_root_pose", + "target_root_velocity", + "category", + "objective", + "objective_raw", + "objective_components", + "grasp_region", + "grasp_side", + "attempt_id", + "particle_layout_id", + "ik_cost", + "ik_position_residual", + "ik_rotation_residual", +) + + +@dataclass(frozen=True) +class FrankaPourResetDatasetGeneratorCfg: + """Configuration for one statically valid reset-dataset candidate pool.""" + + grasping_count: int = RESET_DATASET_GRASPING_COUNT + non_grasping_count: int = RESET_DATASET_NON_GRASPING_COUNT + batch_size: int = 256 + seed: int = 42 + max_attempt_multiplier: int = 100 + + # Reserve an explicit, side-balanced near-goal stratum. These states are generated target-first + # and still pass the same Newton IK, self-collision, obstacle, and support-surface rejection as + # every broad grasping state. + near_pour_grasp_count: int = RESET_DATASET_NEAR_POUR_COUNT + near_pour_horizontal_radius: float = 0.02 + near_pour_height_range: tuple[float, float] = (0.15, 0.25) + near_pour_tilt_angle_range: tuple[float, float] = (math.radians(120.0), math.radians(170.0)) + + # A conservative cylindrical envelope around the Panda base. Sampling the central 90% of + # this envelope and rejecting through Newton IK defines the usable kinematic workspace without + # relying on a hand-authored finite pose bank. + workspace_central_fraction: float = 0.90 + workspace_radius_range: tuple[float, float] = (0.20, 0.82) + workspace_azimuth_range: tuple[float, float] = (-math.pi, math.pi) + workspace_height_range: tuple[float, float] = (0.08, 0.82) + + ik_seeds: int = 64 + ik_iterations: int = 160 + ik_noise_std: float = 0.75 + ik_max_cost: float = 1.0e-3 + ik_joint_margin: float = 0.015 + ik_max_position_residual: float = 0.003 + ik_max_rotation_residual: float = math.radians(3.0) + ik_max_home_distance: float = 6.0 + + # TCP-local +Y is the jaw axis. Its very small standard deviation preserves a genuine + # Gaussian without initializing the exact-width cup deeply inside either finger. + grasp_position_std: tuple[float, float, float] = (0.0015, 0.00005, 0.0015) + grasp_seating_max_offset: tuple[float, float, float] = (0.006, 0.00025, 0.006) + grasp_seating_max_rotation_error: float = math.radians(1.0) + non_grasping_min_tcp_source_distance: float = 0.12 + collision_penetration_tolerance: float = 1.0e-4 + finger_contact_penetration_tolerance: float = 0.00025 + obstacle_clearance: float = 0.001 + + objective_distance_threshold: float = 0.15 + objective_target_horizontal_threshold: float = 0.15 + objective_target_height_threshold: float = 0.15 + objective_inversion_gate_horizontal_threshold: float = 0.07 + objective_weights: tuple[float, float, float] = (1.0 / 3.0, 1.0 / 3.0, 1.0 / 3.0) + + def __post_init__(self) -> None: + """Reject invalid sampling contracts before allocating simulation state.""" + for field_name in ("grasping_count", "non_grasping_count", "batch_size", "max_attempt_multiplier"): + value = getattr(self, field_name) + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"{field_name} must be a positive integer.") + if ( + not isinstance(self.near_pour_grasp_count, int) + or isinstance(self.near_pour_grasp_count, bool) + or not 0 <= self.near_pour_grasp_count <= self.grasping_count + ): + raise ValueError("near_pour_grasp_count must be an integer in [0, grasping_count].") + if not 0.0 < self.workspace_central_fraction <= 1.0: + raise ValueError("workspace_central_fraction must lie in (0, 1].") + for field_name in ("workspace_radius_range", "workspace_azimuth_range", "workspace_height_range"): + lower, upper = getattr(self, field_name) + if not math.isfinite(lower) or not math.isfinite(upper) or lower >= upper: + raise ValueError(f"{field_name} must contain two finite increasing values.") + for field_name in ( + "ik_noise_std", + "ik_max_cost", + "ik_joint_margin", + "ik_max_position_residual", + "ik_max_rotation_residual", + "ik_max_home_distance", + "non_grasping_min_tcp_source_distance", + "grasp_seating_max_rotation_error", + "finger_contact_penetration_tolerance", + "near_pour_horizontal_radius", + "objective_distance_threshold", + "objective_target_horizontal_threshold", + "objective_target_height_threshold", + "objective_inversion_gate_horizontal_threshold", + ): + value = getattr(self, field_name) + if not math.isfinite(value) or value <= 0.0: + raise ValueError(f"{field_name} must be finite and positive.") + if self.collision_penetration_tolerance < 0.0 or self.obstacle_clearance < 0.0: + raise ValueError("Collision tolerances must be nonnegative.") + for field_name in ("near_pour_height_range", "near_pour_tilt_angle_range"): + lower, upper = getattr(self, field_name) + if not math.isfinite(lower) or not math.isfinite(upper) or lower <= 0.0 or lower >= upper: + raise ValueError(f"{field_name} must contain two finite increasing positive values.") + if self.near_pour_tilt_angle_range[1] >= math.pi: + raise ValueError("near_pour_tilt_angle_range must remain below pi radians.") + if any(not math.isfinite(value) or value < 0.0 for value in self.grasp_position_std): + raise ValueError("grasp_position_std must contain three finite nonnegative values.") + if any(not math.isfinite(value) or value <= 0.0 for value in self.grasp_seating_max_offset): + raise ValueError("grasp_seating_max_offset must contain three finite positive values.") + if any(not math.isfinite(value) or value < 0.0 for value in self.objective_weights): + raise ValueError("objective_weights must be finite and nonnegative.") + if not math.isclose(sum(self.objective_weights), 1.0, rel_tol=0.0, abs_tol=1.0e-6): + raise ValueError("objective_weights must sum to one.") + if self.objective_inversion_gate_horizontal_threshold > self.objective_target_horizontal_threshold: + raise ValueError("The inversion gate cannot be wider than the target-alignment kernel.") + + +def grasp_objective_components( + source_pose: torch.Tensor, + target_pose: torch.Tensor, + *, + source_region_center: torch.Tensor | Sequence[float], + cup_center_offset: torch.Tensor | Sequence[float], + target_rim_height: float, + distance_threshold: float = 0.15, + target_horizontal_threshold: float = 0.15, + target_height_threshold: float = 0.15, + inversion_gate_horizontal_threshold: float | None = None, +) -> torch.Tensor: + """Return distance, inversion, and target-alignment scores in ``[0, 1]``. + + Poses use environment-frame position followed by an XYZW quaternion. The source asset root + lies at the cup base, so ``cup_center_offset`` is rotated into the world before scoring. + """ + if source_pose.ndim != 2 or source_pose.shape[1] != 7: + raise ValueError(f"source_pose must have shape (N, 7), got {tuple(source_pose.shape)}.") + if target_pose.shape != source_pose.shape: + raise ValueError("target_pose must have the same (N, 7) shape as source_pose.") + source_center = torch.as_tensor(source_region_center, device=source_pose.device, dtype=source_pose.dtype) + center_offset_value = torch.as_tensor(cup_center_offset, device=source_pose.device, dtype=source_pose.dtype) + if source_center.shape != (3,) or center_offset_value.shape != (3,): + raise ValueError("source_region_center and cup_center_offset must each have shape (3,).") + if inversion_gate_horizontal_threshold is None: + inversion_gate_horizontal_threshold = target_horizontal_threshold + thresholds = ( + distance_threshold, + target_horizontal_threshold, + target_height_threshold, + inversion_gate_horizontal_threshold, + ) + if any(not math.isfinite(value) or value <= 0.0 for value in thresholds): + raise ValueError("Objective thresholds must be finite and positive.") + + count = source_pose.shape[0] + center_offset = center_offset_value.expand(count, -1) + cup_center = source_pose[:, :3] + math_utils.quat_apply(source_pose[:, 3:7], center_offset) + distance_score = torch.linalg.vector_norm(cup_center - source_center, dim=-1).div(distance_threshold).clamp_(0, 1) + + local_up = torch.zeros((count, 3), device=source_pose.device, dtype=source_pose.dtype) + local_up[:, 2] = 1.0 + world_up = math_utils.quat_apply(source_pose[:, 3:7], local_up) + inversion_score = ((1.0 - world_up[:, 2]) * 0.5).clamp_(0, 1) + + horizontal_distance = torch.linalg.vector_norm(cup_center[:, :2] - target_pose[:, :2], dim=-1) + horizontal_score = (1.0 - horizontal_distance / target_horizontal_threshold).clamp_(0, 1) + target_rim_z = target_pose[:, 2] + float(target_rim_height) + height_score = ((cup_center[:, 2] - target_rim_z) / target_height_threshold).clamp_(0, 1) + target_score = horizontal_score * height_score + # Inverting the cup away from the receiver is destructive, not task progress. Couple tilt + # credit to the complete receiver-alignment score so it vanishes unless the cup is both above + # and horizontally aligned with the bowl. + centered_above_target = horizontal_distance <= inversion_gate_horizontal_threshold + target_gated_inversion_score = inversion_score * target_score * centered_above_target + return torch.stack((distance_score, target_gated_inversion_score, target_score), dim=-1) + + +def source_root_position_from_tcp_grasp( + tcp_position: torch.Tensor, + tcp_quaternion: torch.Tensor, + source_quaternion: torch.Tensor, + grasp_offset_source: torch.Tensor, + seating_offset_tcp: torch.Tensor, +) -> torch.Tensor: + """Place a source root so its configured grasp point is seated at the tool centre. + + ``seating_offset_tcp`` is the sampled Gaussian displacement of the cup grasp point from the + tool centre, expressed in the TCP frame. Keeping this transform explicit prevents confusing + the cup's geometric centre with the point along the fingers that is actually grasped. + """ + expected_shape = tcp_position.shape + if expected_shape[-1:] != (3,) or source_quaternion.shape != (*expected_shape[:-1], 4): + raise ValueError("TCP positions and source quaternions must have matching (..., 3/4) shapes.") + if tcp_quaternion.shape != source_quaternion.shape or seating_offset_tcp.shape != expected_shape: + raise ValueError("TCP quaternions and seating offsets must match the source batch shape.") + grasp_offset_source = torch.as_tensor( + grasp_offset_source, + device=tcp_position.device, + dtype=tcp_position.dtype, + ) + if grasp_offset_source.shape != (3,): + raise ValueError("grasp_offset_source must have shape (3,).") + grasp_position = tcp_position + math_utils.quat_apply(tcp_quaternion, seating_offset_tcp) + return grasp_position - math_utils.quat_apply( + source_quaternion, + grasp_offset_source.expand_as(tcp_position), + ) + + +def above_target_tilted_mask( + source_pose: torch.Tensor, + target_pose: torch.Tensor, + *, + cup_center_offset: torch.Tensor | Sequence[float], + target_rim_height: float, + max_horizontal_distance: float, + min_vertical_clearance: float, + min_tilt_angle: float, +) -> torch.Tensor: + """Return states whose cup center is above the receiver and whose cup is sufficiently tilted.""" + if source_pose.ndim != 2 or source_pose.shape[1] != 7 or target_pose.shape != source_pose.shape: + raise ValueError("source_pose and target_pose must have matching shape (N, 7).") + thresholds = (target_rim_height, max_horizontal_distance, min_vertical_clearance, min_tilt_angle) + if any(not math.isfinite(value) or value <= 0.0 for value in thresholds): + raise ValueError("Near-pour geometry thresholds must be finite and positive.") + if min_tilt_angle >= math.pi: + raise ValueError("min_tilt_angle must remain below pi radians.") + count = source_pose.shape[0] + center_offset = torch.as_tensor(cup_center_offset, device=source_pose.device, dtype=source_pose.dtype) + if center_offset.shape != (3,): + raise ValueError("cup_center_offset must have shape (3,).") + cup_center = source_pose[:, :3] + math_utils.quat_apply(source_pose[:, 3:7], center_offset.expand(count, -1)) + horizontal_distance = torch.linalg.vector_norm(cup_center[:, :2] - target_pose[:, :2], dim=-1) + vertical_clearance = cup_center[:, 2] - (target_pose[:, 2] + float(target_rim_height)) + local_up = torch.zeros((count, 3), device=source_pose.device, dtype=source_pose.dtype) + local_up[:, 2] = 1.0 + world_up = math_utils.quat_apply(source_pose[:, 3:7], local_up) + return ( + (horizontal_distance <= max_horizontal_distance) + & (vertical_clearance >= min_vertical_clearance) + & (world_up[:, 2] <= math.cos(min_tilt_angle)) + ) + + +def oriented_box_supported_by_bounds( + pose: torch.Tensor, + half_extents: torch.Tensor | Sequence[float], + support_lower_xy: torch.Tensor | Sequence[float], + support_upper_xy: torch.Tensor | Sequence[float], + *, + clearance: float = 0.0, +) -> torch.Tensor: + """Return whether each oriented box footprint lies completely on a rectangular support.""" + if pose.ndim != 2 or pose.shape[1] != 7: + raise ValueError("pose must have shape (N, 7).") + if not math.isfinite(clearance) or clearance < 0.0: + raise ValueError("clearance must be finite and nonnegative.") + half = torch.as_tensor(half_extents, device=pose.device, dtype=pose.dtype) + lower = torch.as_tensor(support_lower_xy, device=pose.device, dtype=pose.dtype) + upper = torch.as_tensor(support_upper_xy, device=pose.device, dtype=pose.dtype) + if half.shape != (3,) or lower.shape != (2,) or upper.shape != (2,): + raise ValueError("half_extents, support_lower_xy, and support_upper_xy must have shapes (3,), (2,), (2,).") + if bool(torch.any(half <= 0.0)) or bool(torch.any(lower >= upper)): + raise ValueError("Support bounds and half extents must define positive regions.") + center_offset = torch.zeros((pose.shape[0], 3), device=pose.device, dtype=pose.dtype) + center_offset[:, 2] = half[2] + center = pose[:, :3] + math_utils.quat_apply(pose[:, 3:7], center_offset) + rotation = math_utils.matrix_from_quat(pose[:, 3:7]).abs() + planar_radius = (rotation[:, :2, :] * half).sum(dim=-1) + clearance + return ((center[:, :2] - planar_radius >= lower) & (center[:, :2] + planar_radius <= upper)).all(dim=-1) + + +def normalize_grasp_objectives(raw_objective: torch.Tensor) -> torch.Tensor: + """Normalize one non-degenerate vector so its exact extrema are zero and one.""" + if raw_objective.ndim != 1 or raw_objective.numel() < 2: + raise ValueError("raw_objective must be a one-dimensional tensor with at least two entries.") + if not bool(torch.isfinite(raw_objective).all()): + raise ValueError("raw_objective must contain only finite values.") + minimum = raw_objective.min() + span = raw_objective.max() - minimum + if float(span) <= torch.finfo(raw_objective.dtype).eps: + raise ValueError("Cannot normalize degenerate grasp objectives with zero range.") + normalized = (raw_objective - minimum) / span + # Avoid roundoff obscuring the cache contract's exact endpoints. + normalized[torch.argmin(raw_objective)] = 0.0 + normalized[torch.argmax(raw_objective)] = 1.0 + return normalized + + +def oriented_boxes_overlap( + center_a: torch.Tensor, + quaternion_a: torch.Tensor, + half_extents_a: Sequence[float], + center_b: torch.Tensor, + quaternion_b: torch.Tensor, + half_extents_b: Sequence[float], + *, + clearance: float = 0.0, +) -> torch.Tensor: + """Return batched OBB intersection using the complete 15-axis separating-axis test.""" + if center_a.ndim != 2 or center_a.shape[1] != 3 or center_b.shape != center_a.shape: + raise ValueError("Both center tensors must have shape (N, 3).") + count = center_a.shape[0] + if quaternion_a.shape != (count, 4) or quaternion_b.shape != (count, 4): + raise ValueError("Both quaternion tensors must have shape (N, 4).") + if clearance < 0.0: + raise ValueError("clearance must be nonnegative.") + + rotation_a = math_utils.matrix_from_quat(quaternion_a) + rotation_b = math_utils.matrix_from_quat(quaternion_b) + relative_rotation = rotation_a.transpose(-1, -2) @ rotation_b + absolute_rotation = relative_rotation.abs() + 1.0e-6 + translation = (rotation_a.transpose(-1, -2) @ (center_b - center_a).unsqueeze(-1)).squeeze(-1) + half_a = torch.as_tensor(half_extents_a, device=center_a.device, dtype=center_a.dtype) + clearance * 0.5 + half_b = torch.as_tensor(half_extents_b, device=center_a.device, dtype=center_a.dtype) + clearance * 0.5 + + separated = torch.zeros(count, device=center_a.device, dtype=torch.bool) + for axis in range(3): + radius_b = (absolute_rotation[:, axis, :] * half_b).sum(dim=-1) + separated |= translation[:, axis].abs() > half_a[axis] + radius_b + for axis in range(3): + projection = (translation * relative_rotation[:, :, axis]).sum(dim=-1).abs() + radius_a = (absolute_rotation[:, :, axis] * half_a).sum(dim=-1) + separated |= projection > radius_a + half_b[axis] + for axis_a in range(3): + for axis_b in range(3): + other_a = (axis_a + 1) % 3 + last_a = (axis_a + 2) % 3 + other_b = (axis_b + 1) % 3 + last_b = (axis_b + 2) % 3 + projection = ( + translation[:, last_a] * relative_rotation[:, other_a, axis_b] + - translation[:, other_a] * relative_rotation[:, last_a, axis_b] + ).abs() + radius_a = ( + half_a[other_a] * absolute_rotation[:, last_a, axis_b] + + half_a[last_a] * absolute_rotation[:, other_a, axis_b] + ) + radius_b = ( + half_b[other_b] * absolute_rotation[:, axis_a, last_b] + + half_b[last_b] * absolute_rotation[:, axis_a, other_b] + ) + separated |= projection > radius_a + radius_b + return ~separated + + +def _derive_tabletop_support_bounds( + env: FrankaPourEnv, + source_env_path: str | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Derive the transformed SeattleLab tabletop footprint from the live USD stage.""" + from pxr import Usd, UsdGeom + + import isaaclab.sim as sim_utils + from isaaclab.cloner import resolve_clone_plan_source + + if source_env_path is None: + plan = sim_utils.SimulationContext.instance().get_clone_plan() + resolved = resolve_clone_plan_source(env._robot.cfg.prim_path, plan) if plan is not None else None + if resolved is None: + raise RuntimeError(f"Could not resolve clone-plan source for {env._robot.cfg.prim_path!r}.") + source_env_path = resolved[0] + + stage = sim_utils.get_current_stage() + collision_path = f"{source_env_path.rstrip('/')}/Table/Collisions/Cube" + collision_prim = stage.GetPrimAtPath(collision_path) + if not collision_prim.IsValid(): + raise RuntimeError(f"Could not find the SeattleLab tabletop collision prim at {collision_path!r}.") + bbox_cache = UsdGeom.BBoxCache( + Usd.TimeCode.Default(), + [UsdGeom.Tokens.default_, UsdGeom.Tokens.render, UsdGeom.Tokens.proxy], + useExtentsHint=False, + ) + world_range = bbox_cache.ComputeWorldBound(collision_prim).ComputeAlignedRange() + lower_w = torch.tensor(tuple(world_range.GetMin()), device=env.device, dtype=torch.float32) + upper_w = torch.tensor(tuple(world_range.GetMax()), device=env.device, dtype=torch.float32) + lower = lower_w[:2] - env.env_origins[0, :2] + upper = upper_w[:2] - env.env_origins[0, :2] + if not bool(torch.isfinite(lower).all() & torch.isfinite(upper).all() & torch.all(lower < upper)): + raise RuntimeError("Could not derive finite positive SeattleLab tabletop support bounds.") + return lower, upper + + +def build_franka_pour_reset_task_contract(env: FrankaPourEnv) -> dict[str, Any]: + """Build the canonical physics, geometry, and frame contract for reset-dataset compatibility.""" + import newton + import warp as wp + + task_cfg = env.cfg + source_half = torch.as_tensor(task_cfg.cup_grasp_box_half, dtype=torch.float32) + target_half = torch.tensor( + ( + task_cfg.target_cup_inner_width * 0.5 + task_cfg.target_cup_wall_thickness, + task_cfg.target_cup_inner_depth * 0.5 + task_cfg.target_cup_wall_thickness, + (task_cfg.target_cup_cavity_depth + task_cfg.target_cup_bottom_thickness) * 0.5, + ), + dtype=torch.float32, + ) + cup_center_offset = torch.tensor((0.0, 0.0, float(source_half[2])), dtype=torch.float32) + source_region_center = torch.as_tensor(task_cfg.cup_reset_pos, dtype=torch.float32) + cup_center_offset + arm_limits = env._joint_pos_limits_t[0, env._arm_joint_ids].detach().cpu() + finger_limits = env._joint_pos_limits_t[0, env._finger_joint_ids].detach().cpu() + finger_position_range = (float(finger_limits[:, 0].max()), float(finger_limits[:, 1].min())) + tabletop_lower, tabletop_upper = _derive_tabletop_support_bounds(env) + material = task_cfg.media_material + particle_spacing = float(task_cfg.voxel_size) / float(task_cfg.particles_per_cell) + robot_spawn = task_cfg.scene.robot.spawn + table_spawn = task_cfg.scene.table.spawn + return { + "robot_asset": str(getattr(robot_spawn, "usd_path", type(robot_spawn).__name__)), + "table_asset": str(getattr(table_spawn, "usd_path", type(table_spawn).__name__)), + "newton_version": str(getattr(newton, "__version__", "unknown")), + "warp_version": str(getattr(wp, "__version__", "unknown")), + "source_box_half": tuple(float(value) for value in source_half), + "target_box_half": tuple(float(value) for value in target_half), + "target_rim_height": float(target_half[2] * 2.0), + "source_mesh_vertices_sha256": reset_dataset_digest(torch.as_tensor(env._cup_vertices)), + "source_mesh_indices_sha256": reset_dataset_digest(torch.as_tensor(env._cup_indices)), + "target_mesh_vertices_sha256": reset_dataset_digest(torch.as_tensor(env._target_vertices)), + "target_mesh_indices_sha256": reset_dataset_digest(torch.as_tensor(env._target_indices)), + "tabletop_support_lower_xy": tuple(float(value) for value in tabletop_lower), + "tabletop_support_upper_xy": tuple(float(value) for value in tabletop_upper), + "cup_grasp_tcp_quat_c": tuple(float(value) for value in task_cfg.cup_grasp_tcp_quat_c), + "cup_grasp_height": float(task_cfg.cup_grasp_height), + "source_region_center": tuple(float(value) for value in source_region_center), + "source_radius_range": task_cfg.curriculum_randomized_source_radius_range, + "source_position_range": tuple(task_cfg.curriculum_randomized_source_position_range), + "source_azimuth_range": float(task_cfg.curriculum_randomized_source_azimuth_range), + "target_center_xy": tuple(task_cfg.curriculum_randomized_target_center_xy), + "target_position_range": tuple(task_cfg.curriculum_randomized_target_position_range), + "tcp_body_name": task_cfg.tcp_body_name, + "tcp_offset_pos": tuple(float(value) for value in task_cfg.tcp_offset_pos), + "tcp_offset_rot": tuple(float(value) for value in task_cfg.tcp_offset_rot), + "arm_home": tuple(float(value) for value in task_cfg.arm_home), + "arm_joint_limits": arm_limits, + "gripper_position_range": finger_position_range, + "gripper_open_pos": float(task_cfg.gripper_open_pos), + "gripper_preload_pos": float(task_cfg.gripper_preload_pos), + "gripper_grasp_reset_target": float(task_cfg.actions.gripper_action.close_position), + "gripper_contact_min_deflection": float(task_cfg.actions.gripper_action.contact_min_deflection), + "cup_mass": float(task_cfg.cup_mass), + "source_cup_friction": float(task_cfg.source_cup_friction), + "target_cup_friction": float(task_cfg.target_cup_friction), + "cup_grasp_box_friction": float(task_cfg.cup_grasp_box_friction), + "grasp_contact_ke": float(task_cfg.grasp_contact_ke), + "grasp_contact_kd": float(task_cfg.grasp_contact_kd), + "grasp_contact_kf": float(task_cfg.grasp_contact_kf), + "collider_margin": float(task_cfg.collider_margin), + "simulation_dt": float(task_cfg.sim.dt), + "gravity": tuple(float(value) for value in task_cfg.sim.gravity), + "policy_decimation": int(task_cfg.decimation), + "physics_substeps": int(task_cfg.physics_substeps), + "rigid_entry_substeps": int(task_cfg.rigid_entry_substeps), + "mpm_entry_substeps": int(task_cfg.mpm_entry_substeps), + "mpm_iterations": int(task_cfg.mpm_iterations), + "proxy_iterations": int(task_cfg.proxy_iterations), + "proxy_mass_scale": float(task_cfg.proxy_mass_scale), + "particle_workspace_lower_bound": tuple(task_cfg.particle_workspace_lower_bound), + "particle_workspace_upper_bound": tuple(task_cfg.particle_workspace_upper_bound), + "particle_max_velocity": float(task_cfg.particle_max_velocity), + "particle_count": int(env._media_local_points_t.shape[0]), + "particle_spacing": particle_spacing, + "particle_mass": particle_spacing**3 * float(material.density), + "particle_radius": 0.5 * particle_spacing, + "media_fill_fraction": float(task_cfg.media_fill_frac), + "media_material": { + name: float(getattr(material, name)) + for name in ( + "density", + "young_modulus", + "poisson_ratio", + "viscosity", + "friction", + "damping", + "yield_pressure", + "tensile_yield_ratio", + "yield_stress", + "hardening", + "dilatancy", + ) + }, + "particle_layout_sha256": reset_dataset_digest(env._media_local_points_t), + } + + +def reset_dataset_content_sha256(payload: Mapping[str, Any]) -> str: + """Return the content hash while excluding the hash field itself.""" + return reset_dataset_content_digest(payload) + + +def build_reset_dataset_payload( + states: Mapping[str, torch.Tensor], + particle_local_positions: torch.Tensor, + metadata: Mapping[str, Any], + cfg: FrankaPourResetDatasetGeneratorCfg, +) -> dict[str, Any]: + """Build, hash, and validate a CPU-only direct-state cache payload.""" + missing = sorted(set(_STATE_KEYS) - set(states)) + extra = sorted(set(states) - set(_STATE_KEYS)) + if missing or extra: + raise ValueError(f"State tensor keys differ from the schema: missing={missing}, extra={extra}.") + cpu_states = {key: value.detach().cpu().contiguous() for key, value in states.items()} + local_positions = particle_local_positions.detach().cpu().to(dtype=torch.float32).contiguous() + if local_positions.ndim == 3 and local_positions.shape[0] == 1: + local_positions = local_positions[0] + if local_positions.ndim != 2 or local_positions.shape[1] != 3 or local_positions.shape[0] == 0: + raise ValueError("particle_local_positions must have shape (P, 3) or (1, P, 3), with P > 0.") + particle_layouts = { + "local_position": local_positions.unsqueeze(0), + "local_velocity": torch.zeros_like(local_positions).unsqueeze(0), + } + sampler_cfg = asdict(cfg) + task_contract = metadata.get("task_contract", {}) + contract_sha256 = reset_dataset_digest({"sampler_cfg": sampler_cfg, "task_contract": task_contract}) + payload: dict[str, Any] = { + "schema_version": FRANKA_POUR_RESET_DATASET_SCHEMA_VERSION, + "format": FRANKA_POUR_RESET_DATASET_FORMAT, + "contract_sha256": contract_sha256, + "metadata": { + **dict(metadata), + "sampler_cfg": sampler_cfg, + "state_count": cfg.grasping_count + cfg.non_grasping_count, + "category_names": ("non_grasping", "grasping"), + "category_counts": torch.tensor((cfg.non_grasping_count, cfg.grasping_count), dtype=torch.int64), + "joint_names": _ARM_JOINT_NAMES + _FINGER_JOINT_NAMES, + "frame": "environment", + "quaternion_order": "xyzw", + "particle_solver_state": "fresh_zero", + }, + "states": cpu_states, + "particle_layouts": particle_layouts, + } + payload["content_sha256"] = reset_dataset_content_sha256(payload) + validate_reset_dataset( + payload, + expected_grasping_count=cfg.grasping_count, + expected_non_grasping_count=cfg.non_grasping_count, + ) + return payload + + +def validate_reset_dataset( + payload: Mapping[str, Any], + *, + expected_grasping_count: int | None = None, + expected_non_grasping_count: int | None = None, + expected_task_contract: Mapping[str, Any] | None = None, +) -> None: + """Validate schema, state invariants, category quotas, and the complete content hash.""" + metadata, sampler_cfg, task_contract, states = _validate_cache_header(payload) + state_count = _validate_state_tensor_schema(states) + grasping, non_grasping = _validate_category_counts( + states, + metadata, + sampler_cfg, + state_count, + expected_grasping_count, + expected_non_grasping_count, + ) + _validate_objective_values(states, metadata, grasping, non_grasping) + _validate_state_invariants(states, sampler_cfg, task_contract, grasping, non_grasping) + _validate_particle_layouts(payload) + if expected_task_contract is not None: + stored_contract_digest = reset_dataset_digest(task_contract) + expected_contract_digest = reset_dataset_digest(expected_task_contract) + if stored_contract_digest != expected_contract_digest: + raise ValueError( + "Reset dataset task contract does not match the current environment: " + f"{stored_contract_digest} != {expected_contract_digest}. Regenerate candidates " + "with the same task configuration used for validation and training." + ) + expected_hash = payload.get("content_sha256") + if not isinstance(expected_hash, str) or expected_hash != reset_dataset_content_sha256(payload): + raise ValueError("Reset dataset content hash does not match its payload.") + + +def validate_production_reset_dataset( + payload: Mapping[str, Any], + *, + expected_grasping_count: int = RESET_DATASET_GRASPING_COUNT, + expected_non_grasping_count: int = RESET_DATASET_NON_GRASPING_COUNT, + expected_task_contract: Mapping[str, Any] | None = None, +) -> None: + """Validate a reset dataset and its dynamic-simulation provenance. + + Candidate datasets intentionally satisfy the state schema so the offline validator can replay + them. Production training additionally requires the provenance marker written only after every + retained row has passed dynamic validation in the real task. + + Args: + payload: Loaded Franka Pour reset-dataset payload. + expected_grasping_count: Required number of grasping rows. + expected_non_grasping_count: Required number of non-grasping rows. + expected_task_contract: Optional canonical current-task contract to compare in full. + + Raises: + ValueError: If the payload is not a valid reset dataset or lacks valid dynamic-validation + provenance. + """ + validate_reset_dataset( + payload, + expected_grasping_count=expected_grasping_count, + expected_non_grasping_count=expected_non_grasping_count, + expected_task_contract=expected_task_contract, + ) + metadata = payload["metadata"] + marker = metadata.get("dynamic_validation") + if not isinstance(marker, Mapping): + raise ValueError( + "Production reset datasets require a dynamic_validation metadata marker; " + "run validate_franka_pour_reset_dataset.py on the candidate dataset first." + ) + + source_hash = marker.get("source_content_sha256") + if ( + not isinstance(source_hash, str) + or len(source_hash) != 64 + or any(character not in "0123456789abcdef" for character in source_hash) + ): + raise ValueError("dynamic_validation.source_content_sha256 must be a lowercase SHA-256.") + if source_hash == payload["content_sha256"]: + raise ValueError("Dynamic-validation source and production content hashes must differ.") + + integer_fields = { + "steps": 1, + "settle_steps": 0, + "failure_dwell_steps": 1, + "balance_trimmed": 0, + } + for name, minimum in integer_fields.items(): + value = marker.get(name) + if not isinstance(value, int) or isinstance(value, bool) or value < minimum: + raise ValueError(f"dynamic_validation.{name} must be an integer >= {minimum}.") + if marker["settle_steps"] >= marker["steps"]: + raise ValueError("dynamic_validation.settle_steps must be smaller than steps.") + + failure_counts = marker.get("failure_counts") + if not isinstance(failure_counts, Mapping) or not failure_counts: + raise ValueError("dynamic_validation.failure_counts must be a nonempty mapping.") + for name, count in failure_counts.items(): + if not isinstance(name, str) or not name: + raise ValueError("dynamic_validation.failure_counts keys must be nonempty strings.") + if not isinstance(count, int) or isinstance(count, bool) or count < 0: + raise ValueError("dynamic_validation.failure_counts values must be nonnegative integers.") + + +def select_production_reset_rows( + states: Mapping[str, torch.Tensor], + dynamically_valid: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Select the exact balanced production quotas from an oversampled validated pool.""" + category = states["category"] + region = states["grasp_region"] + side = states["grasp_side"] + objective = states["objective"] + if dynamically_valid.shape != category.shape or dynamically_valid.dtype != torch.bool: + raise ValueError("dynamically_valid must be a Boolean vector aligned with reset rows.") + + keep = torch.zeros_like(dynamically_valid) + non_grasping_rows = torch.nonzero( + dynamically_valid & (category == NON_GRASPING_CATEGORY), + as_tuple=False, + ).flatten() + if non_grasping_rows.numel() < RESET_DATASET_NON_GRASPING_COUNT: + raise RuntimeError( + "Dynamic validation retained only " + f"{non_grasping_rows.numel()}/{RESET_DATASET_NON_GRASPING_COUNT} required non-grasping states. " + "Generate a larger candidate pool or improve proposal validity." + ) + keep[non_grasping_rows[:RESET_DATASET_NON_GRASPING_COUNT]] = True + + required_per_side = { + 0: (RESET_DATASET_GRASPING_COUNT - RESET_DATASET_NEAR_POUR_COUNT) // 4, + 1: RESET_DATASET_NEAR_POUR_COUNT // 4, + } + for region_id, required in required_per_side.items(): + for side_id in range(4): + rows = torch.nonzero( + dynamically_valid & (category == GRASPING_CATEGORY) & (region == region_id) & (side == side_id), + as_tuple=False, + ).flatten() + rows = rows[torch.argsort(objective[rows], descending=True, stable=True)] + if rows.numel() < required: + label = "near-pour" if region_id == 1 else "broad" + raise RuntimeError( + f"Dynamic validation retained only {rows.numel()}/{required} required " + f"{label} grasping states for side {side_id}. Generate a larger candidate " + "pool or improve proposal validity." + ) + keep[rows[:required]] = True + + if int(keep.sum()) != RESET_DATASET_GRASPING_COUNT + RESET_DATASET_NON_GRASPING_COUNT: + raise RuntimeError("Balanced validation did not produce the exact 20,000-state production quota.") + return keep, dynamically_valid & ~keep + + +def _validate_cache_header( + payload: Mapping[str, Any], +) -> tuple[Mapping[str, Any], Mapping[str, Any], Mapping[str, Any], Mapping[str, torch.Tensor]]: + """Validate and return the typed top-level cache mappings.""" + metadata = payload.get("metadata") + if not isinstance(metadata, Mapping): + raise ValueError("Reset-dataset metadata must be a mapping.") + sampler_cfg = metadata.get("sampler_cfg") + task_contract = metadata.get("task_contract", {}) + if not isinstance(sampler_cfg, Mapping) or not isinstance(task_contract, Mapping): + raise ValueError("Reset-dataset sampling and task contracts must be mappings.") + metadata, states = reset_dataset_validate_header( + payload, + expected_format=FRANKA_POUR_RESET_DATASET_FORMAT, + expected_schema_version=FRANKA_POUR_RESET_DATASET_SCHEMA_VERSION, + expected_contract={"sampler_cfg": sampler_cfg, "task_contract": task_contract}, + ) + if not isinstance(states, Mapping) or set(states) != set(_STATE_KEYS): + raise ValueError("Reset dataset has invalid state tensor keys.") + if not all(isinstance(value, torch.Tensor) for value in states.values()): + raise TypeError("Every reset-dataset state field must be a tensor.") + return metadata, sampler_cfg, task_contract, states + + +def _validate_state_tensor_schema(states: Mapping[str, torch.Tensor]) -> int: + """Validate all tensor shapes, scalar dtypes, and finite floating-point values.""" + state_count = int(states["category"].numel()) + expected_shapes = { + "arm_joint_position": (state_count, 7), + "arm_joint_velocity": (state_count, 7), + "finger_joint_position": (state_count, 2), + "finger_joint_velocity": (state_count, 2), + "finger_joint_target": (state_count, 2), + "source_root_pose": (state_count, 7), + "source_root_velocity": (state_count, 6), + "target_root_pose": (state_count, 7), + "target_root_velocity": (state_count, 6), + "category": (state_count,), + "objective": (state_count,), + "objective_raw": (state_count,), + "objective_components": (state_count, 3), + "grasp_region": (state_count,), + "grasp_side": (state_count,), + "attempt_id": (state_count,), + "particle_layout_id": (state_count,), + "ik_cost": (state_count,), + "ik_position_residual": (state_count,), + "ik_rotation_residual": (state_count,), + } + for key, shape in expected_shapes.items(): + if tuple(states[key].shape) != shape: + raise ValueError(f"State field {key!r} must have shape {shape}, got {tuple(states[key].shape)}.") + expected_dtypes = { + "category": torch.int8, + "grasp_region": torch.int8, + "grasp_side": torch.int8, + "attempt_id": torch.int64, + "particle_layout_id": torch.int32, + } + for key, dtype in expected_dtypes.items(): + if states[key].dtype != dtype: + raise ValueError(f"State field {key!r} must use dtype {dtype}, got {states[key].dtype}.") + floating_keys = [key for key in _STATE_KEYS if states[key].is_floating_point()] + if any(not bool(torch.isfinite(states[key]).all()) for key in floating_keys): + raise ValueError("Reset dataset contains a non-finite state value.") + return state_count + + +def _validate_category_counts( + states: Mapping[str, torch.Tensor], + metadata: Mapping[str, Any], + sampler_cfg: Mapping[str, Any], + state_count: int, + expected_grasping_count: int | None, + expected_non_grasping_count: int | None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Validate category encoding, metadata census, and requested exact quotas.""" + category = states["category"] + grasping = category == GRASPING_CATEGORY + non_grasping = category == NON_GRASPING_CATEGORY + if not bool((grasping | non_grasping).all()): + raise ValueError("State categories must be zero (non-grasping) or one (grasping).") + grasping_count = int(grasping.sum()) + non_grasping_count = int(non_grasping.sum()) + metadata_counts = metadata.get("category_counts") + if not isinstance(metadata_counts, torch.Tensor) or not torch.equal( + metadata_counts.cpu(), torch.tensor((non_grasping_count, grasping_count), dtype=torch.int64) + ): + raise ValueError("Reset-dataset metadata category counts do not match its states.") + if metadata.get("state_count") != state_count: + raise ValueError("Reset-dataset metadata state count does not match its states.") + if expected_grasping_count is None: + expected_grasping_count = int(sampler_cfg["grasping_count"]) + if expected_non_grasping_count is None: + expected_non_grasping_count = int(sampler_cfg["non_grasping_count"]) + if expected_grasping_count is not None and grasping_count != expected_grasping_count: + raise ValueError(f"Expected {expected_grasping_count} grasping states, got {grasping_count}.") + if expected_non_grasping_count is not None and non_grasping_count != expected_non_grasping_count: + raise ValueError(f"Expected {expected_non_grasping_count} non-grasping states, got {non_grasping_count}.") + return grasping, non_grasping + + +def _validate_objective_values( + states: Mapping[str, torch.Tensor], + metadata: Mapping[str, Any], + grasping: torch.Tensor, + non_grasping: torch.Tensor, +) -> None: + """Validate fixed non-grasp values and globally normalized grasp objectives.""" + grasping_count = int(grasping.sum()) + non_grasping_count = int(non_grasping.sum()) + objective = states["objective"] + if non_grasping_count and not bool((objective[non_grasping] == -1.0).all()): + raise ValueError("Every non-grasping objective must be exactly -1.") + if non_grasping_count and ( + not bool((states["objective_raw"][non_grasping] == -1.0).all()) + or not bool((states["objective_components"][non_grasping] == -1.0).all()) + ): + raise ValueError("Every non-grasping raw objective and component must be exactly -1.") + if grasping_count: + grasp_objective = objective[grasping] + if not bool(((grasp_objective >= 0.0) & (grasp_objective <= 1.0)).all()): + raise ValueError("Every grasping objective must lie in [0, 1].") + if grasping_count > 1 and (float(grasp_objective.min()) != 0.0 or float(grasp_objective.max()) != 1.0): + raise ValueError("Grasping objectives must include exact normalized extrema zero and one.") + grasp_sides = states["grasp_side"][grasping] + if not bool(((grasp_sides >= 0) & (grasp_sides <= 3)).all()): + raise ValueError("Grasping states must use side IDs in [0, 3].") + grasp_regions = states["grasp_region"][grasping] + if not bool(((grasp_regions == 0) | (grasp_regions == 1)).all()): + raise ValueError("Grasping states must use broad region zero or near-pour region one.") + near_pour_count = int((grasp_regions == 1).sum()) + expected_near_pour_count = int(metadata["sampler_cfg"]["near_pour_grasp_count"]) + if near_pour_count != expected_near_pour_count: + raise ValueError(f"Expected {expected_near_pour_count} near-pour grasping states, got {near_pour_count}.") + components = states["objective_components"][grasping] + if not bool(((components >= 0.0) & (components <= 1.0)).all()): + raise ValueError("Every grasping objective component must lie in [0, 1].") + weights = metadata.get("objective_weights") + if not isinstance(weights, torch.Tensor) or weights.shape != (3,): + raise ValueError("Reset-dataset metadata must contain three objective weights.") + if metadata.get("objective_component_names") != ( + "source_distance", + "target_gated_inversion", + "target_alignment", + ): + raise ValueError("Reset-dataset metadata has unsupported objective-component semantics.") + expected_raw = (components * weights.to(dtype=components.dtype)).sum(dim=-1) + if not bool(torch.allclose(states["objective_raw"][grasping], expected_raw, atol=1.0e-6, rtol=0.0)): + raise ValueError("Grasping raw objectives do not match their weighted components.") + expected_normalized = normalize_grasp_objectives(states["objective_raw"][grasping]) + if not bool(torch.allclose(grasp_objective, expected_normalized, atol=1.0e-6, rtol=0.0)): + raise ValueError("Grasping objectives do not match global min/max normalization.") + raw_min_max = metadata.get("objective_raw_min_max") + actual_min_max = torch.stack( + (states["objective_raw"][grasping].min(), states["objective_raw"][grasping].max()) + ).cpu() + if not isinstance(raw_min_max, torch.Tensor) or not bool( + torch.allclose(raw_min_max.cpu(), actual_min_max, atol=1.0e-6, rtol=0.0) + ): + raise ValueError("Reset-dataset metadata raw-objective extrema do not match its states.") + if grasping_count % 4 == 0: + side_counts = torch.bincount(grasp_sides.to(torch.long), minlength=4) + if not bool((side_counts == grasping_count // 4).all()): + raise ValueError("Grasping states must be exactly balanced over all four horizontal sides.") + if expected_near_pour_count and expected_near_pour_count % 4 == 0: + near_pour_side_counts = torch.bincount(grasp_sides[grasp_regions == 1].to(torch.long), minlength=4) + if not bool((near_pour_side_counts == expected_near_pour_count // 4).all()): + raise ValueError("Near-pour grasping states must be exactly balanced over all four sides.") + if non_grasping_count: + if not bool((states["grasp_side"][non_grasping] == -1).all()): + raise ValueError("Non-grasping states must use grasp_side=-1.") + if not bool((states["grasp_region"][non_grasping] == -1).all()): + raise ValueError("Non-grasping states must use grasp_region=-1.") + + +def _validate_state_invariants( + states: Mapping[str, torch.Tensor], + sampler_cfg: Mapping[str, Any], + task_contract: Mapping[str, Any], + grasping: torch.Tensor, + non_grasping: torch.Tensor, +) -> None: + """Validate gripper, velocity, pose, attempt, and particle-layout invariants.""" + state_count = states["category"].shape[0] + grasping_count = int(grasping.sum()) + non_grasping_count = int(non_grasping.sum()) + source_box_half = task_contract.get("source_box_half") if isinstance(task_contract, Mapping) else None + if source_box_half is not None and grasping_count: + exact_finger_position = float(source_box_half[1]) + if not bool((states["finger_joint_position"][grasping] == exact_finger_position).all()): + raise ValueError("Grasping finger positions must exactly match the source-cup width.") + grasp_reset_target = task_contract.get("gripper_grasp_reset_target") + contact_min_deflection = task_contract.get("gripper_contact_min_deflection") + cup_grasp_height = task_contract.get("cup_grasp_height") + if grasp_reset_target is None or contact_min_deflection is None or cup_grasp_height is None: + raise ValueError("The reset-cache task contract must define grasp placement and drive preload.") + grasp_reset_target = float(grasp_reset_target) + contact_min_deflection = float(contact_min_deflection) + if not math.isfinite(float(cup_grasp_height)) or float(cup_grasp_height) <= 0.0: + raise ValueError("The configured cup grasp height must be finite and positive.") + if not 0.0 <= grasp_reset_target < exact_finger_position: + raise ValueError("The grasp reset target must lie inside the geometric source-cup contact position.") + if exact_finger_position - grasp_reset_target < contact_min_deflection: + raise ValueError("The grasp reset target does not retain the required finger-contact deflection.") + if not bool((states["finger_joint_target"][grasping] == grasp_reset_target).all()): + raise ValueError("Grasping finger targets must equal the configured reset target.") + near_pour = grasping & (states["grasp_region"] == 1) + if bool(near_pour.any()): + target_rim_height = float(task_contract["target_rim_height"]) + near_pour_valid = above_target_tilted_mask( + states["source_root_pose"][near_pour], + states["target_root_pose"][near_pour], + cup_center_offset=(0.0, 0.0, float(source_box_half[2])), + target_rim_height=target_rim_height, + max_horizontal_distance=float(sampler_cfg["near_pour_horizontal_radius"]), + min_vertical_clearance=float(sampler_cfg["near_pour_height_range"][0]), + min_tilt_angle=float(sampler_cfg["near_pour_tilt_angle_range"][0]), + ) + if not bool(near_pour_valid.all()): + raise ValueError("A tagged near-pour grasp is not geometrically above and tilted over the target.") + gripper_position_range = task_contract.get("gripper_position_range") if isinstance(task_contract, Mapping) else None + if gripper_position_range is not None and non_grasping_count: + lower, upper = (float(value) for value in gripper_position_range) + non_grasp_fingers = states["finger_joint_position"][non_grasping] + if not bool(((non_grasp_fingers >= lower) & (non_grasp_fingers <= upper)).all()): + raise ValueError("Non-grasping finger positions must lie in the complete valid opening range.") + if not bool(torch.allclose(non_grasp_fingers[:, 0], non_grasp_fingers[:, 1], atol=0.0, rtol=0.0)): + raise ValueError("The two Franka fingers must be restored symmetrically.") + if torch.unique(states["attempt_id"]).numel() != state_count: + raise ValueError("Accepted reset states must have unique attempt IDs.") + if non_grasping_count and not bool( + torch.allclose( + states["finger_joint_target"][non_grasping], + states["finger_joint_position"][non_grasping], + atol=0.0, + rtol=0.0, + ) + ): + raise ValueError("Non-grasping finger targets must preserve the sampled opening.") + if not bool((states["particle_layout_id"] == 0).all()): + raise ValueError("The current cache schema requires the shared particle layout ID zero.") + for velocity_key in ( + "arm_joint_velocity", + "finger_joint_velocity", + "source_root_velocity", + "target_root_velocity", + ): + if not bool((states[velocity_key] == 0.0).all()): + raise ValueError(f"Reset-dataset field {velocity_key!r} must start at zero.") + for pose_key in ("source_root_pose", "target_root_pose"): + norms = torch.linalg.vector_norm(states[pose_key][:, 3:7], dim=-1) + if not bool(torch.allclose(norms, torch.ones_like(norms), atol=1.0e-4, rtol=0.0)): + raise ValueError(f"{pose_key} contains a non-unit quaternion.") + tabletop_lower_xy = task_contract.get("tabletop_support_lower_xy") + tabletop_upper_xy = task_contract.get("tabletop_support_upper_xy") + target_box_half = task_contract.get("target_box_half") + obstacle_clearance = float(sampler_cfg.get("obstacle_clearance", 0.0)) + if tabletop_lower_xy is not None and tabletop_upper_xy is not None and target_box_half is not None: + target_supported = oriented_box_supported_by_bounds( + states["target_root_pose"], + target_box_half, + tabletop_lower_xy, + tabletop_upper_xy, + clearance=obstacle_clearance, + ) + if not bool(target_supported.all()): + raise ValueError("A target bowl reset pose is not fully supported by the tabletop.") + if non_grasping_count and source_box_half is not None: + source_supported = oriented_box_supported_by_bounds( + states["source_root_pose"][non_grasping], + source_box_half, + tabletop_lower_xy, + tabletop_upper_xy, + clearance=obstacle_clearance, + ) + if not bool(source_supported.all()): + raise ValueError("A non-grasping source-cup reset pose is not fully supported by the tabletop.") + + +def _validate_particle_layouts(payload: Mapping[str, Any]) -> None: + """Validate the one shared deterministic local particle layout.""" + layouts = payload.get("particle_layouts") + if not isinstance(layouts, Mapping) or set(layouts) != {"local_position", "local_velocity"}: + raise ValueError("Reset dataset has invalid particle layouts.") + local_position = layouts["local_position"] + local_velocity = layouts["local_velocity"] + if ( + not isinstance(local_position, torch.Tensor) + or local_position.ndim != 3 + or local_position.shape[0] != 1 + or local_position.shape[2] != 3 + or local_velocity.shape != local_position.shape + or not bool(torch.isfinite(local_position).all()) + or not bool((local_velocity == 0.0).all()) + ): + raise ValueError("Reset-dataset particle layout must be one finite (P, 3) layout with zero velocity.") + + +def save_reset_dataset(payload: Mapping[str, Any], output_path: str | Path) -> None: + """Atomically write a validated reset dataset.""" + reset_dataset_save_atomic(payload, output_path, validator=validate_reset_dataset) + + +class FrankaPourResetDatasetGenerator: + """Generate the complete Franka Pour reset dataset through rejection sampling.""" + + def __init__(self, env: FrankaPourEnv, cfg: FrankaPourResetDatasetGeneratorCfg): + if cfg.near_pour_grasp_count in (0, cfg.grasping_count): + raise ValueError("Candidate generation requires both broad and near-pour grasping states.") + if cfg.near_pour_grasp_count % 4 != 0: + raise ValueError("near_pour_grasp_count must be divisible by four for side-balanced proposals.") + if (cfg.grasping_count - cfg.near_pour_grasp_count) % 4 != 0: + raise ValueError("The broad grasping count must be divisible by four for side-balanced proposals.") + self.env = env + self.task_cfg = env.cfg + self.cfg = cfg + self.device = torch.device(env.device) + self.generator = torch.Generator(device=self.device) + self.generator.manual_seed(cfg.seed) + self._next_attempt_id = 0 + self._attempt_counts = {NON_GRASPING_CATEGORY: 0, GRASPING_CATEGORY: 0} + self._rejection_counts: dict[int, dict[str, int]] = { + NON_GRASPING_CATEGORY: defaultdict(int), + GRASPING_CATEGORY: defaultdict(int), + } + self._build_ik_context() + + @torch.inference_mode() + def generate(self) -> dict[str, Any]: + """Generate, score, validate, and return the configured exact category quotas.""" + near_pour_grasping = self._sample_category( + GRASPING_CATEGORY, + self.cfg.near_pour_grasp_count, + near_pour=True, + ) + broad_grasping = self._sample_category( + GRASPING_CATEGORY, + self.cfg.grasping_count - self.cfg.near_pour_grasp_count, + near_pour=False, + ) + grasping = {key: torch.cat((near_pour_grasping[key], broad_grasping[key]), dim=0) for key in _STATE_KEYS} + raw = grasping["objective_raw"] + grasping["objective"] = normalize_grasp_objectives(raw) + non_grasping = self._sample_category(NON_GRASPING_CATEGORY, self.cfg.non_grasping_count) + states = {key: torch.cat((grasping[key], non_grasping[key]), dim=0) for key in _STATE_KEYS} + permutation = torch.randperm( + self.cfg.grasping_count + self.cfg.non_grasping_count, + device=self.device, + generator=self.generator, + ) + states = {key: value[permutation] for key, value in states.items()} + + metadata = { + "seed": self.cfg.seed, + "storage_order": "seeded_random_permutation", + "source_region_center": self._source_region_center.detach().cpu(), + "cup_center_offset": self._cup_center_offset.detach().cpu(), + "objective_weights": torch.tensor(self.cfg.objective_weights, dtype=torch.float32), + "objective_component_names": ( + "source_distance", + "target_gated_inversion", + "target_alignment", + ), + "objective_raw_min_max": torch.stack((raw.min(), raw.max())).detach().cpu(), + "attempt_counts": torch.tensor( + (self._attempt_counts[NON_GRASPING_CATEGORY], self._attempt_counts[GRASPING_CATEGORY]), + dtype=torch.int64, + ), + "rejection_counts": { + name: torch.tensor( + ( + self._rejection_counts[NON_GRASPING_CATEGORY].get(name, 0), + self._rejection_counts[GRASPING_CATEGORY].get(name, 0), + ), + dtype=torch.int64, + ) + for name in sorted( + set(self._rejection_counts[NON_GRASPING_CATEGORY]) | set(self._rejection_counts[GRASPING_CATEGORY]) + ) + }, + "task_contract": self._task_contract(), + } + return build_reset_dataset_payload(states, self.env._media_local_points_t, metadata, self.cfg) + + def _build_ik_context(self) -> None: + import newton + import warp as wp + from isaaclab_newton.cloner import copy_newton_source_builder + from isaaclab_newton.ik.newton_ik_objectives_cfg import ( + NewtonIKJointLimitObjectiveCfg, + NewtonIKPoseObjectiveCfg, + ) + from isaaclab_newton.ik.newton_ik_solver import NewtonIKSolver + from isaaclab_newton.ik.newton_ik_solver_cfg import NewtonIKSolverCfg + + import isaaclab.sim as sim_utils + from isaaclab.cloner import resolve_clone_plan_source + + plan = sim_utils.SimulationContext.instance().get_clone_plan() + resolved = resolve_clone_plan_source(self.env._robot.cfg.prim_path, plan) if plan is not None else None + if resolved is None: + raise RuntimeError(f"Could not resolve clone-plan source for {self.env._robot.cfg.prim_path!r}.") + source_builder = copy_newton_source_builder(resolved[0]) + prototype_origin = -self.env.env_origins[0] + prototype_xform = wp.transform(wp.vec3(*prototype_origin.tolist()), wp.quat_identity()) + self._prototype_builder = newton.ModelBuilder(up_axis=source_builder.up_axis) + self._prototype_builder.add_builder(source_builder, xform=prototype_xform) + if not any( + "/Table/" in str(label) or str(label).endswith("/Table") for label in self._prototype_builder.shape_label + ): + table_prim_path = self.env.scene["table"].cfg.prim_path + table_resolved = resolve_clone_plan_source(table_prim_path, plan) + if table_resolved is None: + raise RuntimeError(f"Could not resolve clone-plan source for {table_prim_path!r}.") + table_builder = copy_newton_source_builder(table_resolved[0]) + self._prototype_builder.add_builder(table_builder, xform=prototype_xform) + if not any( + "/Table/" in str(label) or str(label).endswith("/Table") for label in self._prototype_builder.shape_label + ): + raise RuntimeError("Reset-dataset validation could not import SeattleLab table collision geometry.") + self._tabletop_support_lower_xy, self._tabletop_support_upper_xy = _derive_tabletop_support_bounds( + self.env, resolved[0] + ) + self._ik_model = self._prototype_builder.finalize(device=str(self.device)) + + body_names = [str(label).rsplit("/", 1)[-1] for label in self._ik_model.body_label] + hand_matches = [index for index, name in enumerate(body_names) if name == self.task_cfg.tcp_body_name] + if len(hand_matches) != 1: + raise RuntimeError(f"Expected one IK body named {self.task_cfg.tcp_body_name!r}, got {hand_matches}.") + self._hand_id = hand_matches[0] + joint_names = [str(label).rsplit("/", 1)[-1] for label in self._ik_model.joint_label] + joint_q_start = wp.to_torch(self._ik_model.joint_q_start).to(device=self.device, dtype=torch.long) + + def coordinate_id(name: str) -> int: + matches = [index for index, joint_name in enumerate(joint_names) if joint_name == name] + if len(matches) != 1: + raise RuntimeError(f"Expected one IK joint named {name!r}, got {matches}.") + return int(joint_q_start[matches[0]].item()) + + self._arm_coordinate_ids = torch.tensor( + [coordinate_id(name) for name in _ARM_JOINT_NAMES], device=self.device, dtype=torch.long + ) + self._finger_coordinate_ids = torch.tensor( + [coordinate_id(name) for name in _FINGER_JOINT_NAMES], device=self.device, dtype=torch.long + ) + self._arm_limits = self.env._joint_pos_limits_t[0, self.env._arm_joint_ids].to(self.device) + finger_limits = self.env._joint_pos_limits_t[0, self.env._finger_joint_ids].to(self.device) + self._finger_position_range = (float(finger_limits[:, 0].max()), float(finger_limits[:, 1].min())) + if self._finger_position_range[0] >= self._finger_position_range[1]: + raise RuntimeError(f"The two Franka fingers have no shared valid opening range: {finger_limits}.") + self._arm_home = torch.tensor(self.task_cfg.arm_home, device=self.device, dtype=torch.float32) + self._tcp_offset_pos = torch.tensor(self.task_cfg.tcp_offset_pos, device=self.device, dtype=torch.float32) + self._tcp_offset_quaternion = torch.tensor( + self.task_cfg.tcp_offset_rot, device=self.device, dtype=torch.float32 + ) + self._joint_seed = ( + wp.to_torch(self._ik_model.joint_q) + .to(device=self.device, dtype=torch.float32) + .repeat(self.cfg.batch_size, 1) + ) + self._joint_seed[:, self._arm_coordinate_ids] = self._arm_home + + target_name = "reset_dataset_tcp" + objectives = [ + NewtonIKPoseObjectiveCfg( + body_name=self.task_cfg.tcp_body_name, + name=target_name, + body_offset_pos=self.task_cfg.tcp_offset_pos, + body_offset_rot=self.task_cfg.tcp_offset_rot, + position_weight=100.0, + rotation_weight=5.0, + ), + NewtonIKJointLimitObjectiveCfg(weight=1.0), + ] + self._ik_solver = NewtonIKSolver( + NewtonIKSolverCfg( + optimizer="lm", + jacobian_mode="analytic", + sampler="gauss", + n_seeds=self.cfg.ik_seeds, + noise_std=self.cfg.ik_noise_std, + iterations=self.cfg.ik_iterations, + lambda_initial=0.1, + rng_seed=self.cfg.seed, + ), + model=self._ik_model, + num_envs=self.cfg.batch_size, + device=str(self.device), + objectives=objectives, + link_resolver=lambda _name: self._hand_id, + ) + self._pose_objective = self._ik_solver.objectives_by_name[target_name] + + source_half = torch.tensor(self.task_cfg.cup_grasp_box_half, device=self.device, dtype=torch.float32) + self._source_half = source_half + self._cup_center_offset = torch.tensor((0.0, 0.0, float(source_half[2])), device=self.device) + self._cup_grasp_offset = torch.tensor( + (0.0, 0.0, float(self.task_cfg.cup_grasp_height)), + device=self.device, + ) + source_inner_lower = torch.as_tensor(self.env._source_inner_lo, device=self.device, dtype=torch.float32) + source_inner_upper = torch.as_tensor(self.env._source_inner_hi, device=self.device, dtype=torch.float32) + local_particles = self.env._media_local_points_t.to(self.device) + if not bool(((local_particles >= source_inner_lower) & (local_particles <= source_inner_upper)).all()): + raise RuntimeError("The existing particle sampler produced a point outside the source-cup cavity.") + nominal_source = torch.tensor(self.task_cfg.cup_reset_pos, device=self.device, dtype=torch.float32) + self._source_region_center = nominal_source + self._cup_center_offset + self._target_half = torch.tensor( + ( + self.task_cfg.target_cup_inner_width * 0.5 + self.task_cfg.target_cup_wall_thickness, + self.task_cfg.target_cup_inner_depth * 0.5 + self.task_cfg.target_cup_wall_thickness, + (self.task_cfg.target_cup_cavity_depth + self.task_cfg.target_cup_bottom_thickness) * 0.5, + ), + device=self.device, + dtype=torch.float32, + ) + self._target_rim_height = float(self._target_half[2] * 2.0) + + def _sample_category( + self, + category: int, + required_count: int, + *, + near_pour: bool = False, + ) -> dict[str, torch.Tensor]: + if required_count <= 0: + raise ValueError("A sampled reset category must request at least one state.") + if near_pour and category != GRASPING_CATEGORY: + raise ValueError("Only grasping states can use the near-pour proposal region.") + accepted_count = 0 + accepted_side_counts = torch.zeros(4, device=self.device, dtype=torch.long) + side_quotas = torch.full((4,), required_count // 4, device=self.device, dtype=torch.long) + side_quotas[: required_count % 4] += 1 + max_attempts = required_count * self.cfg.max_attempt_multiplier + initial_attempt_count = self._attempt_counts[category] + + def evaluate_batch(candidate_ids: range) -> dict[str, torch.Tensor]: + nonlocal accepted_count + candidate_count = len(candidate_ids) + proposal = self._propose_batch( + category, + accepted_side_counts, + side_quotas, + near_pour=near_pour, + count=candidate_count, + ) + self._attempt_counts[category] += candidate_count + valid = self._validate_proposal(proposal, category) + valid_before_quota = valid.clone() + if category == GRASPING_CATEGORY: + keep = torch.zeros_like(valid) + for side in range(4): + side_rows = torch.where(valid & (proposal["grasp_side"] == side))[0] + remaining = int(side_quotas[side] - accepted_side_counts[side]) + chosen = side_rows[: max(remaining, 0)] + keep[chosen] = True + accepted_side_counts[side] += chosen.numel() + valid = keep + else: + valid_rows = torch.where(valid)[0][: required_count - accepted_count] + valid = torch.zeros_like(valid) + valid[valid_rows] = True + self._rejection_counts[category]["quota_full"] += int((valid_before_quota & ~valid).sum()) + accepted = int(valid.sum()) + accepted_count += accepted + return {key: proposal[key][valid].detach().clone() for key in _STATE_KEYS} + + def batch_count(batch: dict[str, torch.Tensor]) -> int: + return int(batch["category"].shape[0]) + + def batch_slice(batch: dict[str, torch.Tensor], count: int) -> dict[str, torch.Tensor]: + return {key: value[:count] for key, value in batch.items()} + + try: + batches, _ = reset_dataset_collect_batches( + required_count, + batch_size=self.cfg.batch_size, + max_candidate_count=max_attempts, + evaluate_batch=evaluate_batch, + batch_count=batch_count, + batch_slice=batch_slice, + ) + except RuntimeError as error: + attempted = self._attempt_counts[category] - initial_attempt_count + raise RuntimeError( + f"Reset-dataset sampling exhausted {attempted} {self._category_name(category)} candidates " + f"after accepting {accepted_count}/{required_count}; rejection counts are " + f"{dict(self._rejection_counts[category])}." + ) from error + return {key: torch.cat([batch[key] for batch in batches], dim=0) for key in _STATE_KEYS} + + def _propose_batch( + self, + category: int, + accepted_side_counts: torch.Tensor, + side_quotas: torch.Tensor, + *, + near_pour: bool = False, + count: int | None = None, + ) -> dict[str, torch.Tensor]: + count = self.cfg.batch_size if count is None else count + attempt_ids = torch.arange( + self._next_attempt_id, self._next_attempt_id + count, device=self.device, dtype=torch.int64 + ) + self._next_attempt_id += count + target_positions = self._sample_target_positions(count) + target_quaternions = self._identity_quaternions(count) + finger_position = torch.empty((count, 2), device=self.device) + + if category == GRASPING_CATEGORY: + side_deficit = side_quotas - accepted_side_counts + active_sides = torch.where(side_deficit > 0)[0] + side_order = active_sides[torch.argsort(side_deficit[active_sides], descending=True)] + grasp_side = side_order[torch.arange(count, device=self.device) % side_order.numel()].to(torch.int8) + side_angle = grasp_side.to(torch.float32) * (0.5 * math.pi) + side_axis = torch.zeros((count, 3), device=self.device) + side_axis[:, 2] = 1.0 + side_rotation = math_utils.quat_from_angle_axis(side_angle, side_axis) + base_grasp = torch.tensor( + self.task_cfg.cup_grasp_tcp_quat_c, device=self.device, dtype=torch.float32 + ).expand(count, -1) + cup_to_tcp = math_utils.quat_mul(side_rotation, base_grasp) + offset_tcp = torch.randn((count, 3), device=self.device, generator=self.generator) + offset_tcp *= torch.tensor(self.cfg.grasp_position_std, device=self.device) + if near_pour: + source_quaternions, cup_centers = self._sample_near_pour_cup_poses(target_positions) + tcp_quaternions = math_utils.quat_unique(math_utils.quat_mul(source_quaternions, cup_to_tcp)) + source_positions = cup_centers - math_utils.quat_apply( + source_quaternions, + self._cup_center_offset.expand(count, -1), + ) + grasp_positions = source_positions + math_utils.quat_apply( + source_quaternions, + self._cup_grasp_offset.expand(count, -1), + ) + tcp_positions = grasp_positions - math_utils.quat_apply(tcp_quaternions, offset_tcp) + else: + tcp_positions = self._sample_workspace_positions(count) + tcp_quaternions = self._sample_uniform_quaternions(count) + source_quaternions = math_utils.quat_unique( + math_utils.quat_mul(tcp_quaternions, math_utils.quat_conjugate(cup_to_tcp)) + ) + source_positions = source_root_position_from_tcp_grasp( + tcp_positions, + tcp_quaternions, + source_quaternions, + self._cup_grasp_offset, + offset_tcp, + ) + finger_position.fill_(float(self._source_half[1])) + # Restore the physical fingers tangent to the cup, but command the task's full close + # target from the first physics step. Initializing the physical joints at that target + # would embed the collision meshes in the cup; separating q from its target creates a + # strong bilateral preload without reset penetration. + finger_target = torch.full_like( + finger_position, + float(self.task_cfg.actions.gripper_action.close_position), + ) + grasp_region = torch.full((count,), int(near_pour), device=self.device, dtype=torch.int8) + else: + tcp_positions = self._sample_workspace_positions(count) + tcp_quaternions = self._sample_uniform_quaternions(count) + grasp_side = torch.full((count,), -1, device=self.device, dtype=torch.int8) + grasp_region = torch.full((count,), -1, device=self.device, dtype=torch.int8) + source_positions, source_quaternions = self._sample_table_source_poses(count) + finger_position.uniform_(*self._finger_position_range, generator=self.generator) + # The two fingers are driven symmetrically in the task; one sampled physical opening + # therefore maps to equal joint coordinates rather than two unrelated finger widths. + finger_position[:, 1] = finger_position[:, 0] + finger_target = finger_position.clone() + + ( + robot_q, + ik_cost, + position_residual, + rotation_residual, + actual_tcp_position, + actual_tcp_quaternion, + ik_solution_valid, + ) = self._solve_ik(tcp_positions, tcp_quaternions, finger_position) + source_pose = torch.cat((source_positions, source_quaternions), dim=-1) + target_pose = torch.cat((target_positions, target_quaternions), dim=-1) + objective_components = torch.full((count, 3), -1.0, device=self.device) + objective_raw = torch.full((count,), -1.0, device=self.device) + objective = torch.full((count,), -1.0, device=self.device) + if category == GRASPING_CATEGORY: + objective_components = grasp_objective_components( + source_pose, + target_pose, + source_region_center=self._source_region_center, + cup_center_offset=self._cup_center_offset, + target_rim_height=self._target_rim_height, + distance_threshold=self.cfg.objective_distance_threshold, + target_horizontal_threshold=self.cfg.objective_target_horizontal_threshold, + target_height_threshold=self.cfg.objective_target_height_threshold, + inversion_gate_horizontal_threshold=self.cfg.objective_inversion_gate_horizontal_threshold, + ) + weights = torch.tensor(self.cfg.objective_weights, device=self.device) + objective_raw = (objective_components * weights).sum(dim=-1) + objective = objective_raw.clone() + + return { + "arm_joint_position": robot_q[:, self._arm_coordinate_ids], + "arm_joint_velocity": torch.zeros((count, 7), device=self.device), + "finger_joint_position": finger_position, + "finger_joint_velocity": torch.zeros((count, 2), device=self.device), + "finger_joint_target": finger_target, + "source_root_pose": source_pose, + "source_root_velocity": torch.zeros((count, 6), device=self.device), + "target_root_pose": target_pose, + "target_root_velocity": torch.zeros((count, 6), device=self.device), + "category": torch.full((count,), category, device=self.device, dtype=torch.int8), + "objective": objective, + "objective_raw": objective_raw, + "objective_components": objective_components, + "grasp_region": grasp_region, + "grasp_side": grasp_side, + "attempt_id": attempt_ids, + "particle_layout_id": torch.zeros(count, device=self.device, dtype=torch.int32), + "ik_cost": ik_cost, + "ik_position_residual": position_residual, + "ik_rotation_residual": rotation_residual, + "_robot_q": robot_q, + "_tcp_target_position": tcp_positions, + "_tcp_target_quaternion": tcp_quaternions, + "_tcp_actual_position": actual_tcp_position, + "_tcp_actual_quaternion": actual_tcp_quaternion, + "_ik_solution_valid": ik_solution_valid, + } + + def _solve_ik( + self, tcp_positions: torch.Tensor, tcp_quaternions: torch.Tensor, finger_position: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + import warp as wp + + count = int(tcp_positions.shape[0]) + if not 0 < count <= self.cfg.batch_size: + raise ValueError(f"IK batch count must be in [1, {self.cfg.batch_size}], got {count}.") + if tcp_quaternions.shape != (count, 4) or finger_position.shape != (count, 2): + raise ValueError("IK targets and finger positions have inconsistent batch dimensions.") + + # NewtonIK is constructed once at the configured maximum batch size. Pad only the final + # rejection-sampling batch and discard its duplicate rows after solving; rebuilding the + # solver for a short tail batch would be both expensive and unnecessary. + if count < self.cfg.batch_size: + padding = self.cfg.batch_size - count + tcp_positions = torch.cat((tcp_positions, tcp_positions[-1:].expand(padding, -1)), dim=0) + tcp_quaternions = torch.cat((tcp_quaternions, tcp_quaternions[-1:].expand(padding, -1)), dim=0) + finger_position_padded = torch.cat((finger_position, finger_position[-1:].expand(padding, -1)), dim=0) + else: + finger_position_padded = finger_position + + self._pose_objective.position_objective.set_target_positions( + wp.from_torch(tcp_positions.contiguous(), dtype=wp.vec3) + ) + self._pose_objective.rotation_objective.set_target_rotations( + wp.from_torch(tcp_quaternions.contiguous(), dtype=wp.vec4) + ) + seed = self._joint_seed.clone() + seed[:, self._finger_coordinate_ids] = finger_position_padded + self._ik_solver.solve(wp.from_torch(seed.contiguous(), dtype=wp.float32)) + expanded_joint_q = wp.to_torch(self._ik_solver.joint_q).reshape(self.cfg.batch_size, self.cfg.ik_seeds, -1) + expanded_cost = wp.to_torch(self._ik_solver.costs).reshape(self.cfg.batch_size, self.cfg.ik_seeds) + residuals = wp.to_torch(self._ik_solver.solver.residuals).reshape(self.cfg.batch_size, self.cfg.ik_seeds, -1) + expanded_position_residual = torch.linalg.vector_norm(residuals[:, :, :3] / 100.0, dim=-1) + expanded_rotation_residual = torch.linalg.vector_norm(residuals[:, :, 3:6] / 5.0, dim=-1) + expanded_arm_q = expanded_joint_q[:, :, self._arm_coordinate_ids] + arm_margin = torch.minimum( + expanded_arm_q - self._arm_limits[:, 0], self._arm_limits[:, 1] - expanded_arm_q + ).amin(dim=-1) + seed_valid = ( + torch.isfinite(expanded_joint_q).all(dim=-1) + & torch.isfinite(expanded_cost) + & torch.isfinite(expanded_position_residual) + & torch.isfinite(expanded_rotation_residual) + & (expanded_cost <= self.cfg.ik_max_cost) + & (arm_margin >= self.cfg.ik_joint_margin) + & (expanded_position_residual <= self.cfg.ik_max_position_residual) + & (expanded_rotation_residual <= self.cfg.ik_max_rotation_residual) + & (torch.linalg.vector_norm(expanded_arm_q - self._arm_home, dim=-1) <= self.cfg.ik_max_home_distance) + ) + valid_cost = expanded_cost.masked_fill(~seed_valid, torch.inf) + best_seed = valid_cost.argmin(dim=-1) + solution_valid = seed_valid.any(dim=-1) + fallback_seed = torch.nan_to_num(expanded_cost, nan=torch.inf, posinf=torch.inf, neginf=torch.inf).argmin( + dim=-1 + ) + best_seed = torch.where(solution_valid, best_seed, fallback_seed) + rows = torch.arange(count, device=self.device) + selected_seed = best_seed[rows] + solved = expanded_joint_q[rows, selected_seed].clone() + solved[:, self._finger_coordinate_ids] = finger_position + cost = expanded_cost[rows, selected_seed].clone() + position_residual = expanded_position_residual[rows, selected_seed].clone() + rotation_residual = expanded_rotation_residual[rows, selected_seed].clone() + solution_valid = solution_valid[rows] + body_poses = wp.to_torch(self._ik_solver.solver.body_q).reshape(self.cfg.batch_size, self.cfg.ik_seeds, -1, 7) + hand_pose = body_poses[rows, selected_seed, self._hand_id].clone() + actual_tcp_position, actual_tcp_quaternion = math_utils.combine_frame_transforms( + hand_pose[:, :3], + hand_pose[:, 3:7], + self._tcp_offset_pos.expand(count, -1), + self._tcp_offset_quaternion.expand(count, -1), + ) + return ( + solved, + cost, + position_residual, + rotation_residual, + actual_tcp_position, + actual_tcp_quaternion, + solution_valid, + ) + + def _validate_proposal(self, proposal: dict[str, torch.Tensor], category: int) -> torch.Tensor: + count = int(proposal["category"].shape[0]) + valid = torch.ones(count, device=self.device, dtype=torch.bool) + arm_q = proposal["arm_joint_position"] + arm_margin = torch.minimum(arm_q - self._arm_limits[:, 0], self._arm_limits[:, 1] - arm_q).amin(dim=-1) + checks = ( + ("ik_no_valid_seed", proposal["_ik_solution_valid"]), + ("ik_nonfinite", torch.isfinite(proposal["_robot_q"]).all(dim=-1) & torch.isfinite(proposal["ik_cost"])), + ("ik_cost", proposal["ik_cost"] <= self.cfg.ik_max_cost), + ("ik_joint_limit", arm_margin >= self.cfg.ik_joint_margin), + ("ik_position_residual", proposal["ik_position_residual"] <= self.cfg.ik_max_position_residual), + ("ik_rotation_residual", proposal["ik_rotation_residual"] <= self.cfg.ik_max_rotation_residual), + ( + "ik_discontinuity", + torch.linalg.vector_norm(arm_q - self._arm_home, dim=-1) <= self.cfg.ik_max_home_distance, + ), + ) + for reason, check in checks: + valid = self._reject(valid, check, category, reason) + + source_pose = proposal["source_root_pose"] + target_pose = proposal["target_root_pose"] + source_center = source_pose[:, :3] + math_utils.quat_apply( + source_pose[:, 3:7], self._cup_center_offset.expand(count, -1) + ) + target_center = target_pose[:, :3].clone() + target_center[:, 2] += self._target_half[2] + source_target_clear = ~oriented_boxes_overlap( + source_center, + source_pose[:, 3:7], + self._source_half, + target_center, + target_pose[:, 3:7], + self._target_half, + clearance=self.cfg.obstacle_clearance, + ) + valid = self._reject(valid, source_target_clear, category, "source_target_collision") + valid = self._reject(valid, self._box_above_table(source_pose), category, "source_table_collision") + valid = self._reject(valid, self._box_above_table(target_pose, target=True), category, "target_table_collision") + valid = self._reject( + valid, + oriented_box_supported_by_bounds( + target_pose, + self._target_half, + self._tabletop_support_lower_xy, + self._tabletop_support_upper_xy, + clearance=self.cfg.obstacle_clearance, + ), + category, + "target_table_support", + ) + if category == NON_GRASPING_CATEGORY: + valid = self._reject( + valid, + oriented_box_supported_by_bounds( + source_pose, + self._source_half, + self._tabletop_support_lower_xy, + self._tabletop_support_upper_xy, + clearance=self.cfg.obstacle_clearance, + ), + category, + "source_table_support", + ) + valid = self._reject(valid, self._particles_in_workspace(source_pose), category, "particle_workspace") + if category == GRASPING_CATEGORY and bool((proposal["grasp_region"] == 1).any()): + valid = self._reject( + valid, + above_target_tilted_mask( + source_pose, + target_pose, + cup_center_offset=self._cup_center_offset, + target_rim_height=self._target_rim_height, + max_horizontal_distance=self.cfg.near_pour_horizontal_radius, + min_vertical_clearance=self.cfg.near_pour_height_range[0], + min_tilt_angle=self.cfg.near_pour_tilt_angle_range[0], + ), + category, + "invalid_near_pour_geometry", + ) + if category == GRASPING_CATEGORY: + source_grasp_point = source_pose[:, :3] + math_utils.quat_apply( + source_pose[:, 3:7], + self._cup_grasp_offset.expand(count, -1), + ) + cup_offset_tcp = math_utils.quat_apply_inverse( + proposal["_tcp_actual_quaternion"], + source_grasp_point - proposal["_tcp_actual_position"], + ) + seating_bound = torch.tensor(self.cfg.grasp_seating_max_offset, device=self.device) + side_angle = proposal["grasp_side"].to(torch.float32) * (0.5 * math.pi) + side_axis = torch.zeros((count, 3), device=self.device) + side_axis[:, 2] = 1.0 + expected_cup_to_tcp = math_utils.quat_mul( + math_utils.quat_from_angle_axis(side_angle, side_axis), + torch.tensor(self.task_cfg.cup_grasp_tcp_quat_c, device=self.device).expand(count, -1), + ) + actual_cup_to_tcp = math_utils.quat_mul( + math_utils.quat_conjugate(source_pose[:, 3:7]), proposal["_tcp_actual_quaternion"] + ) + seating_rotation_error = math_utils.quat_error_magnitude(actual_cup_to_tcp, expected_cup_to_tcp) + valid = self._reject( + valid, + (cup_offset_tcp.abs() <= seating_bound).all(dim=-1) + & (seating_rotation_error <= self.cfg.grasp_seating_max_rotation_error), + category, + "invalid_grasp_seating", + ) + else: + cup_tcp_distance = torch.linalg.vector_norm(source_center - proposal["_tcp_actual_position"], dim=-1) + valid = self._reject( + valid, + cup_tcp_distance >= self.cfg.non_grasping_min_tcp_source_distance, + category, + "source_inside_gripper", + ) + + # A grasp candidate gets the required robot-only screen before inserting the source cup. + # The source proxy is parked far outside the workspace for this first pass. + if category == GRASPING_CATEGORY: + pre_indices = torch.where(valid)[0] + if pre_indices.numel(): + far_source = proposal["source_root_pose"][pre_indices, :3].clone() + far_source[:, 2] = 10.0 + pre_clear = self._collision_screen( + proposal["_robot_q"][pre_indices], + far_source, + proposal["source_root_pose"][pre_indices, 3:7], + proposal["target_root_pose"][pre_indices, :3], + allow_finger_contact=False, + ) + full_check = torch.zeros_like(valid) + full_check[pre_indices] = pre_clear + valid = self._reject(valid, full_check, category, "robot_preinsert_collision") + + final_indices = torch.where(valid)[0] + if final_indices.numel(): + collision_clear = self._collision_screen( + proposal["_robot_q"][final_indices], + proposal["source_root_pose"][final_indices, :3], + proposal["source_root_pose"][final_indices, 3:7], + proposal["target_root_pose"][final_indices, :3], + allow_finger_contact=category == GRASPING_CATEGORY, + ) + full_check = torch.zeros_like(valid) + full_check[final_indices] = collision_clear + valid = self._reject(valid, full_check, category, "complete_state_collision") + return valid + + def _collision_screen( + self, + robot_q: torch.Tensor, + source_position: torch.Tensor, + source_quaternion: torch.Tensor, + target_position: torch.Tensor, + *, + allow_finger_contact: bool, + ) -> torch.Tensor: + from ._reset_collision_screen import collision_free_reset_candidates + + return collision_free_reset_candidates( + self._prototype_builder, + robot_q, + source_position, + source_quaternion, + target_position, + source_box_half=tuple(float(value) for value in self._source_half), + target_vertices=self.env._target_vertices, + target_indices=self.env._target_indices, + collider_margin=float(self.task_cfg.collider_margin), + device=str(self.device), + penetration_tolerance=self.cfg.collision_penetration_tolerance, + allow_source_finger_contact=allow_finger_contact, + source_finger_penetration_tolerance=self.cfg.finger_contact_penetration_tolerance, + check_self_collision=True, + check_complete_robot_table=True, + ) + + def _box_above_table(self, pose: torch.Tensor, *, target: bool = False) -> torch.Tensor: + half = self._target_half if target else self._source_half + center = pose[:, :3] + math_utils.quat_apply( + pose[:, 3:7], torch.tensor((0.0, 0.0, float(half[2])), device=self.device).expand(pose.shape[0], -1) + ) + rotation = math_utils.matrix_from_quat(pose[:, 3:7]).abs() + vertical_radius = (rotation[:, 2, :] * half).sum(dim=-1) + return center[:, 2] - vertical_radius >= -self.cfg.collision_penetration_tolerance + + def _particles_in_workspace(self, source_pose: torch.Tensor) -> torch.Tensor: + # Reuse the task's existing particle-sampling transform rather than maintaining a second + # interpretation of the cup-local media layout in this offline tool. + world = self.env._sample_cup_media(source_pose[:, :3], source_pose[:, 3:7]) + lower = torch.tensor(self.task_cfg.particle_workspace_lower_bound, device=self.device) + upper = torch.tensor(self.task_cfg.particle_workspace_upper_bound, device=self.device) + return ((world >= lower) & (world <= upper)).all(dim=-1).all(dim=-1) + + def _sample_workspace_positions(self, count: int) -> torch.Tensor: + fraction = self.cfg.workspace_central_fraction + + def central(bounds: tuple[float, float]) -> tuple[float, float]: + midpoint = 0.5 * (bounds[0] + bounds[1]) + half = 0.5 * (bounds[1] - bounds[0]) * fraction + return midpoint - half, midpoint + half + + radius_lower, radius_upper = central(self.cfg.workspace_radius_range) + # Azimuth is periodic rather than a kinematic limit. Preserve the complete configured + # circle and apply the central-fraction margin only to radial reach and height. + angle_lower, angle_upper = self.cfg.workspace_azimuth_range + height_lower, height_upper = central(self.cfg.workspace_height_range) + radius_square = torch.empty(count, device=self.device).uniform_( + radius_lower**2, radius_upper**2, generator=self.generator + ) + radius = torch.sqrt(radius_square) + angle = torch.empty(count, device=self.device).uniform_(angle_lower, angle_upper, generator=self.generator) + height = torch.empty(count, device=self.device).uniform_(height_lower, height_upper, generator=self.generator) + return torch.stack((radius * torch.cos(angle), radius * torch.sin(angle), height), dim=-1) + + def _sample_uniform_quaternions(self, count: int) -> torch.Tensor: + # Shoemake's transform produces Haar-uniform SO(3) rotations in XYZW order. + random = torch.rand((count, 3), device=self.device, generator=self.generator) + first = torch.sqrt(1.0 - random[:, 0]) + second = torch.sqrt(random[:, 0]) + quaternion = torch.stack( + ( + first * torch.sin(2.0 * math.pi * random[:, 1]), + first * torch.cos(2.0 * math.pi * random[:, 1]), + second * torch.sin(2.0 * math.pi * random[:, 2]), + second * torch.cos(2.0 * math.pi * random[:, 2]), + ), + dim=-1, + ) + return math_utils.quat_unique(quaternion) + + def _sample_table_source_poses(self, count: int) -> tuple[torch.Tensor, torch.Tensor]: + radius_range = self.task_cfg.curriculum_randomized_source_radius_range + if radius_range is None: + center = torch.tensor(self.task_cfg.cup_reset_pos, device=self.device) + extent = torch.tensor(self.task_cfg.curriculum_randomized_source_position_range, device=self.device) + random_offset = 2.0 * torch.rand((count, 2), device=self.device, generator=self.generator) - 1.0 + xy = center[:2] + random_offset * extent + else: + radius_square = torch.empty(count, device=self.device).uniform_( + radius_range[0] ** 2, radius_range[1] ** 2, generator=self.generator + ) + radius = torch.sqrt(radius_square) + angle = torch.empty(count, device=self.device).uniform_( + -self.task_cfg.curriculum_randomized_source_azimuth_range, + self.task_cfg.curriculum_randomized_source_azimuth_range, + generator=self.generator, + ) + xy = torch.stack((radius * torch.cos(angle), radius * torch.sin(angle)), dim=-1) + positions = torch.zeros((count, 3), device=self.device) + positions[:, :2] = xy + yaw = torch.empty(count, device=self.device).uniform_(-math.pi, math.pi, generator=self.generator) + quaternions = torch.zeros((count, 4), device=self.device) + quaternions[:, 2] = torch.sin(0.5 * yaw) + quaternions[:, 3] = torch.cos(0.5 * yaw) + return positions, quaternions + + def _sample_target_positions(self, count: int) -> torch.Tensor: + center = torch.tensor(self.task_cfg.curriculum_randomized_target_center_xy, device=self.device) + extent = torch.tensor(self.task_cfg.curriculum_randomized_target_position_range, device=self.device) + configured_lower = center - extent + configured_upper = center + extent + supported_lower = self._tabletop_support_lower_xy + self._target_half[:2] + self.cfg.obstacle_clearance + supported_upper = self._tabletop_support_upper_xy - self._target_half[:2] - self.cfg.obstacle_clearance + lower = torch.maximum(configured_lower, supported_lower) + upper = torch.minimum(configured_upper, supported_upper) + if not bool(torch.all(lower < upper)): + raise RuntimeError("The configured target region has no fully supported tabletop area.") + positions = torch.zeros((count, 3), device=self.device) + positions[:, :2] = lower + torch.rand((count, 2), device=self.device, generator=self.generator) * ( + upper - lower + ) + return positions + + def _sample_near_pour_cup_poses(self, target_positions: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Sample tilted cup centers directly above already-supported receiver poses.""" + count = target_positions.shape[0] + azimuth = torch.empty(count, device=self.device).uniform_(-math.pi, math.pi, generator=self.generator) + tilt = torch.empty(count, device=self.device).uniform_( + *self.cfg.near_pour_tilt_angle_range, + generator=self.generator, + ) + tilt_axes = torch.stack((torch.cos(azimuth), torch.sin(azimuth), torch.zeros_like(azimuth)), dim=-1) + source_quaternions = math_utils.quat_unique(math_utils.quat_from_angle_axis(tilt, tilt_axes)) + radial_unit = torch.rand(count, device=self.device, generator=self.generator) + radial_distance = self.cfg.near_pour_horizontal_radius * torch.sqrt(radial_unit) + radial_angle = torch.empty(count, device=self.device).uniform_( + -math.pi, + math.pi, + generator=self.generator, + ) + cup_centers = target_positions.clone() + cup_centers[:, 0] += radial_distance * torch.cos(radial_angle) + cup_centers[:, 1] += radial_distance * torch.sin(radial_angle) + clearance = torch.empty(count, device=self.device).uniform_( + *self.cfg.near_pour_height_range, + generator=self.generator, + ) + cup_centers[:, 2] += self._target_rim_height + clearance + return source_quaternions, cup_centers + + def _identity_quaternions(self, count: int) -> torch.Tensor: + quaternions = torch.zeros((count, 4), device=self.device) + quaternions[:, 3] = 1.0 + return quaternions + + def _reject(self, current_valid: torch.Tensor, check: torch.Tensor, category: int, reason: str) -> torch.Tensor: + rejected = current_valid & ~check + self._rejection_counts[category][reason] += int(rejected.sum()) + return current_valid & check + + def _task_contract(self) -> dict[str, Any]: + """Return the same canonical contract checked by the runtime loader.""" + return build_franka_pour_reset_task_contract(self.env) + + @staticmethod + def _category_name(category: int, *, near_pour: bool = False) -> str: + if category == GRASPING_CATEGORY: + return "near-pour grasping" if near_pour else "broad grasping" + return "non-grasping" diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/reset_utils.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/reset_utils.py new file mode 100644 index 000000000000..6bf70b91fc43 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/reset_utils.py @@ -0,0 +1,532 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Small tensor utilities shared by Franka Pour reset orchestration.""" + +import math + +import torch + + +def polar_workspace_cells( + nominal_position: tuple[float, float, float] | torch.Tensor, + *, + radius_range: tuple[float, float] | torch.Tensor, + azimuth_half_range: float, + grid_size: int, + device: str | torch.device | None = None, + dtype: torch.dtype | None = None, +) -> torch.Tensor: + """Return deterministic source-workspace positions on a polar grid [m]. + + The polar origin is the fixed robot-base origin in the environment XY plane. When the nominal + radius is strictly inside the configured range, radial samples are piecewise linear from the + lower bound through the nominal radius to the upper bound. When the nominal radius is either + endpoint, a single linear grid spans the range without duplicating a radial sample. Azimuth + samples are symmetric about the nominal XY bearing. The result is flattened radius-major with + azimuth varying fastest; sampling rows uniformly therefore samples discrete rings uniformly + rather than sampling the sector uniformly by area. + + Args: + nominal_position: Authored nominal source position ``(x, y, z)`` [m]. Its XY bearing + centers the sector, its radius is an exact grid sample, and its z coordinate is used + by every returned cell. + radius_range: Strictly increasing positive lower and upper workspace radii [m]. The + nominal radius must lie within their closed interval. + azimuth_half_range: Positive angular half-range about the nominal bearing [rad]. Must be + less than pi radians. + grid_size: Odd number of radial samples and azimuth samples. Must be at least three. + device: Optional output device. Tensor inputs are converted to this device. When omitted, + the device of ``nominal_position`` is preserved, or CPU is used for tuple input. + dtype: Optional floating-point output dtype. When omitted, the dtype of a tensor + ``nominal_position`` is preserved, or the default floating-point dtype is used. + + Returns: + Source positions [m], shape ``(grid_size**2, 3)``. One row is bit-for-bit equal to + ``nominal_position`` after any requested device or dtype conversion. It is the center row + when the nominal radius is strictly interior and lies on the first or last ring when the + nominal radius anchors the corresponding range endpoint. + """ + if isinstance(grid_size, bool) or not isinstance(grid_size, int) or grid_size < 3 or grid_size % 2 == 0: + raise ValueError("grid_size must be an odd integer of at least three.") + if not math.isfinite(azimuth_half_range) or not 0.0 < azimuth_half_range < math.pi: + raise ValueError("azimuth_half_range must be finite and lie in (0, pi).") + + if isinstance(nominal_position, torch.Tensor): + resolved_device = nominal_position.device if device is None else torch.device(device) + resolved_dtype = nominal_position.dtype if dtype is None else dtype + else: + resolved_device = torch.device("cpu") if device is None else torch.device(device) + resolved_dtype = torch.get_default_dtype() if dtype is None else dtype + if not torch.empty((), dtype=resolved_dtype).is_floating_point(): + raise ValueError("nominal_position and the output must use a floating-point dtype.") + + nominal = torch.as_tensor(nominal_position, device=resolved_device, dtype=resolved_dtype) + radii = torch.as_tensor(radius_range, device=resolved_device, dtype=resolved_dtype) + if nominal.shape != (3,): + raise ValueError(f"nominal_position must contain three coordinates, got shape {tuple(nominal.shape)}.") + if radii.shape != (2,): + raise ValueError(f"radius_range must contain two values, got shape {tuple(radii.shape)}.") + if not bool(torch.isfinite(nominal).all()): + raise ValueError("nominal_position must contain finite values.") + if not bool(torch.isfinite(radii).all()): + raise ValueError("radius_range must contain finite values.") + + nominal_radius = torch.linalg.vector_norm(nominal[:2]) + radius_lower, radius_upper = radii.unbind() + if not bool((radius_lower > 0.0) & (radius_lower < radius_upper)): + raise ValueError("radius_range must contain positive values in strictly increasing order.") + if not bool((radius_lower <= nominal_radius) & (nominal_radius <= radius_upper)): + raise ValueError("radius_range must contain the nominal XY radius within its closed interval.") + + half_size = grid_size // 2 + nominal_at_lower = bool(nominal_radius == radius_lower) + nominal_at_upper = bool(nominal_radius == radius_upper) + if nominal_at_lower or nominal_at_upper: + radial_samples = torch.linspace( + float(radius_lower), + float(radius_upper), + grid_size, + device=resolved_device, + dtype=resolved_dtype, + ) + nominal_radius_index = 0 if nominal_at_lower else grid_size - 1 + else: + lower_radii = torch.linspace( + float(radius_lower), + float(nominal_radius), + half_size + 1, + device=resolved_device, + dtype=resolved_dtype, + )[:-1] + upper_radii = torch.linspace( + float(nominal_radius), + float(radius_upper), + half_size + 1, + device=resolved_device, + dtype=resolved_dtype, + ) + radial_samples = torch.cat((lower_radii, upper_radii)) + nominal_radius_index = half_size + nominal_azimuth = torch.atan2(nominal[1], nominal[0]) + azimuth_offsets = torch.linspace( + -azimuth_half_range, + azimuth_half_range, + grid_size, + device=resolved_device, + dtype=resolved_dtype, + ) + radius_grid, azimuth_grid = torch.meshgrid( + radial_samples, + nominal_azimuth + azimuth_offsets, + indexing="ij", + ) + cells = torch.stack( + ( + radius_grid * torch.cos(azimuth_grid), + radius_grid * torch.sin(azimuth_grid), + torch.full_like(radius_grid, nominal[2]), + ), + dim=-1, + ).reshape(-1, 3) + cells[nominal_radius_index * grid_size + half_size] = nominal + return cells + + +def asymmetric_reset_offset_samples( + lower_bound: tuple[float, float, float] | torch.Tensor, + upper_bound: tuple[float, float, float] | torch.Tensor, + sample_count: int, + *, + device: str | torch.device | None = None, + dtype: torch.dtype | None = None, +) -> torch.Tensor: + """Return deterministic, balanced reset-TCP offset samples [m]. + + The first three rows are the exact zero, lower-bound, and upper-bound offsets. Remaining rows + are low-discrepancy interior samples in complementary pairs. In normalized box coordinates, + every interior pair sums to one; in physical coordinates, it sums to the element-wise sum of + ``lower_bound`` and ``upper_bound``. This is balance within an asymmetric box, not a zero-mean + guarantee. + + Args: + lower_bound: Inclusive lower offset bound ``(x, y, z)`` [m]. + upper_bound: Inclusive upper offset bound ``(x, y, z)`` [m]. + sample_count: Odd number of samples. Must be at least three. + device: Optional output device. When omitted, a tensor lower bound supplies the device, + followed by a tensor upper bound; otherwise CPU is used. + dtype: Optional floating-point output dtype. When omitted, a tensor lower bound supplies + the dtype, followed by a tensor upper bound; otherwise the default dtype is used. + + Returns: + Reset-TCP offsets [m], shape ``(sample_count, 3)``. Row zero is exactly zero, rows one and + two exactly span the configured bounds, and later rows are complementary interior pairs. + """ + if isinstance(sample_count, bool) or not isinstance(sample_count, int) or sample_count < 3 or sample_count % 2 == 0: + raise ValueError("sample_count must be an odd integer of at least three.") + + tensor_inputs = tuple(value for value in (lower_bound, upper_bound) if isinstance(value, torch.Tensor)) + if device is None: + if len(tensor_inputs) == 2 and tensor_inputs[0].device != tensor_inputs[1].device: + raise ValueError("Tensor offset bounds must use the same device unless device is specified.") + resolved_device = tensor_inputs[0].device if tensor_inputs else torch.device("cpu") + else: + resolved_device = torch.device(device) + if dtype is None: + if len(tensor_inputs) == 2 and tensor_inputs[0].dtype != tensor_inputs[1].dtype: + raise ValueError("Tensor offset bounds must use the same dtype unless dtype is specified.") + resolved_dtype = tensor_inputs[0].dtype if tensor_inputs else torch.get_default_dtype() + else: + resolved_dtype = dtype + if not torch.empty((), dtype=resolved_dtype).is_floating_point(): + raise ValueError("Offset bounds and the output must use a floating-point dtype.") + + lower = torch.as_tensor(lower_bound, device=resolved_device, dtype=resolved_dtype) + upper = torch.as_tensor(upper_bound, device=resolved_device, dtype=resolved_dtype) + if lower.shape != (3,) or upper.shape != (3,): + raise ValueError( + "lower_bound and upper_bound must each contain three coordinates; " + f"got {tuple(lower.shape)} and {tuple(upper.shape)}." + ) + if not bool(torch.isfinite(lower).all()) or not bool(torch.isfinite(upper).all()): + raise ValueError("Offset bounds must contain finite values.") + if bool(torch.any(lower > 0.0)) or bool(torch.any(upper < 0.0)) or bool(torch.any(lower > upper)): + raise ValueError("Offset bounds must be ordered coordinate-wise and contain zero.") + if not bool(torch.any(lower < upper)): + raise ValueError("Offset bounds must have positive width on at least one coordinate.") + + samples = [torch.zeros((1, 3), device=resolved_device, dtype=resolved_dtype), lower[None], upper[None]] + interior_pair_count = (sample_count - 3) // 2 + if interior_pair_count: + pair_index = torch.arange(interior_pair_count, device=resolved_device, dtype=resolved_dtype) + 0.5 + # Three-dimensional Kronecker sequence used by the legacy reset bank, restricted to the + # open unit cube so endpoint coverage remains the responsibility of the explicit bounds. + multipliers = lower.new_tensor((0.754877666, 0.569840296, 0.438447187)) + unit = torch.frac(pair_index[:, None] * multipliers[None, :]) + epsilon = torch.finfo(resolved_dtype).eps + unit = torch.clamp(unit, min=epsilon, max=1.0 - epsilon) + span = upper - lower + interior = lower + unit * span + complement = lower + (1.0 - unit) * span + samples.append(torch.stack((interior, complement), dim=1).reshape(-1, 3)) + return torch.cat(samples, dim=0) + + +def reset_rotation_vector_samples( + angle_range: tuple[float, float] | torch.Tensor, + sample_count: int, + *, + device: str | torch.device | None = None, + dtype: torch.dtype | None = None, +) -> torch.Tensor: + """Return deterministic reset-orientation perturbations as rotation vectors [rad]. + + Rotation axes follow a Fibonacci sphere, avoiding a preferred roll, pitch, or yaw direction. + Magnitudes cover the configured interval and remain nonzero when its lower bound is positive. + Curriculum code can therefore scale the vectors continuously from an exactly aligned pose at + zero extent to a directionally diverse, deliberately misaligned pose at full extent. + + Args: + angle_range: Inclusive lower and upper rotation magnitudes [rad]. Values must be finite, + nonnegative, ordered, and no greater than pi radians. + sample_count: Number of deterministic samples. Must be positive. + device: Optional output device. A tensor range supplies the default; otherwise CPU is used. + dtype: Optional floating-point output dtype. A tensor range supplies the default; otherwise + the default floating-point dtype is used. + + Returns: + Rotation vectors [rad], shape ``(sample_count, 3)``. For more than one sample, the first and + last rows attain the lower and upper configured magnitudes exactly. + """ + if isinstance(sample_count, bool) or not isinstance(sample_count, int) or sample_count <= 0: + raise ValueError("sample_count must be a positive integer.") + + if isinstance(angle_range, torch.Tensor): + resolved_device = angle_range.device if device is None else torch.device(device) + resolved_dtype = angle_range.dtype if dtype is None else dtype + else: + resolved_device = torch.device("cpu") if device is None else torch.device(device) + resolved_dtype = torch.get_default_dtype() if dtype is None else dtype + if not torch.empty((), dtype=resolved_dtype).is_floating_point(): + raise ValueError("angle_range and the output must use a floating-point dtype.") + + bounds = torch.as_tensor(angle_range, device=resolved_device, dtype=resolved_dtype) + if bounds.shape != (2,): + raise ValueError(f"angle_range must contain two values, got shape {tuple(bounds.shape)}.") + if not bool(torch.isfinite(bounds).all()): + raise ValueError("angle_range must contain finite values.") + lower, upper = bounds.unbind() + if not bool((lower >= 0.0) & (lower <= upper) & (upper <= math.pi)): + raise ValueError("angle_range must be ordered within [0, pi].") + + sample_ids = torch.arange(sample_count, device=resolved_device, dtype=resolved_dtype) + unit_height = (sample_ids + 0.5) / sample_count + axis_z = 1.0 - 2.0 * unit_height + azimuth = 2.0 * math.pi * torch.frac((sample_ids + 0.5) * 0.6180339887498949) + axis_radius = torch.sqrt(torch.clamp(1.0 - axis_z.square(), min=0.0)) + axes = torch.stack( + (axis_radius * torch.cos(azimuth), axis_radius * torch.sin(azimuth), axis_z), + dim=-1, + ) + + if sample_count == 1: + angles = ((lower + upper) * 0.5).reshape(1) + else: + angle_unit = torch.frac((sample_ids + 0.5) * 0.7548776662466927) + angles = lower + angle_unit * (upper - lower) + angles[0] = lower + angles[-1] = upper + return axes * angles.unsqueeze(-1) + + +def boolean_selection_mask(count: int, selected: torch.Tensor) -> torch.Tensor: + """Return a fixed-size boolean mask selecting the supplied indices.""" + if count < 0: + raise ValueError(f"Mask length must be non-negative, got {count}.") + mask = torch.zeros(count, dtype=torch.bool, device=selected.device) + mask[selected.reshape(-1).long()] = True + return mask + + +def balanced_cyclic_permutations(values: torch.Tensor, group_count: int) -> torch.Tensor: + """Return deterministic cyclic permutations with balanced column-wise value counts.""" + if values.ndim != 1 or values.numel() == 0: + raise ValueError(f"values must be a nonempty one-dimensional tensor, got shape {tuple(values.shape)}.") + if group_count < 0: + raise ValueError(f"group_count must be nonnegative, got {group_count}.") + group_ids = torch.arange(group_count, device=values.device).unsqueeze(-1) + value_ids = torch.arange(values.numel(), device=values.device).unsqueeze(0) + return values[(group_ids + value_ids) % values.numel()] + + +def scale_randomization_rows_by_extent( + rows: torch.Tensor, + extent_levels: tuple[float, ...], +) -> torch.Tensor: + """Scale one balanced offset bank independently at every curriculum extent. + + Args: + rows: Full-amplitude offsets whose first dimension is the balanced bank-row dimension. + extent_levels: Strictly increasing amplitude multipliers in ``[0, 1]`` ending at ``1``. + A leading zero provides an exact continuity anchor before randomization is introduced. + + Returns: + Scaled rows with shape ``(len(extent_levels), *rows.shape)``. Every level therefore + preserves the complete row marginal, and the final level exactly preserves ``rows``. + """ + if rows.ndim == 0 or rows.shape[0] == 0: + raise ValueError(f"rows must have a nonempty bank-row dimension, got shape {tuple(rows.shape)}.") + if not torch.is_floating_point(rows) or not bool(torch.isfinite(rows).all()): + raise ValueError("rows must be a finite floating-point tensor.") + if not extent_levels: + raise ValueError("extent_levels must not be empty.") + levels = tuple(float(extent) for extent in extent_levels) + if any(not math.isfinite(extent) or extent < 0.0 or extent > 1.0 for extent in levels): + raise ValueError("extent_levels must contain finite values in [0, 1].") + if any(right <= left for left, right in zip(levels, levels[1:], strict=False)): + raise ValueError("extent_levels must be strictly increasing.") + if not math.isclose(levels[-1], 1.0, rel_tol=0.0, abs_tol=1.0e-9): + raise ValueError("extent_levels must end at 1.0 to preserve the full randomization domain.") + levels = (*levels[:-1], 1.0) + + extent = rows.new_tensor(levels).reshape((len(levels),) + (1,) * rows.ndim) + return rows.unsqueeze(0) * extent + + +def randomization_extent_index_pools( + source_positions: torch.Tensor, + source_yaws: torch.Tensor, + target_positions: torch.Tensor, + tcp_jitter: torch.Tensor, + *, + source_center: tuple[float, float] | torch.Tensor, + source_half_range: tuple[float, float] | torch.Tensor, + source_yaw_half_range: float, + target_center: tuple[float, float] | torch.Tensor, + target_half_range: tuple[float, float] | torch.Tensor, + tcp_jitter_half_range: tuple[float, float, float] | torch.Tensor, + extent_levels: tuple[float, ...], + tolerance: float = 1.0e-6, +) -> tuple[torch.Tensor, ...]: + """Return nested bank indices within combined normalized reset extents. + + Each extent is a Chebyshev radius over source XY position [m], source yaw [rad], target XY + position [m], and TCP jitter [m], normalized by their configured half-ranges. A zero-range axis + contributes zero difficulty at its center and excludes rows displaced from that center. + """ + if source_positions.ndim != 2 or source_positions.shape[-1] < 2: + raise ValueError(f"source_positions must have shape (N, D) with D >= 2, got {tuple(source_positions.shape)}.") + if source_yaws.ndim != 1: + raise ValueError(f"source_yaws must have shape (N,), got {tuple(source_yaws.shape)}.") + if target_positions.ndim != 2 or target_positions.shape[-1] < 2: + raise ValueError(f"target_positions must have shape (N, D) with D >= 2, got {tuple(target_positions.shape)}.") + if tcp_jitter.ndim != 2 or tcp_jitter.shape[-1] != 3: + raise ValueError(f"tcp_jitter must have shape (N, 3), got {tuple(tcp_jitter.shape)}.") + row_count = source_positions.shape[0] + if source_yaws.shape[0] != row_count or target_positions.shape[0] != row_count or tcp_jitter.shape[0] != row_count: + raise ValueError( + "source_positions, source_yaws, target_positions, and tcp_jitter must have the same row count." + ) + if ( + source_yaws.device != source_positions.device + or target_positions.device != source_positions.device + or tcp_jitter.device != source_positions.device + ): + raise ValueError("source_positions, source_yaws, target_positions, and tcp_jitter must be on the same device.") + if not math.isfinite(tolerance) or tolerance < 0.0: + raise ValueError("tolerance must be finite and nonnegative.") + if not extent_levels: + raise ValueError("extent_levels must not be empty.") + levels = tuple(float(extent) for extent in extent_levels) + if any(not math.isfinite(extent) or extent < 0.0 for extent in levels): + raise ValueError("extent_levels must contain finite nonnegative values.") + if any(right <= left for left, right in zip(levels, levels[1:], strict=False)): + raise ValueError("extent_levels must be strictly increasing.") + + def normalized_offsets( + values: torch.Tensor, + center_values: tuple[float, ...] | torch.Tensor, + half_range_values: tuple[float, ...] | torch.Tensor, + name: str, + ) -> torch.Tensor: + center = torch.as_tensor(center_values, device=values.device, dtype=values.dtype) + half_range = torch.as_tensor(half_range_values, device=values.device, dtype=values.dtype) + if center.shape != (values.shape[1],) or half_range.shape != (values.shape[1],): + raise ValueError(f"{name}_center and {name}_half_range must each contain {values.shape[1]} coordinates.") + if bool(torch.any(~torch.isfinite(values))) or bool(torch.any(~torch.isfinite(center))): + raise ValueError(f"{name} values and center must be finite.") + if bool(torch.any(~torch.isfinite(half_range))) or bool(torch.any(half_range < 0.0)): + raise ValueError(f"{name}_half_range must contain finite nonnegative values.") + + offsets = torch.abs(values - center) + positive_range = half_range > 0.0 + result = torch.zeros_like(offsets) + result[:, positive_range] = offsets[:, positive_range] / half_range[positive_range] + if bool(torch.any(~positive_range)): + result[:, ~positive_range] = torch.where( + offsets[:, ~positive_range] <= tolerance, + torch.zeros_like(offsets[:, ~positive_range]), + torch.full_like(offsets[:, ~positive_range], torch.inf), + ) + return result + + normalized_source = normalized_offsets(source_positions[:, :2], source_center, source_half_range, "source") + normalized_source_yaw = normalized_offsets( + source_yaws.unsqueeze(-1), + torch.zeros(1, device=source_yaws.device, dtype=source_yaws.dtype), + torch.as_tensor((source_yaw_half_range,), device=source_yaws.device, dtype=source_yaws.dtype), + "source_yaw", + ) + normalized_target = normalized_offsets(target_positions[:, :2], target_center, target_half_range, "target") + normalized_tcp_jitter = normalized_offsets( + tcp_jitter, + torch.zeros(3, device=tcp_jitter.device, dtype=tcp_jitter.dtype), + tcp_jitter_half_range, + "tcp_jitter", + ) + difficulty = torch.cat( + (normalized_source, normalized_source_yaw, normalized_target, normalized_tcp_jitter), dim=-1 + ).amax(dim=-1) + + pools = tuple(torch.nonzero(difficulty <= extent + tolerance, as_tuple=False).flatten() for extent in levels) + if any(pool.numel() == 0 for pool in pools): + raise ValueError("Every randomization extent level must select at least one bank row.") + return pools + + +def sample_index_pools( + index_pools: tuple[torch.Tensor, ...], + pool_ids: torch.Tensor, + *, + weights: tuple[torch.Tensor, ...] | None = None, +) -> torch.Tensor: + """Sample one global bank index per row from its selected device-resident pool. + + Args: + index_pools: Global bank indices for each pool. + pool_ids: Pool selected by each output row. + weights: Optional nonnegative row weights aligned with :paramref:`index_pools`. Each pool + is sampled proportionally to its weights. If omitted, rows are sampled uniformly. + + Returns: + One sampled global bank index per input row. + """ + if pool_ids.ndim != 1: + raise ValueError(f"pool_ids must be one-dimensional, got shape {tuple(pool_ids.shape)}.") + if weights is not None and len(weights) != len(index_pools): + raise ValueError("weights must contain one tensor per index pool.") + result = torch.empty_like(pool_ids, dtype=torch.long) + for pool_id, index_pool in enumerate(index_pools): + if index_pool.ndim != 1 or index_pool.numel() == 0: + raise ValueError("Every index pool must be a nonempty one-dimensional tensor.") + if index_pool.device != pool_ids.device: + raise ValueError("index pools and pool_ids must be on the same device.") + rows = torch.nonzero(pool_ids == pool_id, as_tuple=False).flatten() + if rows.numel() == 0: + continue + if weights is None: + slots = torch.randint(index_pool.numel(), (rows.numel(),), device=pool_ids.device) + else: + pool_weights = weights[pool_id] + if pool_weights.shape != index_pool.shape or pool_weights.device != pool_ids.device: + raise ValueError("Every weight tensor must match its index pool's shape and device.") + if ( + not torch.is_floating_point(pool_weights) + or not bool(torch.isfinite(pool_weights).all()) + or bool(torch.any(pool_weights < 0.0)) + or not bool(torch.any(pool_weights > 0.0)) + ): + raise ValueError("Pool weights must be finite, nonnegative floating-point values with positive sum.") + slots = torch.multinomial(pool_weights, rows.numel(), replacement=True) + result[rows] = index_pool[slots] + return result + + +def target_xy_behind_source( + source_xy: torch.Tensor, + *, + target_center: tuple[float, float] | torch.Tensor, + target_half_range: tuple[float, float] | torch.Tensor, + minimum_y_separation: float | torch.Tensor, + unit_samples: torch.Tensor, +) -> torch.Tensor: + """Map unit-square samples to target positions safely behind each source cup [m].""" + if source_xy.ndim != 2 or source_xy.shape[-1] != 2: + raise ValueError(f"source_xy must have shape (N, 2), got {tuple(source_xy.shape)}.") + if unit_samples.shape != source_xy.shape: + raise ValueError( + f"unit_samples must match source_xy shape {tuple(source_xy.shape)}, got {tuple(unit_samples.shape)}." + ) + separation = torch.as_tensor(minimum_y_separation, device=source_xy.device, dtype=source_xy.dtype) + if separation.ndim == 0: + separation = separation.expand(source_xy.shape[0]) + elif separation.shape != (source_xy.shape[0],): + raise ValueError("minimum_y_separation must be a scalar or contain one value per source row.") + if bool(torch.any(~torch.isfinite(separation))) or bool(torch.any(separation < 0.0)): + raise ValueError("minimum_y_separation must be finite and nonnegative.") + if bool(torch.any((unit_samples < 0.0) | (unit_samples > 1.0))): + raise ValueError("unit_samples must lie in [0, 1].") + + center = torch.as_tensor(target_center, device=source_xy.device, dtype=source_xy.dtype) + half_range = torch.as_tensor(target_half_range, device=source_xy.device, dtype=source_xy.dtype) + if center.shape != (2,) or half_range.shape != (2,): + raise ValueError("target_center and target_half_range must each contain two coordinates.") + if bool(torch.any(~torch.isfinite(center))) or bool(torch.any(~torch.isfinite(half_range))): + raise ValueError("Target randomization bounds must be finite.") + if bool(torch.any(half_range < 0.0)): + raise ValueError("target_half_range must be nonnegative.") + + lower = center - half_range + upper = center + half_range + allowed_y_upper = torch.minimum( + torch.full_like(source_xy[:, 1], upper[1]), + source_xy[:, 1] - separation, + ) + if bool(torch.any(allowed_y_upper < lower[1])): + raise ValueError("No target y-position satisfies the configured range and source-cup separation.") + + target_x = lower[0] + unit_samples[:, 0] * (upper[0] - lower[0]) + target_y = lower[1] + unit_samples[:, 1] * (allowed_y_upper - lower[1]) + return torch.stack((target_x, target_y), dim=-1) diff --git a/source/isaaclab_tasks/isaaclab_tasks/utils/__init__.pyi b/source/isaaclab_tasks/isaaclab_tasks/utils/__init__.pyi index bfb7af42f762..1108b8bbdc54 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/utils/__init__.pyi +++ b/source/isaaclab_tasks/isaaclab_tasks/utils/__init__.pyi @@ -4,6 +4,13 @@ # SPDX-License-Identifier: BSD-3-Clause __all__ = [ + "AdaptiveResetSampler", + "AdaptiveResetSamplerCfg", + "reset_dataset_collect_batches", + "reset_dataset_content_digest", + "reset_dataset_digest", + "reset_dataset_save_atomic", + "reset_dataset_validate_header", "import_packages", "get_checkpoint_path", "load_cfg_from_registry", @@ -16,7 +23,15 @@ __all__ = [ "setup_preset_cli", ] -from .hydra import PresetCfg, preset, hydra_task_config, resolve_task_config, resolve_presets +from .adaptive_reset_sampler import AdaptiveResetSampler, AdaptiveResetSamplerCfg +from .hydra import PresetCfg, hydra_task_config, preset, resolve_presets, resolve_task_config from .importer import import_packages from .parse_cfg import get_checkpoint_path, load_cfg_from_registry, parse_env_cfg from .preset_cli import setup_preset_cli +from .reset_dataset import ( + reset_dataset_collect_batches, + reset_dataset_content_digest, + reset_dataset_digest, + reset_dataset_save_atomic, + reset_dataset_validate_header, +) diff --git a/source/isaaclab_tasks/isaaclab_tasks/utils/adaptive_reset_sampler.py b/source/isaaclab_tasks/isaaclab_tasks/utils/adaptive_reset_sampler.py new file mode 100644 index 000000000000..9fc0b4822108 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/utils/adaptive_reset_sampler.py @@ -0,0 +1,464 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Adaptive sampling over a difficulty-ordered reset-state cache.""" + +from __future__ import annotations + +import math +from collections.abc import Mapping + +import torch + +from isaaclab.utils.configclass import configclass + +__all__ = ["AdaptiveResetSampler", "AdaptiveResetSamplerCfg"] + + +@configclass +class AdaptiveResetSamplerCfg: + """Configuration for :class:`AdaptiveResetSampler`. + + The sampler targets a mixture of successful and unsuccessful episodes while retaining uniform + replay and a small amount of probing immediately beyond its monotonic difficulty frontier. + """ + + target_success_rate: float = 0.5 + """Desired predicted success rate of sampled resets.""" + + temperature: float = 0.1 + """Minimum softmax temperature used to bound sampling concentration.""" + + history_capacity: int = 32 + """Maximum effective recent outcome count retained for each reset row.""" + + prior_strength: float = 4.0 + """Effective observation count assigned to the cold-start success prior.""" + + initial_frontier_size: int = 128 + """Number of easiest rows initially included in the active frontier.""" + + probe_size: int = 256 + """Number of rows immediately beyond the active frontier eligible for probing.""" + + probe_fraction: float = 0.1 + """Sampling probability reserved for rows beyond the active frontier.""" + + replay_fraction: float = 0.1 + """Uniform replay floor within the active frontier.""" + + frontier_evidence: float = 2.0 + """Excess successful outcomes required to expose one additional reset row.""" + + def __post_init__(self) -> None: + """Validate configuration values.""" + self.validate_values() + + def validate_values(self) -> None: + """Validate the adaptive sampling parameters after runtime overrides.""" + if not math.isfinite(float(self.target_success_rate)) or not 0.0 < self.target_success_rate < 1.0: + raise ValueError("target_success_rate must lie strictly between zero and one.") + for name in ("temperature", "prior_strength", "frontier_evidence"): + value = getattr(self, name) + if isinstance(value, bool) or not math.isfinite(float(value)) or value <= 0.0: + raise ValueError(f"{name} must be finite and positive.") + for name in ("history_capacity", "initial_frontier_size"): + value = getattr(self, name) + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"{name} must be a positive integer.") + if not isinstance(self.probe_size, int) or isinstance(self.probe_size, bool) or self.probe_size < 0: + raise ValueError("probe_size must be a nonnegative integer.") + if not math.isfinite(float(self.probe_fraction)) or not 0.0 <= self.probe_fraction < 1.0: + raise ValueError("probe_fraction must lie in [0, 1).") + if not math.isfinite(float(self.replay_fraction)) or not 0.0 <= self.replay_fraction < 1.0: + raise ValueError("replay_fraction must lie in [0, 1).") + + +class AdaptiveResetSampler: + """Sample reset rows near a requested success rate. + + Reset rows are identified by arbitrary non-negative integer IDs. The caller provides those IDs + in easiest-to-hardest order and reports Boolean episode outcomes using the same IDs. No task or + environment semantics are assumed. + + Outcome counts are kept in bounded, exponentially truncated buffers. Sampling combines a + calibrated softmax over the active frontier, uniform replay, and uniform probes from the next + rows in the difficulty ordering. The frontier can only advance. + + Args: + difficulty_order: Unique raw reset-row IDs ordered from easiest to hardest. + cfg: Sampler configuration. + """ + + _STATE_VERSION = 1 + _BISECTION_STEPS = 24 + _STATE_KEYS = ( + "version", + "difficulty_order", + "effective_successes", + "effective_attempts", + "total_successes", + "total_attempts", + "has_outcome", + "latest_success", + "frontier_size", + "frontier_credit", + ) + + def __init__(self, difficulty_order: torch.Tensor, cfg: AdaptiveResetSamplerCfg | None = None): + self.cfg = cfg if cfg is not None else AdaptiveResetSamplerCfg() + self.cfg.validate_values() + if difficulty_order.ndim != 1 or difficulty_order.numel() == 0: + raise ValueError("difficulty_order must be a non-empty one-dimensional tensor.") + if difficulty_order.dtype != torch.long: + raise TypeError("difficulty_order must have dtype torch.long.") + if bool(torch.any(difficulty_order < 0)): + raise ValueError("difficulty_order must contain non-negative raw row IDs.") + + self._difficulty_order = difficulty_order.detach().clone() + self._sorted_row_ids, self._sorted_to_rank = torch.sort(self._difficulty_order) + if self._sorted_row_ids.numel() > 1 and bool(torch.any(self._sorted_row_ids[1:] == self._sorted_row_ids[:-1])): + raise ValueError("difficulty_order must contain unique raw row IDs.") + + row_count = self._difficulty_order.numel() + device = self._difficulty_order.device + self._effective_successes = torch.zeros(row_count, device=device, dtype=torch.float32) + self._effective_attempts = torch.zeros_like(self._effective_successes) + self._total_successes = torch.zeros(row_count, device=device, dtype=torch.long) + self._total_attempts = torch.zeros_like(self._total_successes) + self._has_outcome = torch.zeros(row_count, device=device, dtype=torch.bool) + self._latest_success = torch.zeros_like(self._has_outcome) + + self._frontier_size = min(self.cfg.initial_frontier_size, row_count) + self._frontier_credit = torch.zeros((), device=device, dtype=torch.float32) + + rank = torch.arange(row_count, device=device, dtype=torch.float32) + remaining = max(row_count - self._frontier_size, 1) + hardness = (rank - self._frontier_size + 1.0).clamp_min(0.0) / remaining + self._prior_success = self.cfg.target_success_rate * (1.0 - hardness).clamp_min(0.0) + self._probabilities_dirty = True + self._probabilities = torch.zeros(row_count, device=device, dtype=torch.float32) + + @property + def difficulty_order(self) -> torch.Tensor: + """Raw row IDs in easiest-to-hardest order.""" + return self._difficulty_order + + @property + def frontier_size(self) -> int: + """Number of rows currently exposed by the monotonic frontier.""" + return self._frontier_size + + @property + def sampling_probabilities(self) -> torch.Tensor: + """Sampling probabilities aligned with :attr:`difficulty_order`.""" + if self._probabilities_dirty: + self._probabilities = self._compute_probabilities() + self._probabilities_dirty = False + return self._probabilities + + @property + def success_estimates(self) -> torch.Tensor: + """Bayesian-smoothed success estimates aligned with :attr:`difficulty_order`.""" + return (self._effective_successes + self.cfg.prior_strength * self._prior_success) / ( + self._effective_attempts + self.cfg.prior_strength + ) + + def sample( + self, + count: int, + forced_row_ids: torch.Tensor | None = None, + *, + generator: torch.Generator | None = None, + ) -> torch.Tensor: + """Sample raw reset-row IDs. + + Args: + count: Number of row IDs to return. + forced_row_ids: Optional row IDs with shape ``(count,)``. Non-negative entries are + returned exactly, including rows outside the current frontier. ``-1`` entries are + sampled adaptively. + generator: Optional random-number generator passed to :func:`torch.multinomial`. + + Returns: + Raw reset-row IDs with shape ``(count,)``. + """ + if count < 0: + raise ValueError("count cannot be negative.") + if forced_row_ids is not None: + if forced_row_ids.shape != (count,) or forced_row_ids.dtype != torch.long: + raise ValueError("forced_row_ids must have shape (count,) and dtype torch.long.") + if forced_row_ids.device != self._difficulty_order.device: + raise ValueError("forced_row_ids must be on the sampler device.") + if bool(torch.any(forced_row_ids < -1)): + raise ValueError("forced_row_ids may contain only known raw row IDs or -1.") + forced = forced_row_ids >= 0 + if bool(torch.any(forced)): + self._row_ids_to_ranks(forced_row_ids[forced]) + else: + forced = None + + if count == 0: + return self._difficulty_order.new_empty(0) + ranks = torch.multinomial(self.sampling_probabilities, count, replacement=True, generator=generator) + row_ids = self._difficulty_order[ranks] + if forced_row_ids is not None: + row_ids = torch.where(forced, forced_row_ids, row_ids) + return row_ids + + def record(self, row_ids: torch.Tensor, successes: torch.Tensor) -> None: + """Record completed episode outcomes and advance the frontier. + + Args: + row_ids: Raw reset-row IDs with shape ``(N,)``. + successes: Boolean success outcomes with shape ``(N,)``. + """ + if row_ids.ndim != 1 or row_ids.dtype != torch.long: + raise ValueError("row_ids must be a one-dimensional torch.long tensor.") + if successes.shape != row_ids.shape or successes.dtype != torch.bool: + raise ValueError("successes must be a Boolean tensor aligned with row_ids.") + if row_ids.device != self._difficulty_order.device or successes.device != row_ids.device: + raise ValueError("row_ids and successes must be on the sampler device.") + if row_ids.numel() == 0: + return + + ranks = self._row_ids_to_ranks(row_ids) + row_count = self._difficulty_order.numel() + batch_attempts = torch.bincount(ranks, minlength=row_count) + batch_successes = torch.bincount(ranks, weights=successes.float(), minlength=row_count) + touched = batch_attempts > 0 + + capacity = float(self.cfg.history_capacity) + kept_batch_attempts = batch_attempts[touched].float().clamp_max(capacity) + kept_batch_scale = kept_batch_attempts / batch_attempts[touched].float() + kept_batch_successes = batch_successes[touched] * kept_batch_scale + old_attempts = self._effective_attempts[touched] + old_scale = ((capacity - kept_batch_attempts) / old_attempts.clamp_min(1.0)).clamp_(0.0, 1.0) + self._effective_successes[touched] = self._effective_successes[touched] * old_scale + kept_batch_successes + self._effective_attempts[touched] = old_attempts * old_scale + kept_batch_attempts + self._total_attempts.add_(batch_attempts) + self._total_successes.add_(torch.bincount(ranks[successes], minlength=row_count)) + + occurrence = torch.arange(row_ids.numel(), device=row_ids.device, dtype=torch.long) + latest_occurrence = torch.full((row_count,), -1, device=row_ids.device, dtype=torch.long) + latest_occurrence.scatter_reduce_(0, ranks, occurrence, reduce="amax", include_self=True) + latest_ranks = torch.nonzero(latest_occurrence >= 0, as_tuple=False).flatten() + self._latest_success[latest_ranks] = successes[latest_occurrence[latest_ranks]] + self._has_outcome[latest_ranks] = True + + self._advance_frontier(ranks, successes) + self._probabilities_dirty = True + + def metrics(self) -> dict[str, float]: + """Return compact sampler metrics using one device-to-host transfer.""" + probabilities = self.sampling_probabilities + success = self.success_estimates + effective_attempts = self._effective_attempts.sum() + values = torch.stack( + ( + torch.as_tensor(self.cfg.target_success_rate, device=success.device), + torch.dot(probabilities, success), + self._effective_successes.sum() / effective_attempts.clamp_min(1.0), + self._latest_success.sum() / success.numel(), + (self._total_successes > 0).sum() / success.numel(), + self._has_outcome.sum() / success.numel(), + torch.as_tensor(self._frontier_size / success.numel(), device=success.device), + probabilities.square().sum().reciprocal(), + ) + ) + names = ( + "target_success_rate", + "predicted_success_rate", + "bounded_success_rate", + "cache_success_rate", + "ever_solved_fraction", + "evaluated_fraction", + "frontier_fraction", + "effective_pool_size", + ) + return dict(zip(names, values.tolist(), strict=True)) + + def state_dict(self) -> dict[str, torch.Tensor]: + """Return a detached copy of the adaptive sampling state.""" + device = self._difficulty_order.device + return { + "version": torch.tensor(self._STATE_VERSION, device=device, dtype=torch.long), + "difficulty_order": self._difficulty_order.detach().clone(), + "effective_successes": self._effective_successes.detach().clone(), + "effective_attempts": self._effective_attempts.detach().clone(), + "total_successes": self._total_successes.detach().clone(), + "total_attempts": self._total_attempts.detach().clone(), + "has_outcome": self._has_outcome.detach().clone(), + "latest_success": self._latest_success.detach().clone(), + "frontier_size": torch.tensor(self._frontier_size, device=device, dtype=torch.long), + "frontier_credit": self._frontier_credit.detach().clone(), + } + + def load_state_dict(self, state_dict: Mapping[str, torch.Tensor]) -> None: + """Restore adaptive sampling state produced by :meth:`state_dict`. + + Args: + state_dict: Mapping containing all sampler state tensors. + """ + expected = set(self._STATE_KEYS) + if set(state_dict) != expected: + missing = sorted(expected - set(state_dict)) + unexpected = sorted(set(state_dict) - expected) + raise ValueError(f"Invalid sampler state keys; missing={missing}, unexpected={unexpected}.") + device = self._difficulty_order.device + self._validate_state_tensor("version", torch.Size([]), torch.long, state_dict) + self._validate_state_tensor("difficulty_order", self._difficulty_order.shape, torch.long, state_dict) + self._validate_state_tensor( + "effective_successes", self._effective_successes.shape, self._effective_successes.dtype, state_dict + ) + self._validate_state_tensor( + "effective_attempts", self._effective_attempts.shape, self._effective_attempts.dtype, state_dict + ) + self._validate_state_tensor("total_successes", self._total_successes.shape, torch.long, state_dict) + self._validate_state_tensor("total_attempts", self._total_attempts.shape, torch.long, state_dict) + self._validate_state_tensor("has_outcome", self._has_outcome.shape, torch.bool, state_dict) + self._validate_state_tensor("latest_success", self._latest_success.shape, torch.bool, state_dict) + self._validate_state_tensor("frontier_size", torch.Size([]), torch.long, state_dict) + self._validate_state_tensor("frontier_credit", torch.Size([]), torch.float32, state_dict) + + version = int(state_dict["version"].to(device=device).item()) + if version != self._STATE_VERSION: + raise ValueError(f"Unsupported adaptive reset sampler state version {version}.") + restored_order = state_dict["difficulty_order"].to(device=device, dtype=torch.long) + if not torch.equal(restored_order, self._difficulty_order): + raise ValueError("Sampler state difficulty_order does not match this reset cache.") + + effective_successes = state_dict["effective_successes"].to(device=device) + effective_attempts = state_dict["effective_attempts"].to(device=device) + total_successes = state_dict["total_successes"].to(device=device) + total_attempts = state_dict["total_attempts"].to(device=device) + invalid_effective = ( + ~torch.isfinite(effective_successes) + | ~torch.isfinite(effective_attempts) + | (effective_successes < 0.0) + | (effective_attempts < effective_successes) + | (effective_attempts > self.cfg.history_capacity) + ) + if bool(torch.any(invalid_effective)): + raise ValueError("Sampler state contains invalid bounded outcome counts.") + if bool(torch.any((total_successes < 0) | (total_attempts < total_successes))): + raise ValueError("Sampler state contains invalid lifetime outcome counts.") + frontier_size = int(state_dict["frontier_size"].to(device=device).item()) + if not 1 <= frontier_size <= self._difficulty_order.numel(): + raise ValueError("Sampler state frontier_size is outside the reset cache.") + frontier_credit = state_dict["frontier_credit"].to(device=device) + if ( + not bool(torch.isfinite(frontier_credit)) + or not 0.0 <= float(frontier_credit.item()) < self.cfg.frontier_evidence + ): + raise ValueError("Sampler state frontier_credit is outside its valid range.") + + self._copy_state_tensor("effective_successes", self._effective_successes, state_dict) + self._copy_state_tensor("effective_attempts", self._effective_attempts, state_dict) + self._copy_state_tensor("total_successes", self._total_successes, state_dict) + self._copy_state_tensor("total_attempts", self._total_attempts, state_dict) + self._copy_state_tensor("has_outcome", self._has_outcome, state_dict) + self._copy_state_tensor("latest_success", self._latest_success, state_dict) + self._frontier_size = frontier_size + self._frontier_credit.copy_(frontier_credit) + self._probabilities_dirty = True + + def _compute_probabilities(self) -> torch.Tensor: + """Build the target-success mixture in difficulty-rank order.""" + success = self.success_estimates + active_rows = torch.arange(self._frontier_size, device=success.device) + active_success = success[active_rows] + probe_size = min(self.cfg.probe_size, success.numel() - self._frontier_size) + probe_fraction = self.cfg.probe_fraction if probe_size > 0 else 0.0 + probabilities = torch.zeros_like(success) + + if probe_size > 0: + probe_rows = torch.arange( + self._frontier_size, + self._frontier_size + probe_size, + device=success.device, + ) + probabilities[probe_rows] = probe_fraction / probe_size + probe_success = success[probe_rows].mean() + active_target = (self.cfg.target_success_rate - probe_fraction * probe_success) / (1.0 - probe_fraction) + else: + active_target = torch.as_tensor(self.cfg.target_success_rate, device=success.device) + + max_inverse_temperature = 1.0 / self.cfg.temperature + lower = torch.full((), -max_inverse_temperature, device=success.device) + upper = torch.full((), max_inverse_temperature, device=success.device) + uniform = torch.full_like(active_success, 1.0 / self._frontier_size) + active_target = active_target.clamp(0.0, 1.0) + for _ in range(self._BISECTION_STEPS): + inverse_temperature = 0.5 * (lower + upper) + softmax = torch.softmax(-inverse_temperature * active_success, dim=0) + candidate = (1.0 - self.cfg.replay_fraction) * softmax + self.cfg.replay_fraction * uniform + predicted = torch.dot(candidate, active_success) + lower = torch.where(predicted > active_target, inverse_temperature, lower) + upper = torch.where(predicted > active_target, upper, inverse_temperature) + + softmax = torch.softmax(-0.5 * (lower + upper) * active_success, dim=0) + active_probabilities = (1.0 - self.cfg.replay_fraction) * softmax + self.cfg.replay_fraction * uniform + probabilities[active_rows] = (1.0 - probe_fraction) * active_probabilities + return probabilities + + def _advance_frontier(self, ranks: torch.Tensor, successes: torch.Tensor) -> None: + """Advance, but never retract, the active difficulty frontier.""" + row_count = self._difficulty_order.numel() + if self._frontier_size >= row_count: + return + window_size = max(self.cfg.probe_size, 1) + frontier_start = max(self._frontier_size - window_size, 0) + frontier_end = min(self._frontier_size + window_size, row_count) + near_frontier = (ranks >= frontier_start) & (ranks < frontier_end) + contribution = torch.where( + near_frontier, + successes.float() - self.cfg.target_success_rate, + torch.zeros_like(successes, dtype=torch.float32), + ).sum() + self._frontier_credit.add_(contribution).clamp_min_(0.0) + advance = min( + int(torch.floor(self._frontier_credit / self.cfg.frontier_evidence).item()), + row_count - self._frontier_size, + ) + if advance > 0: + self._frontier_size += advance + self._frontier_credit.sub_(advance * self.cfg.frontier_evidence) + + def _row_ids_to_ranks(self, row_ids: torch.Tensor) -> torch.Tensor: + """Resolve raw row IDs to positions in the difficulty ordering.""" + positions = torch.searchsorted(self._sorted_row_ids, row_ids) + in_range = positions < self._sorted_row_ids.numel() + safe_positions = positions.clamp_max(self._sorted_row_ids.numel() - 1) + valid = in_range & (self._sorted_row_ids[safe_positions] == row_ids) + if not bool(torch.all(valid)): + invalid = row_ids[~valid].detach().cpu().tolist() + raise ValueError(f"Unknown raw reset-row IDs: {invalid}.") + return self._sorted_to_rank[safe_positions] + + @staticmethod + def _validate_state_tensor( + name: str, + shape: torch.Size, + dtype: torch.dtype, + state_dict: Mapping[str, torch.Tensor], + ) -> None: + """Validate the metadata of one serialized state tensor.""" + source = state_dict[name] + if not isinstance(source, torch.Tensor) or source.shape != shape or source.dtype != dtype: + source_shape = source.shape if isinstance(source, torch.Tensor) else None + source_dtype = source.dtype if isinstance(source, torch.Tensor) else None + raise ValueError( + f"Sampler state {name!r} has shape/dtype {source_shape}/{source_dtype}; expected {shape}/{dtype}." + ) + + @staticmethod + def _copy_state_tensor( + name: str, + target: torch.Tensor, + state_dict: Mapping[str, torch.Tensor], + ) -> None: + """Validate and copy one state tensor.""" + target.copy_(state_dict[name].to(device=target.device)) diff --git a/source/isaaclab_tasks/isaaclab_tasks/utils/reset_dataset.py b/source/isaaclab_tasks/isaaclab_tasks/utils/reset_dataset.py new file mode 100644 index 000000000000..36e97f6c238f --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/utils/reset_dataset.py @@ -0,0 +1,295 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Task-agnostic utilities for generating and persisting reset-state datasets.""" + +from __future__ import annotations + +import hashlib +import os +import tempfile +from collections.abc import Callable, Mapping +from io import BytesIO +from pathlib import Path +from typing import Any, TypeVar + +import torch + +_BatchT = TypeVar("_BatchT") + + +class _HashWriter: + """Adapt a hashlib digest to the byte-writer interface used by the encoder.""" + + def __init__(self) -> None: + self.digest = hashlib.sha256() + + def write(self, value: bytes) -> None: + """Append bytes to the digest.""" + self.digest.update(value) + + +def _write_bytes(sink: Any, value: bytes) -> None: + """Write one length-delimited byte string to a hash or byte buffer.""" + sink.write(len(value).to_bytes(8, byteorder="big", signed=False)) + sink.write(value) + + +def _write_value(sink: Any, value: Any) -> None: + """Write the canonical representation of one supported value.""" + if value is None: + sink.write(b"none") + elif isinstance(value, bool): + sink.write(b"bool") + sink.write(b"1" if value else b"0") + elif isinstance(value, int): + sink.write(b"int") + _write_bytes(sink, str(value).encode("ascii")) + elif isinstance(value, float): + sink.write(b"float") + _write_bytes(sink, value.hex().encode("ascii")) + elif isinstance(value, str): + sink.write(b"str") + _write_bytes(sink, value.encode("utf-8")) + elif isinstance(value, bytes): + sink.write(b"bytes") + _write_bytes(sink, value) + elif isinstance(value, torch.Tensor): + if value.layout != torch.strided or value.is_quantized: + raise TypeError("Reset-dataset hashes support only dense, strided tensors.") + tensor = value.detach().cpu().contiguous() + sink.write(b"tensor") + _write_bytes(sink, str(tensor.dtype).encode("ascii")) + _write_value(sink, tuple(tensor.shape)) + raw_bytes = tensor.reshape(-1).view(torch.uint8).numpy().tobytes() + _write_bytes(sink, raw_bytes) + elif isinstance(value, Mapping): + sink.write(b"mapping") + encoded_items: list[tuple[bytes, Any]] = [] + for key, item in value.items(): + key_buffer = BytesIO() + _write_value(key_buffer, key) + encoded_items.append((key_buffer.getvalue(), item)) + encoded_items.sort(key=lambda pair: pair[0]) + _write_value(sink, len(encoded_items)) + for encoded_key, item in encoded_items: + _write_bytes(sink, encoded_key) + _write_value(sink, item) + elif isinstance(value, tuple): + sink.write(b"tuple") + _write_value(sink, len(value)) + for item in value: + _write_value(sink, item) + elif isinstance(value, list): + sink.write(b"list") + _write_value(sink, len(value)) + for item in value: + _write_value(sink, item) + else: + raise TypeError(f"Unsupported reset-dataset hash value type: {type(value).__name__}.") + + +def reset_dataset_digest(value: Any) -> str: + """Return a stable SHA-256 digest for nested primitive and tensor data. + + Mapping insertion order and tensor device do not affect the result. Concrete sequence types, + tensor dtypes, and tensor shapes remain part of the digest. + + Args: + value: Nested mappings, lists, tuples, primitive values, or dense tensors. + + Returns: + The lowercase hexadecimal SHA-256 digest. + + Raises: + TypeError: If the value contains an unsupported type or tensor layout. + """ + writer = _HashWriter() + _write_value(writer, value) + return writer.digest.hexdigest() + + +def reset_dataset_content_digest( + payload: Mapping[str, Any], + *, + digest_key: str = "content_sha256", +) -> str: + """Return a dataset payload digest while excluding its own digest field. + + Args: + payload: Dataset payload to hash. + digest_key: Top-level field that stores the resulting content digest. + + Returns: + The content SHA-256 digest. + """ + if not isinstance(payload, Mapping): + raise TypeError("Reset dataset payload must be a mapping.") + if not isinstance(digest_key, str) or not digest_key: + raise ValueError("digest_key must be a non-empty string.") + return reset_dataset_digest({key: value for key, value in payload.items() if key != digest_key}) + + +def reset_dataset_validate_header( + payload: Mapping[str, Any], + *, + expected_format: str, + expected_schema_version: int, + expected_contract: Any, +) -> tuple[Mapping[str, Any], Mapping[str, Any]]: + """Validate the common envelope and integrity hashes of a reset dataset. + + Task-specific validators remain responsible for state keys, tensor shapes, physical + invariants, and metadata semantics. + + Args: + payload: Dataset payload to validate. + expected_format: Exact dataset format identifier. + expected_schema_version: Exact schema version supported by the caller. + expected_contract: Generator and task contract represented by ``contract_sha256``. + + Returns: + The validated metadata and states mappings. + + Raises: + TypeError: If the payload, metadata, or states are not mappings. + ValueError: If a header field or integrity digest does not match. + """ + if not isinstance(payload, Mapping): + raise TypeError("Reset dataset payload must be a mapping.") + if not isinstance(expected_format, str) or not expected_format: + raise ValueError("expected_format must be a non-empty string.") + if not isinstance(expected_schema_version, int) or isinstance(expected_schema_version, bool): + raise TypeError("expected_schema_version must be an integer.") + if payload.get("format") != expected_format: + raise ValueError(f"Expected reset dataset format {expected_format!r}.") + if payload.get("schema_version") != expected_schema_version: + raise ValueError(f"Expected reset dataset schema version {expected_schema_version}.") + + metadata = payload.get("metadata") + states = payload.get("states") + if not isinstance(metadata, Mapping): + raise TypeError("Reset dataset metadata must be a mapping.") + if not isinstance(states, Mapping): + raise TypeError("Reset dataset states must be a mapping.") + + expected_contract_digest = reset_dataset_digest(expected_contract) + if payload.get("contract_sha256") != expected_contract_digest: + raise ValueError("Reset dataset contract digest does not match the expected contract.") + if payload.get("content_sha256") != reset_dataset_content_digest(payload): + raise ValueError("Reset dataset content digest does not match its payload.") + return metadata, states + + +def reset_dataset_save_atomic( + payload: Mapping[str, Any], + output_path: str | Path, + *, + validator: Callable[[Mapping[str, Any]], None], +) -> Path: + """Validate and atomically save a reset dataset with :func:`torch.save`. + + Args: + payload: Complete dataset payload. + output_path: Destination file path. + validator: Dataset-specific validator called before the destination is modified. + + Returns: + The resolved destination path. + """ + if not callable(validator): + raise TypeError("validator must be callable.") + validator(payload) + path = Path(output_path).expanduser().resolve() + path.parent.mkdir(parents=True, exist_ok=True) + descriptor, temporary_name = tempfile.mkstemp(prefix=f".{path.name}.", suffix=".tmp", dir=path.parent) + os.close(descriptor) + temporary_path = Path(temporary_name) + try: + torch.save(dict(payload), temporary_path) + temporary_path.replace(path) + finally: + temporary_path.unlink(missing_ok=True) + return path + + +def reset_dataset_collect_batches( + target_count: int, + *, + batch_size: int, + max_candidate_count: int, + evaluate_batch: Callable[[range], _BatchT], + batch_count: Callable[[_BatchT], int], + batch_slice: Callable[[_BatchT, int], _BatchT], +) -> tuple[list[_BatchT], int]: + """Collect exactly the requested accepted samples from rejection-sampled batches. + + ``evaluate_batch`` receives a unique, contiguous range of candidate IDs. It owns proposal + generation and validation, and returns only accepted samples. Keeping batch representation + behind callbacks lets tasks retain efficient tensor dictionaries or custom batch types. + + Args: + target_count: Number of accepted samples to collect. + batch_size: Maximum number of candidates evaluated per callback. + max_candidate_count: Maximum candidates evaluated before failing. + evaluate_batch: Callback that proposes and validates one candidate-ID range. + batch_count: Callback returning the accepted sample count in a batch. + batch_slice: Callback retaining the first requested samples of a batch. + + Returns: + A list of accepted batches and the number of candidates evaluated. + + Raises: + ValueError: If counts are invalid or a callback reports an impossible batch size. + RuntimeError: If the candidate budget is exhausted before reaching ``target_count``. + """ + for name, value in ( + ("target_count", target_count), + ("batch_size", batch_size), + ("max_candidate_count", max_candidate_count), + ): + if not isinstance(value, int) or isinstance(value, bool): + raise TypeError(f"{name} must be an integer.") + if target_count < 0: + raise ValueError("target_count must be nonnegative.") + if batch_size <= 0: + raise ValueError("batch_size must be positive.") + if max_candidate_count < target_count: + raise ValueError("max_candidate_count cannot be smaller than target_count.") + if target_count == 0: + return [], 0 + + batches: list[_BatchT] = [] + accepted_total = 0 + evaluated_total = 0 + while accepted_total < target_count and evaluated_total < max_candidate_count: + candidate_count = min(batch_size, max_candidate_count - evaluated_total) + candidate_ids = range(evaluated_total, evaluated_total + candidate_count) + batch = evaluate_batch(candidate_ids) + accepted_count = batch_count(batch) + if not isinstance(accepted_count, int) or isinstance(accepted_count, bool): + raise TypeError("batch_count must return an integer.") + if not 0 <= accepted_count <= candidate_count: + raise ValueError( + f"Accepted batch count {accepted_count} is outside the candidate range [0, {candidate_count}]." + ) + evaluated_total += candidate_count + if accepted_count == 0: + continue + remaining = target_count - accepted_total + if accepted_count > remaining: + batch = batch_slice(batch, remaining) + accepted_count = batch_count(batch) + if accepted_count != remaining: + raise ValueError("batch_slice did not return the requested accepted sample count.") + batches.append(batch) + accepted_total += accepted_count + + if accepted_total != target_count: + raise RuntimeError( + f"Reset-dataset rejection sampling accepted {accepted_total}/{target_count} samples " + f"after evaluating {evaluated_total} candidates." + ) + return batches, evaluated_total diff --git a/source/isaaclab_tasks/test/contrib/test_franka_pour_cube_bowl_mesh.py b/source/isaaclab_tasks/test/contrib/test_franka_pour_cube_bowl_mesh.py new file mode 100644 index 000000000000..351e7bea505d --- /dev/null +++ b/source/isaaclab_tasks/test/contrib/test_franka_pour_cube_bowl_mesh.py @@ -0,0 +1,82 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Unit tests for the Franka pour hollow-cube bowl mesh (no simulator).""" + +import ast +from pathlib import Path + +import numpy as np +import pytest + +from isaaclab_tasks.contrib.franka_pour import cube_bowl_mesh +from isaaclab_tasks.contrib.franka_pour.cube_bowl_mesh import cube_bowl_inner_bounds, make_cube_bowl_mesh + +DIMS = dict( + inner_width=0.037, + inner_depth=0.037, + cavity_depth=0.045, + wall_thickness=0.009, + bottom_thickness=0.009, +) + + +def test_returns_float32_flat_int32_arrays(): + v, f = make_cube_bowl_mesh(**DIMS) + assert v.dtype == np.float32 and v.ndim == 2 and v.shape[1] == 3 + assert f.dtype == np.int32 and f.ndim == 1 and f.size % 3 == 0 + assert int(f.max()) < len(v) + + +def test_is_watertight_outward_manifold(): + # validate=True runs the task-local validator and raises if the shell is not a closed, + # consistently-wound manifold. + make_cube_bowl_mesh(**DIMS, validate=True) + + +def test_mesh_validation_is_task_local_and_rejects_open_shells(): + source_path = Path(cube_bowl_mesh.__file__) + tree = ast.parse(source_path.read_text(encoding="utf-8"), filename=str(source_path)) + imported_modules = { + node.module for node in ast.walk(tree) if isinstance(node, ast.ImportFrom) and node.module is not None + } + assert not any("franka_scoop" in module for module in imported_modules) + + validator = getattr(cube_bowl_mesh, "_validate_closed_oriented_mesh", None) + assert callable(validator) + _, indices = make_cube_bowl_mesh(**DIMS, validate=False) + with pytest.raises(RuntimeError, match="not watertight"): + validator(indices[:-3], "Open cube bowl") + + +def test_outward_orientation_positive_signed_volume(): + v, f = make_cube_bowl_mesh(**DIMS) + tris = v[f.reshape(-1, 3)] + signed_vol = float(np.einsum("ij,ij->i", tris[:, 0], np.cross(tris[:, 1], tris[:, 2])).sum()) + assert signed_vol > 0.0 + + +def test_outer_footprint_fits_gripper(): + v, _ = make_cube_bowl_mesh(**DIMS) + outer = float(v[:, 0].max() - v[:, 0].min()) + assert abs(outer - (DIMS["inner_width"] + 2 * DIMS["wall_thickness"])) < 1e-6 + assert outer <= 0.06 # fits the ~0.08 m Franka opening with closure margin + + +def test_height_and_floor(): + v, _ = make_cube_bowl_mesh(**DIMS) + assert abs(float(v[:, 2].min()) - 0.0) < 1e-6 + assert abs(float(v[:, 2].max()) - (DIMS["bottom_thickness"] + DIMS["cavity_depth"])) < 1e-6 + + +def test_inner_bounds_inside_outer(): + lo, hi = cube_bowl_inner_bounds( + DIMS["inner_width"], DIMS["inner_depth"], DIMS["cavity_depth"], DIMS["bottom_thickness"] + ) + assert np.all(hi - lo > 0) + assert abs(float(hi[0] - lo[0]) - DIMS["inner_width"]) < 1e-6 + assert abs(float(hi[1] - lo[1]) - DIMS["inner_depth"]) < 1e-6 + assert abs(float(lo[2]) - DIMS["bottom_thickness"]) < 1e-6 + assert abs(float(hi[2]) - (DIMS["bottom_thickness"] + DIMS["cavity_depth"])) < 1e-6 diff --git a/source/isaaclab_tasks/test/contrib/test_franka_pour_cube_bowl_spawner.py b/source/isaaclab_tasks/test/contrib/test_franka_pour_cube_bowl_spawner.py new file mode 100644 index 000000000000..3b5b6d755d2c --- /dev/null +++ b/source/isaaclab_tasks/test/contrib/test_franka_pour_cube_bowl_spawner.py @@ -0,0 +1,213 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Stage-authoring tests for the Franka pour cube-bowl spawner.""" + +import subprocess +import sys +import textwrap + +from isaaclab.app import AppLauncher + +# Launch Omniverse before importing simulator or USD modules. +simulation_app = AppLauncher(headless=True).app + +import numpy as np +import pytest + +from pxr import UsdGeom, UsdPhysics, UsdShade + +import isaaclab.sim as sim_utils +from isaaclab.sim import SimulationCfg, SimulationContext +from isaaclab.sim.spawners.materials import RigidBodyMaterialBaseCfg + +from isaaclab_tasks.contrib.franka_pour.cube_bowl_mesh import make_cube_bowl_mesh +from isaaclab_tasks.contrib.franka_pour.cube_bowl_spawner_cfg import CubeBowlSpawnerCfg + +pytestmark = pytest.mark.isaacsim_ci + +_BOWL_DIMS = { + "inner_width": 0.037, + "inner_depth": 0.039, + "cavity_depth": 0.045, + "wall_thickness": 0.009, + "bottom_thickness": 0.008, +} + + +@pytest.fixture +def sim(): + """Create a fresh stage and simulation context for each test.""" + sim_utils.create_new_stage() + sim = SimulationContext(SimulationCfg(dt=0.01)) + sim_utils.update_stage() + yield sim + sim._disable_app_control_on_stop_handle = True # prevent timeout + sim.stop() + sim.clear_instance() + + +def _make_cfg(**kwargs) -> CubeBowlSpawnerCfg: + """Build a bowl config with task-realistic dimensions.""" + return CubeBowlSpawnerCfg(**_BOWL_DIMS, **kwargs) + + +def _quat_xyzw(prim) -> tuple[float, float, float, float]: + """Read a prim's authored local orientation in Isaac Lab order.""" + quat = prim.GetAttribute("xformOp:orient").Get() + imaginary = quat.GetImaginary() + return (float(imaginary[0]), float(imaginary[1]), float(imaginary[2]), float(quat.GetReal())) + + +def test_config_import_and_instantiation_do_not_require_physx(): + """The task-local config remains importable when the optional PhysX package is absent.""" + code = textwrap.dedent( + """ + import builtins + import sys + + class _PhysxBlocker: + def find_spec(self, name, path=None, target=None): + if name == "isaaclab_physx" or name.startswith("isaaclab_physx."): + raise ImportError(f"blocked optional PhysX import: {name}") + return None + + sys.meta_path.insert(0, _PhysxBlocker()) + builtins._isaaclab_tasks_registered = True + + from isaaclab.sim.spawners.materials import RigidBodyMaterialBaseCfg + from isaaclab_tasks.contrib.franka_pour.cube_bowl_spawner_cfg import CubeBowlSpawnerCfg + + cfg = CubeBowlSpawnerCfg( + inner_width=0.037, + inner_depth=0.039, + cavity_depth=0.045, + wall_thickness=0.009, + bottom_thickness=0.008, + physics_material=RigidBodyMaterialBaseCfg(static_friction=0.6, dynamic_friction=0.5), + ) + assert isinstance(cfg.physics_material, RigidBodyMaterialBaseCfg) + assert cfg.physics_material.static_friction == 0.6 + assert "isaaclab_physx" not in sys.modules + """ + ) + + result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=False) + + assert result.returncode == 0, result.stdout + result.stderr + + +def test_dynamic_source_authors_visual_mesh_grasp_proxy_and_material(sim): + """The source bowl has a visual shell and an invisible rigid grasp proxy.""" + half_extents = (0.028, 0.029, 0.031) + color = (0.15, 0.55, 0.85) + cfg = _make_cfg( + display_color=color, + grasp_proxy_half_extents=half_extents, + mass_props=sim_utils.MassPropertiesCfg(mass=0.24), + rigid_props=sim_utils.RigidBodyPropertiesCfg(rigid_body_enabled=True), + collision_props=sim_utils.CollisionPropertiesCfg(collision_enabled=True), + physics_material=RigidBodyMaterialBaseCfg( + static_friction=0.83, + dynamic_friction=0.71, + restitution=0.04, + ), + ) + + root = cfg.func("/World/Cup", cfg) + stage = sim.stage + mesh_prim = stage.GetPrimAtPath("/World/Cup/geometry/mesh") + proxy_prim = stage.GetPrimAtPath("/World/Cup/geometry/grasp_proxy") + + assert root.IsValid() + assert root.HasAPI(UsdPhysics.RigidBodyAPI) + assert root.HasAPI(UsdPhysics.MassAPI) + assert root.GetAttribute("physics:mass").Get() == pytest.approx(0.24) + + assert mesh_prim.IsA(UsdGeom.Mesh) + assert not mesh_prim.HasAPI(UsdPhysics.CollisionAPI) + mesh = UsdGeom.Mesh(mesh_prim) + expected_points, expected_indices = make_cube_bowl_mesh(**_BOWL_DIMS) + np.testing.assert_allclose(np.asarray(mesh.GetPointsAttr().Get()), expected_points) + np.testing.assert_array_equal(np.asarray(mesh.GetFaceVertexIndicesAttr().Get()), expected_indices) + assert len(mesh.GetPointsAttr().Get()) == 16 + assert list(mesh.GetFaceVertexCountsAttr().Get()) == [3] * (expected_indices.size // 3) + assert tuple(mesh.GetDisplayColorAttr().Get()[0]) == pytest.approx(color) + + visual_binding = UsdShade.MaterialBindingAPI(mesh_prim).GetDirectBinding() + assert visual_binding.GetMaterialPath() == "/World/Cup/geometry/visual_material" + + assert proxy_prim.IsA(UsdGeom.Cube) + assert proxy_prim.HasAPI(UsdPhysics.CollisionAPI) + assert not proxy_prim.HasAPI(UsdPhysics.RigidBodyAPI) + assert UsdGeom.Imageable(proxy_prim).ComputeVisibility() == UsdGeom.Tokens.invisible + assert UsdGeom.Cube(proxy_prim).GetSizeAttr().Get() == pytest.approx(1.0) + assert [tuple(point) for point in UsdGeom.Cube(proxy_prim).GetExtentAttr().Get()] == pytest.approx( + [(-0.5, -0.5, -0.5), (0.5, 0.5, 0.5)] + ) + assert tuple(proxy_prim.GetAttribute("xformOp:translate").Get()) == pytest.approx((0.0, 0.0, half_extents[2])) + assert tuple(proxy_prim.GetAttribute("xformOp:scale").Get()) == pytest.approx( + tuple(2.0 * value for value in half_extents) + ) + + material_path = "/World/Cup/geometry/material" + material_prim = stage.GetPrimAtPath(material_path) + assert material_prim.IsA(UsdShade.Material) + assert material_prim.HasAPI(UsdPhysics.MaterialAPI) + assert material_prim.GetAttribute("physics:staticFriction").Get() == pytest.approx(0.83) + physics_binding = UsdShade.MaterialBindingAPI(proxy_prim).GetDirectBinding("physics") + assert physics_binding.GetMaterialPath() == material_path + + +def test_kinematic_target_without_proxy_authors_pose_and_display_color(sim): + """The target bowl remains a valid kinematic rigid body without a grasp proxy.""" + translation = (0.42, -0.17, 0.09) + orientation = (0.0, 0.0, 0.38268343, 0.92387953) + color = (0.92, 0.31, 0.12) + cfg = _make_cfg( + display_color=color, + grasp_proxy_half_extents=None, + rigid_props=sim_utils.RigidBodyPropertiesCfg(rigid_body_enabled=True, kinematic_enabled=True), + physics_material=RigidBodyMaterialBaseCfg(static_friction=0.64, dynamic_friction=0.52), + ) + + root = cfg.func("/World/TargetCup", cfg, translation=translation, orientation=orientation) + mesh_prim = sim.stage.GetPrimAtPath("/World/TargetCup/geometry/mesh") + + assert root.IsValid() + assert root.HasAPI(UsdPhysics.RigidBodyAPI) + assert UsdPhysics.RigidBodyAPI(root).GetKinematicEnabledAttr().Get() is True + assert tuple(root.GetAttribute("xformOp:translate").Get()) == pytest.approx(translation) + assert _quat_xyzw(root) == pytest.approx(orientation) + assert not sim.stage.GetPrimAtPath("/World/TargetCup/geometry/grasp_proxy").IsValid() + assert mesh_prim.IsA(UsdGeom.Mesh) + assert not mesh_prim.HasAPI(UsdPhysics.CollisionAPI) + assert tuple(UsdGeom.Mesh(mesh_prim).GetDisplayColorAttr().Get()[0]) == pytest.approx(color) + physics_binding = UsdShade.MaterialBindingAPI(root).GetDirectBinding("physics") + assert physics_binding.GetMaterialPath() == "/World/TargetCup/geometry/material" + + +def test_rejects_existing_root_prim(sim): + """Spawning never mutates an existing root prim.""" + sim_utils.create_prim("/World/ExistingCup", "Xform") + + with pytest.raises(ValueError, match="already exists"): + _make_cfg().func("/World/ExistingCup", _make_cfg()) + + +def test_clone_decorator_spawns_under_matching_parents(sim): + """A regex path clones the authored bowl hierarchy into every matching parent.""" + sim_utils.create_prim("/World/env_0", "Xform") + sim_utils.create_prim("/World/env_1", "Xform") + cfg = _make_cfg(rigid_props=sim_utils.RigidBodyPropertiesCfg(rigid_body_enabled=True)) + + source = cfg.func("/World/env_.*/Cup", cfg) + + assert str(source.GetPath()) == "/World/env_0/Cup" + for env_index in range(2): + root = sim.stage.GetPrimAtPath(f"/World/env_{env_index}/Cup") + mesh = sim.stage.GetPrimAtPath(f"/World/env_{env_index}/Cup/geometry/mesh") + assert root.HasAPI(UsdPhysics.RigidBodyAPI) + assert mesh.IsA(UsdGeom.Mesh) diff --git a/source/isaaclab_tasks/test/contrib/test_franka_pour_curriculum.py b/source/isaaclab_tasks/test/contrib/test_franka_pour_curriculum.py new file mode 100644 index 000000000000..b0d495f2f9e7 --- /dev/null +++ b/source/isaaclab_tasks/test/contrib/test_franka_pour_curriculum.py @@ -0,0 +1,862 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Unit tests for the Franka Pour curriculum and reset-relative actions.""" + +from types import SimpleNamespace + +import pytest +import torch + +from isaaclab.managers import CurriculumTermCfg + +from isaaclab_tasks.contrib.franka_pour import pour_env as pour_env_module +from isaaclab_tasks.contrib.franka_pour.mdp.actions import ( + CurriculumGripperPositionAction, + CurriculumJointPositionAction, + TrajectoryJointPositionAction, + _bilateral_gripper_preload, +) +from isaaclab_tasks.contrib.franka_pour.mdp.curriculums import PourCurriculum +from isaaclab_tasks.contrib.franka_pour.pour_env import FrankaPourEnv + +_STAGE_NAMES = ( + "drain", + "deep_tilt", + "tilt", + "pour", + "near_carry", + "mid_carry", + "carry", + "grasp", + "approach_1", + "approach_2", + "approach_3", + "approach_4", + "approach_5", + "approach_6", + "full", + "randomized", +) +_RANDOMIZED_STAGE = len(_STAGE_NAMES) - 1 +_EXTENT_LEVELS = (0.0, 0.50, 1.0) + + +class FakeCurriculumEnv: + """Minimal vectorized environment state consumed by :class:`PourCurriculum`.""" + + def __init__( + self, + *, + frozen: bool = False, + start_stage: int = 0, + start_randomization_level: int = 0, + replay_fraction: float = 0.0, + entry_replay_fraction: float | None = None, + extent_levels: tuple[float, ...] = _EXTENT_LEVELS, + ): + self.num_envs = 4 + self.device = "cpu" + self.cfg = SimpleNamespace( + curriculum_stage_names=_STAGE_NAMES, + curriculum_target_frac=( + 0.05, + 0.08, + 0.10, + 0.15, + 0.15, + 0.18, + 0.20, + 0.30, + 0.30, + 0.30, + 0.30, + 0.30, + 0.30, + 0.30, + 0.30, + 0.30, + ), + curriculum_start_stage=start_stage, + curriculum_randomization_extent_levels=extent_levels, + curriculum_independent_arm_fraction_levels=tuple( + index / max(len(extent_levels) - 1, 1) for index in range(len(extent_levels)) + ), + curriculum_independent_target_fraction_levels=tuple( + index / max(len(extent_levels) - 1, 1) for index in range(len(extent_levels)) + ), + curriculum_randomization_start_level=start_randomization_level, + curriculum_freeze=frozen, + curriculum_success_threshold=0.75, + curriculum_randomization_promotion_threshold=0.65, + curriculum_min_resets_per_stage=2, + curriculum_min_reset_cohorts_per_stage=0.0, + curriculum_previous_stage_replay_fraction=replay_fraction, + curriculum_frontier_entry_replay_fraction=( + replay_fraction if entry_replay_fraction is None else entry_replay_fraction + ), + ) + self.curriculum_stage = torch.zeros(self.num_envs, dtype=torch.long) + self.curriculum_randomization_level = torch.zeros(self.num_envs, dtype=torch.long) + self.pour_target_frac = torch.zeros(self.num_envs) + self.episode_succeeded = torch.zeros(self.num_envs, dtype=torch.bool) + self.episode_length_buf = torch.zeros(self.num_envs, dtype=torch.long) + self.ep_max_target_frac = torch.zeros(self.num_envs) + + def set_curriculum_stage(self, env_ids, stage: int) -> None: + self.curriculum_stage[env_ids] = stage + self.pour_target_frac[env_ids] = self.cfg.curriculum_target_frac[stage] + + def set_curriculum_randomization_level(self, env_ids, level: int) -> None: + self.curriculum_randomization_level[env_ids] = level + + +def test_curriculum_ignores_initial_reset_and_advances_only_reset_worlds(): + env = FakeCurriculumEnv() + term = PourCurriculum(CurriculumTermCfg(func=PourCurriculum), env) + + initial = term(env, torch.arange(env.num_envs)) + assert set(initial) == { + "stage", + "randomization_level", + "success_rate", + "completed_episodes", + "required_completed_episodes", + "mastered", + } + assert term.resets_in_stage == 0 + + env.episode_length_buf[:2] = 10 + env.episode_succeeded[:2] = True + env.ep_max_target_frac[:2] = torch.tensor([0.4, 0.5]) + metrics = term(env, torch.tensor([0, 1])) + + assert term.stage == 1 + assert term.resets_in_stage == 0 + assert env.curriculum_stage.tolist() == [1, 1, 0, 0] + assert env.pour_target_frac.tolist() == pytest.approx([0.08, 0.08, 0.05, 0.05]) + assert metrics == pytest.approx( + { + "stage": 1.0, + "randomization_level": 0.0, + "success_rate": 0.0, + "completed_episodes": 0.0, + "required_completed_episodes": 2.0, + "mastered": 0.0, + "mean_peak_target_frac": 0.45, + } + ) + + +def test_curriculum_requires_configured_number_of_completed_episodes(): + env = FakeCurriculumEnv() + env.cfg.curriculum_min_resets_per_stage = 500 + term = PourCurriculum(CurriculumTermCfg(func=PourCurriculum), env) + env.episode_length_buf[:] = 10 + env.episode_succeeded[:] = True + + for _ in range(124): + term(env, torch.arange(env.num_envs)) + term(env, torch.tensor([0, 1, 2])) + assert term.resets_in_stage == 499 + assert term.stage == 0 + + term(env, torch.tensor([3])) + assert term.stage == 1 + assert term.resets_in_stage == 0 + + +def test_curriculum_requires_environment_scaled_reset_cohorts(): + env = FakeCurriculumEnv() + env.cfg.curriculum_min_resets_per_stage = 2 + env.cfg.curriculum_min_reset_cohorts_per_stage = 3.0 + term = PourCurriculum(CurriculumTermCfg(func=PourCurriculum), env) + env.episode_length_buf[:] = 10 + env.episode_succeeded[:] = True + + metrics = term(env, torch.arange(env.num_envs)) + assert term.stage == 0 + assert term.resets_in_stage == 4 + assert metrics["required_completed_episodes"] == 12.0 + + term(env, torch.arange(env.num_envs)) + assert term.stage == 0 + assert term.resets_in_stage == 8 + + term(env, torch.arange(env.num_envs)) + assert term.stage == 1 + assert term.resets_in_stage == 0 + + +def test_lagging_old_stage_episode_does_not_change_new_stage_statistics(): + env = FakeCurriculumEnv() + term = PourCurriculum(CurriculumTermCfg(func=PourCurriculum), env) + env.episode_length_buf[:2] = 10 + env.episode_succeeded[:2] = True + term(env, torch.tensor([0, 1])) + + env.episode_length_buf[2] = 10 + env.episode_succeeded[2] = True + term(env, torch.tensor([2])) + + assert env.curriculum_stage.tolist() == [1, 1, 1, 0] + assert term.stage == 1 + assert term.resets_in_stage == 0 + assert term.success_rate == 0.0 + + +def test_curriculum_replays_previous_stage_without_counting_it(monkeypatch): + env = FakeCurriculumEnv(start_stage=1, replay_fraction=0.5) + term = PourCurriculum(CurriculumTermCfg(func=PourCurriculum), env) + draws = iter((torch.tensor([0.1, 0.9, 0.2, 0.8]), torch.tensor([0.9, 0.9]))) + monkeypatch.setattr(torch, "rand", lambda *args, **kwargs: next(draws)) + + term(env, torch.arange(env.num_envs)) + assert env.curriculum_stage.tolist() == [0, 1, 0, 1] + + env.episode_length_buf[[0, 2]] = 10 + env.episode_succeeded[[0, 2]] = True + term(env, torch.tensor([0, 2])) + assert term.resets_in_stage == 0 + + +def test_curriculum_decays_entry_replay_toward_retention_floor(monkeypatch): + env = FakeCurriculumEnv( + start_stage=1, + replay_fraction=0.1, + entry_replay_fraction=0.5, + ) + env.cfg.curriculum_min_resets_per_stage = 4 + term = PourCurriculum(CurriculumTermCfg(func=PourCurriculum), env) + draws = iter((torch.tensor([0.3, 0.45, 0.55, 0.8]), torch.tensor([0.2, 0.35]))) + monkeypatch.setattr(torch, "rand", lambda *args, **kwargs: next(draws)) + + term(env, torch.arange(env.num_envs)) + assert term._previous_frontier_replay_fraction(env, 4) == pytest.approx(0.5) + assert env.curriculum_stage.tolist() == [0, 0, 1, 1] + + env.episode_length_buf[2:] = 10 + env.episode_succeeded[2:] = False + term(env, torch.tensor([2, 3])) + + assert term.resets_in_stage == 2 + assert term._previous_frontier_replay_fraction(env, 4) == pytest.approx(0.3) + assert env.curriculum_stage[2:].tolist() == [0, 1] + + +def test_randomized_curriculum_replays_previous_extent_without_counting_it(monkeypatch): + frontier = len(_EXTENT_LEVELS) // 2 + env = FakeCurriculumEnv( + start_stage=_RANDOMIZED_STAGE, + start_randomization_level=frontier, + replay_fraction=0.5, + ) + term = PourCurriculum(CurriculumTermCfg(func=PourCurriculum), env) + draws = iter((torch.tensor([0.1, 0.9, 0.2, 0.8]), torch.tensor([0.9, 0.9]))) + monkeypatch.setattr(torch, "rand", lambda *args, **kwargs: next(draws)) + + term(env, torch.arange(env.num_envs)) + assert env.curriculum_stage.tolist() == [_RANDOMIZED_STAGE] * env.num_envs + assert env.curriculum_randomization_level.tolist() == [frontier - 1, frontier, frontier - 1, frontier] + + env.episode_length_buf[[0, 2]] = 10 + env.episode_succeeded[[0, 2]] = True + term(env, torch.tensor([0, 2])) + assert term.resets_in_stage == 0 + + +def test_frozen_curriculum_stays_at_configured_stage(monkeypatch): + max_randomization_level = len(_EXTENT_LEVELS) - 1 + env = FakeCurriculumEnv( + frozen=True, + start_stage=_RANDOMIZED_STAGE, + start_randomization_level=max_randomization_level, + replay_fraction=0.5, + ) + term = PourCurriculum(CurriculumTermCfg(func=PourCurriculum), env) + + monkeypatch.setattr(torch, "rand", lambda *args, **kwargs: pytest.fail("frozen curriculum sampled replay")) + term(env, torch.arange(env.num_envs)) + assert env.curriculum_stage.tolist() == [_RANDOMIZED_STAGE] * env.num_envs + assert env.curriculum_randomization_level.tolist() == [max_randomization_level] * env.num_envs + + +def test_curriculum_success_window_weights_each_completed_episode_equally(): + env = FakeCurriculumEnv(frozen=True) + env.cfg.curriculum_min_resets_per_stage = 5 + term = PourCurriculum(CurriculumTermCfg(func=PourCurriculum), env) + env.episode_length_buf[:] = 10 + + env.episode_succeeded[:] = True + term(env, torch.arange(env.num_envs)) + env.episode_succeeded[0] = False + metrics = term(env, torch.tensor([0])) + + assert term.resets_in_stage == 5 + assert term.success_rate == pytest.approx(0.8) + assert metrics["success_rate"] == pytest.approx(0.8) + + +def test_curriculum_success_window_evicts_oldest_completed_episodes(): + env = FakeCurriculumEnv(frozen=True) + env.cfg.curriculum_min_resets_per_stage = 4 + term = PourCurriculum(CurriculumTermCfg(func=PourCurriculum), env) + env.episode_length_buf[:] = 10 + + env.episode_succeeded[:] = True + term(env, torch.arange(env.num_envs)) + env.episode_succeeded[:2] = False + term(env, torch.tensor([0, 1])) + + assert term.resets_in_stage == 6 + assert term.success_rate == pytest.approx(0.5) + + +def test_curriculum_promotes_at_exact_window_threshold(): + env = FakeCurriculumEnv() + env.cfg.curriculum_min_resets_per_stage = 5 + env.cfg.curriculum_success_threshold = 0.8 + term = PourCurriculum(CurriculumTermCfg(func=PourCurriculum), env) + env.episode_length_buf[:] = 10 + + env.episode_succeeded[:] = True + term(env, torch.arange(env.num_envs)) + env.episode_succeeded[0] = False + term(env, torch.tensor([0])) + + assert term.stage == 1 + assert term.resets_in_stage == 0 + assert term.success_rate == 0.0 + + +def test_randomization_frontier_uses_lower_promotion_threshold_without_lowering_mastery(): + env = FakeCurriculumEnv(start_stage=_RANDOMIZED_STAGE) + env.cfg.curriculum_min_resets_per_stage = 5 + env.cfg.curriculum_success_threshold = 0.8 + env.cfg.curriculum_randomization_promotion_threshold = 0.6 + term = PourCurriculum(CurriculumTermCfg(func=PourCurriculum), env) + env.episode_length_buf[:] = 10 + + env.episode_succeeded[:3] = True + env.episode_succeeded[3] = False + term(env, torch.arange(env.num_envs)) + env.episode_succeeded[0] = False + term(env, torch.tensor([0])) + + assert term.randomization_level == 1 + assert term.resets_in_stage == 0 + assert term.success_rate == 0.0 + + +def test_curriculum_mastery_requires_a_full_success_window(): + env = FakeCurriculumEnv( + frozen=True, + start_stage=_RANDOMIZED_STAGE, + start_randomization_level=len(_EXTENT_LEVELS) - 1, + ) + env.cfg.curriculum_min_resets_per_stage = 4 + term = PourCurriculum(CurriculumTermCfg(func=PourCurriculum), env) + env.episode_length_buf[:] = 10 + + env.episode_succeeded[0] = True + metrics = term(env, torch.tensor([0])) + assert metrics["success_rate"] == 1.0 + assert metrics["mastered"] == 0.0 + + env.episode_succeeded[1:4] = torch.tensor([True, True, False]) + metrics = term(env, torch.tensor([1, 2, 3])) + assert metrics["success_rate"] == pytest.approx(0.75) + assert metrics["mastered"] == 1.0 + + +@pytest.mark.parametrize( + "extent_levels", + [ + (1.0,), + (0.4, 0.7, 1.0), + _EXTENT_LEVELS, + ], +) +def test_final_stage_advances_nested_randomization_frontiers_before_mastery(extent_levels): + env = FakeCurriculumEnv(start_stage=_RANDOMIZED_STAGE - 1, extent_levels=extent_levels) + env.cfg.curriculum_success_threshold = 1.0 + term = PourCurriculum(CurriculumTermCfg(func=PourCurriculum), env) + env.episode_length_buf[:] = 10 + env.episode_succeeded[:] = True + + metrics = term(env, torch.tensor([0, 1])) + assert term.stage == _RANDOMIZED_STAGE + assert term.randomization_level == 0 + assert term.resets_in_stage == 0 + assert metrics["mastered"] == 0.0 + assert env.curriculum_randomization_level[:2].tolist() == [0, 0] + + # Reset one lagging full-task world onto level zero without counting its old episode. + term(env, torch.tensor([2])) + assert env.curriculum_stage[2].item() == _RANDOMIZED_STAGE + assert env.curriculum_randomization_level[2].item() == 0 + + # Advance through every configured randomization frontier. At each promotion, one in-flight + # episode from the preceding level must be ignored and then retagged to the new frontier. + for next_level in range(1, len(env.cfg.curriculum_randomization_extent_levels)): + metrics = term(env, torch.tensor([0, 1])) + assert term.stage == _RANDOMIZED_STAGE + assert term.randomization_level == next_level + assert term.resets_in_stage == 0 + assert term.success_rate == 0.0 + assert metrics["mastered"] == 0.0 + assert env.curriculum_randomization_level[:2].tolist() == [next_level, next_level] + + assert env.curriculum_randomization_level[2].item() == next_level - 1 + term(env, torch.tensor([2])) + assert term.randomization_level == next_level + assert term.resets_in_stage == 0 + assert term.success_rate == 0.0 + assert env.curriculum_randomization_level[2].item() == next_level + + metrics = term(env, torch.tensor([0, 1])) + assert term.stage == _RANDOMIZED_STAGE + assert term.randomization_level == len(env.cfg.curriculum_randomization_extent_levels) - 1 + assert term.resets_in_stage == 2 + assert term.success_rate == 1.0 + assert metrics["mastered"] == 1.0 + + +def test_curriculum_rejects_nonpositive_success_window(): + env = FakeCurriculumEnv() + env.cfg.curriculum_min_resets_per_stage = 0 + + with pytest.raises(ValueError, match="curriculum_min_resets_per_stage must be positive"): + PourCurriculum(CurriculumTermCfg(func=PourCurriculum), env) + + +def test_final_randomization_samples_arm_and_target_independently_with_clearance(monkeypatch): + """The final frontier must break paired reset rows while retaining conservative clearance.""" + env = SimpleNamespace( + device="cpu", + cfg=SimpleNamespace( + curriculum_independent_sample_attempts=3, + curriculum_independent_arm_min_tcp_distance=0.5, + curriculum_independent_arm_fraction_levels=(1.0,), + curriculum_independent_target_fraction_levels=(1.0,), + source_cup_inner_width=0.10, + source_cup_inner_depth=0.10, + source_cup_wall_thickness=0.01, + target_cup_inner_width=0.12, + target_cup_inner_depth=0.12, + target_cup_wall_thickness=0.01, + curriculum_randomized_cup_clearance=0.02, + ), + _select_first_safe_candidate=FrankaPourEnv._select_first_safe_candidate, + ) + positions = torch.tensor( + ( + (0.0, 0.0, 0.0), + (1.0, 0.0, 0.0), + (2.0, 0.0, 0.0), + (3.0, 0.0, 0.0), + ) + ) + env._randomized_source_pos_bank_t = positions + env._randomized_source_yaw_bank_t = torch.zeros(4) + env._randomized_tcp_pos_bank_t = positions + env._randomized_target_pos_bank_t = positions + env._randomized_extent_index_pools = (torch.arange(4),) + env._randomized_extent_index_weights = (torch.ones(4),) + + source_indices = torch.arange(4) + independent_indices = torch.remainder(source_indices + 1, 4) + env._independent_arm_fallback_index_t = independent_indices + env._independent_target_fallback_index_t = independent_indices + env._independent_target_clearance = lambda source_rows, target_rows: FrankaPourEnv._independent_target_clearance( + env, + source_rows, + target_rows, + ) + monkeypatch.setattr( + pour_env_module, + "sample_index_pools", + lambda index_pools, pool_ids, *, weights=None: source_indices.clone(), + ) + + arm_indices, target_indices = FrankaPourEnv._sample_independent_reset_indices( + env, + source_indices, + torch.zeros(4, dtype=torch.long), + ) + + torch.testing.assert_close(arm_indices, independent_indices) + torch.testing.assert_close(target_indices, independent_indices) + assert bool(torch.all(arm_indices != source_indices)) + assert bool(torch.all(target_indices != source_indices)) + + +def test_curriculum_joint_action_offset_updates_selected_worlds_only(): + action = CurriculumJointPositionAction.__new__(CurriculumJointPositionAction) + action._offset = torch.zeros((4, 7)) + target = torch.arange(14, dtype=torch.float32).reshape(2, 7) + + action.set_action_offset(target, env_ids=torch.tensor([1, 3])) + + torch.testing.assert_close(action.action_offset[[1, 3]], target) + torch.testing.assert_close(action.action_offset[[0, 2]], torch.zeros((2, 7))) + with pytest.raises(ValueError, match="shape"): + action.set_action_offset(torch.zeros((1, 7)), env_ids=torch.tensor([1, 3])) + + +def test_trajectory_phase_action_modulates_around_nominal_speed(): + phase_action = torch.tensor([-1.0, 0.0, 1.0]) + + phase_speed = TrajectoryJointPositionAction._phase_speed_command(phase_action) + + torch.testing.assert_close(phase_speed, torch.tensor([0.75, 1.0, 1.25])) + + +def test_phase_gate_never_rewinds_a_later_curriculum_reset(): + current = torch.tensor([0.0, 0.10, 0.30, 0.46, 0.66]) + + approach_limit = TrajectoryJointPositionAction._monotonic_gate_limit(current, 0.10) + grasp_limit = TrajectoryJointPositionAction._monotonic_gate_limit(current, 0.30) + + torch.testing.assert_close(approach_limit, torch.tensor([0.10, 0.10, 0.30, 0.46, 0.66])) + torch.testing.assert_close(grasp_limit, torch.tensor([0.30, 0.30, 0.30, 0.46, 0.66])) + + +def test_trajectory_approach_gate_accepts_axial_standoff_and_rejects_cross_track_error(): + action = TrajectoryJointPositionAction.__new__(TrajectoryJointPositionAction) + axial_error = torch.tensor([-0.12, -0.12]) + cross_track_error = torch.tensor([0.0, 0.011]) + action._env = SimpleNamespace( + grasp_approach_error=lambda: (axial_error, cross_track_error), + cup_velocity_w=lambda: torch.zeros((2, 6)), + ) + action._asset = SimpleNamespace(data=SimpleNamespace(joint_pos=SimpleNamespace(torch=torch.zeros((2, 7))))) + action._joint_ids = list(range(7)) + action._processed_actions = torch.zeros((2, 7)) + action._approach_max_lateral_distance = 0.01 + action._approach_max_joint_error = 0.08 + action._approach_max_linear_velocity = 0.01 + action._approach_max_angular_velocity = 0.1 + + assert action._approach_ready().tolist() == [True, False] + + +def test_carry_stage_requires_fresh_grasp_dwell_despite_starting_after_grasp_waypoint(): + action = TrajectoryJointPositionAction.__new__(TrajectoryJointPositionAction) + action._env = SimpleNamespace(curriculum_stage=torch.tensor([0, 1])) + action._waypoint_count = 7 + action._num_joints = 7 + action._grasp_gate_stage = 1 + action._approach_phase = 0.10 + action._grasp_phase = 0.30 + action._lift_phase = 0.46 + action._align_phase = 0.66 + action._reference_waypoints = torch.zeros((2, 7, 7)) + action._reference_phase = torch.zeros(2) + action._minimum_phase = torch.zeros(2) + action._grasp_dwell_count = torch.zeros(2, dtype=torch.long) + action._approach_dwell_count = torch.zeros(2, dtype=torch.long) + action._approach_unlocked = torch.zeros(2, dtype=torch.bool) + action._grasp_unlocked = torch.zeros(2, dtype=torch.bool) + action._lift_unlocked = torch.zeros(2, dtype=torch.bool) + action._align_unlocked = torch.zeros(2, dtype=torch.bool) + action._processed_actions = torch.zeros((2, 7)) + action._filtered_residual = torch.zeros((2, 7)) + phase = torch.tensor([0.46, 0.46]) + + action.set_reference( + torch.zeros((2, 7, 7)), + phase, + torch.zeros((2, 7)), + ) + + assert action._grasp_unlocked.tolist() == [True, False] + assert action._lift_unlocked.tolist() == [True, False] + assert action._align_unlocked.tolist() == [False, False] + + +def test_bilateral_gripper_preload_rejects_unilateral_empty_transient_and_open_states(): + target = torch.tensor( + [ + [0.024, 0.024], + [0.024, 0.024], + [0.024, 0.024], + [0.024, 0.024], + [0.040, 0.040], + [0.024, 0.024], + ] + ) + position = torch.tensor( + [ + [0.026, 0.026], + [0.027, 0.0242], + [0.024, 0.024], + [0.027, 0.027], + [0.042, 0.042], + [float("nan"), 0.026], + ] + ) + velocity = torch.zeros_like(position) + velocity[3] = 0.04 + + deflection, bilateral = _bilateral_gripper_preload( + position, + velocity, + target, + min_deflection=0.001, + max_velocity=0.005, + max_command=0.025, + ) + + assert bilateral.tolist() == [True, False, False, False, False, False] + torch.testing.assert_close(deflection[0], torch.tensor([0.002, 0.002])) + torch.testing.assert_close(deflection[-1], torch.tensor([0.0, 0.002])) + + +def test_curriculum_joint_action_smooths_targets_and_reset_clears_history(): + action = CurriculumJointPositionAction.__new__(CurriculumJointPositionAction) + action.cfg = SimpleNamespace(clip=None) + action._raw_actions = torch.zeros((2, 2)) + action._processed_actions = torch.zeros((2, 2)) + action._previous_target = torch.zeros((2, 2)) + action._offset = torch.zeros((2, 2)) + action._scale = 0.5 + action._alpha = 0.2 + action._project_reference_through_stage = -1 + action._reference_action_magnitude = 1.0 + action._reference_action_index = 0 + action._reference_target = None + + action.process_actions(torch.ones((2, 2))) + torch.testing.assert_close(action.processed_actions, torch.full((2, 2), 0.1)) + action.process_actions(torch.ones((2, 2))) + torch.testing.assert_close(action.processed_actions, torch.full((2, 2), 0.18)) + + action.set_action_offset(torch.tensor([[0.3, 0.4]]), env_ids=torch.tensor([1])) + action.reset(torch.tensor([1])) + torch.testing.assert_close(action.processed_actions[1], torch.tensor([0.3, 0.4])) + torch.testing.assert_close(action._previous_target[1], torch.tensor([0.3, 0.4])) + + +def test_curriculum_joint_action_projects_only_early_stage_onto_validated_segment(): + action = CurriculumJointPositionAction.__new__(CurriculumJointPositionAction) + action.cfg = SimpleNamespace(clip=None) + action._env = SimpleNamespace(curriculum_stage=torch.tensor([0, 1, 0])) + action._raw_actions = torch.zeros((3, 2)) + action._processed_actions = torch.zeros((3, 2)) + action._previous_target = torch.zeros((3, 2)) + action._offset = torch.zeros((3, 2)) + action._scale = torch.tensor([[2.0, 1.0]]).repeat(3, 1) + action._alpha = 1.0 + action._project_reference_through_stage = 0 + action._reference_action_magnitude = 1.0 + action._reference_action_index = 0 + action._reference_target = torch.tensor([[2.0, 2.0]]).repeat(3, 1) + + # The first coordinate is a stage-stable scalar phase; stage one keeps the normal full-rank action. + raw = torch.tensor([[0.2, 0.9], [0.4, -0.2], [-1.0, 2.0]]) + action.process_actions(raw) + + torch.testing.assert_close(action.raw_actions, raw) + torch.testing.assert_close(action.processed_actions[0], torch.tensor([0.4, 0.4])) + torch.testing.assert_close(action.processed_actions[1], torch.tensor([0.8, -0.2])) + torch.testing.assert_close(action.processed_actions[2], torch.zeros(2)) + + # A later low command cannot reverse an early-stage pour, while the unrestricted stage still + # follows its ordinary reset-relative joint command. + action.process_actions(torch.zeros_like(raw)) + torch.testing.assert_close(action.processed_actions[0], torch.tensor([0.4, 0.4])) + torch.testing.assert_close(action.processed_actions[1], torch.zeros(2)) + + +def test_curriculum_gripper_zero_action_tracks_nominal_preload_after_reset(): + action = CurriculumGripperPositionAction.__new__(CurriculumGripperPositionAction) + action._env = SimpleNamespace(device="cpu") + action._raw_actions = torch.zeros((4, 1)) + action._processed_actions = torch.zeros((4, 2)) + action._action_offset = torch.full((4, 1), 0.024) + action._scale = 0.001 + action._alpha = 1.0 + action._use_incremental_target = False + action._binary_threshold = None + action._close_position = 0.024 + action._neutral_position = 0.025 + action._open_position = 0.04 + action._force_open_stage = -1 + action._capture_unlocked = torch.ones(4, dtype=torch.bool) + action._capture_dwell_count = torch.zeros(4, dtype=torch.long) + action._num_joints = 2 + + action.set_reset_position(torch.tensor([[0.04], [0.04]]), env_ids=torch.tensor([1, 3])) + action.reset(torch.arange(4)) + action.process_actions(torch.zeros((4, 1))) + + torch.testing.assert_close(action.processed_actions, torch.full((4, 2), 0.024)) + action.process_actions(torch.tensor([[0.0], [-0.25], [0.0], [1.0]])) + torch.testing.assert_close(action.processed_actions[1], torch.full((2,), 0.024)) + torch.testing.assert_close(action.processed_actions[3], torch.full((2,), 0.025)) + + +def test_curriculum_gripper_incremental_target_holds_open_or_preloaded_state(): + action = CurriculumGripperPositionAction.__new__(CurriculumGripperPositionAction) + action._env = SimpleNamespace(device="cpu") + action._raw_actions = torch.zeros((2, 1)) + action._processed_actions = torch.full((2, 2), 0.024) + action._action_offset = torch.full((2, 1), 0.024) + action._scale = 0.004 + action._alpha = 0.2 + action._use_incremental_target = True + action._binary_threshold = None + action._close_position = 0.021 + action._neutral_position = 0.04 + action._open_position = 0.04 + action._force_open_stage = -1 + action._capture_unlocked = torch.ones(2, dtype=torch.bool) + action._capture_dwell_count = torch.zeros(2, dtype=torch.long) + action._num_joints = 2 + + action.set_reset_position(torch.tensor([[0.04], [0.024]])) + action.reset() + action.process_actions(torch.zeros((2, 1))) + torch.testing.assert_close(action.processed_actions, torch.tensor([[0.04, 0.04], [0.024, 0.024]])) + + action.process_actions(torch.tensor([[-1.0], [0.0]])) + torch.testing.assert_close(action.processed_actions, torch.tensor([[0.0392, 0.0392], [0.024, 0.024]])) + action.process_actions(torch.zeros((2, 1))) + torch.testing.assert_close(action.processed_actions, torch.tensor([[0.0392, 0.0392], [0.024, 0.024]])) + + +def test_curriculum_gripper_binary_action_filters_close_and_open_targets(): + action = CurriculumGripperPositionAction.__new__(CurriculumGripperPositionAction) + action._env = SimpleNamespace(device="cpu") + action._raw_actions = torch.zeros((5, 1)) + action._processed_actions = torch.full((5, 2), 0.03) + action._action_offset = torch.full((5, 1), 0.024) + action._scale = 0.016 + action._alpha = 0.2 + action._use_incremental_target = False + action._binary_threshold = 0.0 + action._close_position = 0.021 + action._neutral_position = 0.04 + action._open_position = 0.04 + action._force_open_stage = -1 + action._capture_unlocked = torch.ones(5, dtype=torch.bool) + action._capture_dwell_count = torch.zeros(5, dtype=torch.long) + action._num_joints = 2 + + raw = torch.tensor([[-1.0], [-1.0e-6], [0.0], [1.0e-6], [1.0]]) + action.process_actions(raw) + + torch.testing.assert_close(action.raw_actions, raw) + torch.testing.assert_close( + action.processed_actions[:, 0], + torch.tensor([0.0282, 0.0282, 0.032, 0.032, 0.032]), + ) + torch.testing.assert_close(action.processed_actions[:, 0], action.processed_actions[:, 1]) + assert action.action_dim == 1 + + +def test_curriculum_gripper_action_filters_bounded_position_residual(): + action = CurriculumGripperPositionAction.__new__(CurriculumGripperPositionAction) + action._env = SimpleNamespace() + action._raw_actions = torch.zeros((2, 1)) + action._processed_actions = torch.full((2, 2), 0.024) + action._action_offset = torch.full((2, 1), 0.024) + action._scale = 0.001 + action._alpha = 0.2 + action._use_incremental_target = False + action._binary_threshold = None + action._close_position = 0.024 + action._neutral_position = 0.025 + action._open_position = 0.04 + action._force_open_stage = -1 + action._capture_unlocked = torch.ones(2, dtype=torch.bool) + action._capture_dwell_count = torch.zeros(2, dtype=torch.long) + action._num_joints = 2 + + action.process_actions(torch.ones((2, 1))) + torch.testing.assert_close(action.processed_actions, torch.full((2, 2), 0.0242)) + action.process_actions(torch.ones((2, 1))) + torch.testing.assert_close(action.processed_actions, torch.full((2, 2), 0.02436)) + + action.reset(torch.tensor([1])) + torch.testing.assert_close(action.processed_actions[1], torch.full((2,), 0.02436)) + + +def test_curriculum_gripper_caps_policy_opening_at_safe_preload_in_every_stage(): + action = CurriculumGripperPositionAction.__new__(CurriculumGripperPositionAction) + action._env = SimpleNamespace(step_dt=1.0 / 60.0, curriculum_stage=torch.tensor([0, 2, 3, 4])) + action._raw_actions = torch.zeros((4, 1)) + action._processed_actions = torch.full((4, 2), 0.025) + action._action_offset = torch.full((4, 1), 0.024) + action._scale = 0.001 + action._alpha = 0.2 + action._use_incremental_target = False + action._binary_threshold = None + action._close_position = 0.024 + action._neutral_position = 0.025 + action._open_position = 0.04 + action._force_open_stage = -1 + action._capture_unlocked = torch.ones(4, dtype=torch.bool) + action._capture_dwell_count = torch.zeros(4, dtype=torch.long) + action._num_joints = 2 + + action.process_actions(torch.ones((4, 1))) + + torch.testing.assert_close(action.raw_actions, torch.ones((4, 1))) + torch.testing.assert_close( + action.processed_actions, + torch.tensor( + [ + [0.025, 0.025], + [0.025, 0.025], + [0.025, 0.025], + [0.025, 0.025], + ] + ), + ) + + +def test_curriculum_gripper_capture_requires_near_zero_axial_and_cross_track_error(): + arm_action = SimpleNamespace( + reference_error=torch.zeros((3, 7)), + reference_phase=torch.ones(3), + ) + axial_error = torch.tensor([0.004, -0.12, 0.0]) + cross_track_error = torch.tensor([0.003, 0.0, 0.011]) + action = CurriculumGripperPositionAction.__new__(CurriculumGripperPositionAction) + action._env = SimpleNamespace( + action_manager=SimpleNamespace(get_term=lambda name: arm_action), + curriculum_stage=torch.full((3,), 3, dtype=torch.long), + grasp_approach_error=lambda: (axial_error, cross_track_error), + cup_velocity_w=lambda: torch.zeros((3, 6)), + ) + action._raw_actions = torch.zeros((3, 1)) + action._processed_actions = torch.full((3, 2), 0.04) + action._action_offset = torch.full((3, 1), 0.024) + action._scale = 0.001 + action._alpha = 1.0 + action._use_incremental_target = False + action._binary_threshold = None + action._close_position = 0.024 + action._neutral_position = 0.025 + action._open_position = 0.04 + action._force_open_stage = 2 + action._force_open_phase = 0.30 + action._capture_max_lateral_distance = 0.005 + action._capture_max_vertical_distance = 0.008 + action._capture_max_joint_error = 0.08 + action._capture_dwell_steps = 1 + action._capture_max_linear_velocity = 0.02 + action._capture_max_angular_velocity = 0.2 + action._capture_unlocked = torch.zeros(3, dtype=torch.bool) + action._capture_dwell_count = torch.zeros(3, dtype=torch.long) + action._num_joints = 2 + + action.process_actions(torch.zeros((3, 1))) + + assert action._capture_unlocked.tolist() == [True, False, False] + torch.testing.assert_close(action.processed_actions[0], torch.full((2,), 0.024)) + torch.testing.assert_close(action.processed_actions[1:], torch.full((2, 2), 0.04)) diff --git a/source/isaaclab_tasks/test/contrib/test_franka_pour_env_cfg.py b/source/isaaclab_tasks/test/contrib/test_franka_pour_env_cfg.py new file mode 100644 index 000000000000..3595e9059887 --- /dev/null +++ b/source/isaaclab_tasks/test/contrib/test_franka_pour_env_cfg.py @@ -0,0 +1,2095 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Construction tests for the Franka pour env config (no simulator).""" + +import ast +import math +import subprocess +import sys +from pathlib import Path +from types import SimpleNamespace + +import gymnasium as gym +import pytest +import torch +from isaaclab_newton.assets import MPMObjectCfg +from isaaclab_newton.sim.schemas import MujocoJointCfg + +from isaaclab.assets import RigidObjectCfg +from isaaclab.managers import CurriculumTermCfg, RewardTermCfg, SceneEntityCfg, TerminationTermCfg +from isaaclab.sim.schemas import MassCfg, UsdPhysicsCollisionCfg, UsdPhysicsRigidBodyCfg +from isaaclab.sim.spawners.materials import RigidBodyMaterialBaseCfg + +from isaaclab_contrib.coupling import CouplerProxyCfg, NewtonCoupler + +import isaaclab_tasks.contrib.franka_pour as franka_pour +import isaaclab_tasks.contrib.franka_pour.config.franka # noqa: F401 +import isaaclab_tasks.contrib.franka_pour.mdp as mdp +import isaaclab_tasks.contrib.franka_pour.pour_env_cfg as pour_env_cfg +from isaaclab_tasks.contrib.franka_pour.config.franka.agents.rsl_rl_ppo_cfg import ( + FrankaPourPPORunnerCfg, + FrankaPourResetDatasetPPORunnerCfg, + FrankaPourResetMixturePPORunnerCfg, +) +from isaaclab_tasks.contrib.franka_pour.cube_bowl_spawner_cfg import CubeBowlSpawnerCfg +from isaaclab_tasks.contrib.franka_pour.cup_media import cup_cavity_lattice +from isaaclab_tasks.contrib.franka_pour.pour_env_cfg import ( + FrankaPourEnvCfg, + FrankaPourEnvCfg_PLAY, + FrankaPourEnvCfg_RESET_DATASET, + FrankaPourEnvCfg_RESET_DATASET_EVAL, + FrankaPourEnvCfg_RESET_DATASET_PLAY, + FrankaPourEnvCfg_RESET_MIXTURE, + FrankaPourEnvCfg_RESET_MIXTURE_EVAL, + FrankaPourEnvCfg_RESET_MIXTURE_PLAY, + FrankaPourEnvCfg_TELEOP, + _resolve_mpm_cell_cap, +) +from isaaclab_tasks.contrib.franka_pour.reset_utils import ( + balanced_cyclic_permutations, + boolean_selection_mask, + randomization_extent_index_pools, + sample_index_pools, + target_xy_behind_source, +) + + +def test_boolean_selection_mask_preserves_device_shape_and_dtype(): + selected = torch.tensor([[1, 3], [3, 4]], dtype=torch.long) + + mask = boolean_selection_mask(6, selected) + + assert mask.device == selected.device + assert mask.dtype == torch.bool + assert mask.shape == (6,) + assert mask.tolist() == [False, True, False, True, True, False] + + +def test_franka_pour_config_import_does_not_preload_usd_before_kit(): + result = subprocess.run( + [ + sys.executable, + "-c", + ( + "import sys; import isaaclab_tasks.contrib.franka_pour.pour_env_cfg; " + "assert 'newton' not in sys.modules; assert 'pxr' not in sys.modules" + ), + ], + check=False, + capture_output=True, + text=True, + ) + + assert result.returncode == 0, result.stderr + + +def test_balanced_cyclic_permutations_preserve_marginals_and_balance_pairings(): + values = torch.tensor([0.0, 0.5, 1.0, -0.5, -1.0]) + + permutations = balanced_cyclic_permutations(values, group_count=49) + + assert permutations.shape == (49, 5) + expected = torch.sort(values).values.expand(49, -1) + torch.testing.assert_close(torch.sort(permutations, dim=-1).values, expected) + for column in range(permutations.shape[1]): + _, counts = torch.unique(permutations[:, column], return_counts=True) + assert int(counts.amax() - counts.amin()) <= 1 + assert torch.unique(permutations, dim=0).shape[0] == values.numel() + + +def test_randomization_extent_pools_combine_all_axes_and_are_nested(): + source_positions = torch.tensor([[0.50, 0.00], [0.525, 0.00], [0.55, 0.00], [0.575, 0.00], [0.60, 0.00]]) + target_positions = torch.tensor([[0.50, -0.20], [0.51, -0.20], [0.54, -0.20], [0.52, -0.20], [0.50, -0.15]]) + source_yaws = torch.tensor([0.00, 0.10, 0.05, 0.15, 0.02]) + tcp_jitter = torch.tensor( + [[0.00, 0.00, 0.00], [0.012, 0.00, 0.00], [0.006, 0.00, 0.00], [0.018, 0.00, 0.00], [0.004, 0.00, 0.00]] + ) + + pools = randomization_extent_index_pools( + source_positions, + source_yaws, + target_positions, + tcp_jitter, + source_center=(0.50, 0.00), + source_half_range=(0.10, 0.10), + source_yaw_half_range=0.20, + target_center=(0.50, -0.20), + target_half_range=(0.05, 0.05), + tcp_jitter_half_range=(0.02, 0.02, 0.02), + extent_levels=(0.5, 0.8, 1.0), + ) + + assert [pool.tolist() for pool in pools] == [[0], [0, 1, 2], [0, 1, 2, 3, 4]] + assert bool(torch.all(torch.isin(pools[0], pools[1]))) + assert bool(torch.all(torch.isin(pools[1], pools[2]))) + torch.testing.assert_close(pools[2], torch.arange(5)) + + +def test_randomization_extent_pools_handle_zero_range_axes_without_nan(): + source_positions = torch.tensor([[0.5, -0.1], [0.5, 0.0], [0.5, 0.1], [0.5, 0.0]]) + target_positions = torch.tensor([[0.4, -0.2], [0.4, -0.1], [0.4, 0.0], [0.4, -0.1]]) + source_yaws = torch.zeros(4) + tcp_jitter = torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.01, 0.0]]) + + pools = randomization_extent_index_pools( + source_positions, + source_yaws, + target_positions, + tcp_jitter, + source_center=(0.5, 0.0), + source_half_range=(0.0, 0.1), + source_yaw_half_range=0.0, + target_center=(0.4, -0.1), + target_half_range=(0.0, 0.1), + tcp_jitter_half_range=(0.0, 0.0, 0.0), + extent_levels=(0.5, 1.0), + ) + + assert [pool.tolist() for pool in pools] == [[1], [0, 1, 2]] + + +def test_randomization_extent_pools_include_source_yaw(): + pools = randomization_extent_index_pools( + torch.zeros((3, 2)), + torch.tensor([0.0, 0.05, -0.10]), + torch.zeros((3, 2)), + torch.zeros((3, 3)), + source_center=(0.0, 0.0), + source_half_range=(0.0, 0.0), + source_yaw_half_range=0.10, + target_center=(0.0, 0.0), + target_half_range=(0.0, 0.0), + tcp_jitter_half_range=(0.0, 0.0, 0.0), + extent_levels=(0.5, 1.0), + ) + + assert [pool.tolist() for pool in pools] == [[0, 1], [0, 1, 2]] + + +def test_randomization_extent_pools_require_aligned_bank_rows(): + with pytest.raises(ValueError, match="same row count"): + randomization_extent_index_pools( + torch.zeros((2, 3)), + torch.zeros(2), + torch.zeros((1, 3)), + torch.zeros((2, 3)), + source_center=(0.0, 0.0), + source_half_range=(0.1, 0.1), + source_yaw_half_range=0.1, + target_center=(0.0, 0.0), + target_half_range=(0.1, 0.1), + tcp_jitter_half_range=(0.1, 0.1, 0.1), + extent_levels=(1.0,), + ) + + +def test_index_pool_sampling_maps_local_slots_back_to_global_bank_rows(monkeypatch): + pools = (torch.tensor([4, 8]), torch.tensor([1, 5, 9]), torch.tensor([0, 3, 6, 9])) + pool_ids = torch.tensor([0, 2, 1, 0, 2]) + + def sample_last_slot(high, size, *, device): + return torch.full(size, high - 1, device=device, dtype=torch.long) + + monkeypatch.setattr(torch, "randint", sample_last_slot) + sampled = sample_index_pools(pools, pool_ids) + + assert sampled.tolist() == [8, 9, 9, 8, 9] + + +def test_target_randomization_is_bounded_and_separated_behind_source(): + azimuth = math.radians(20.0) + source_xy = torch.tensor( + [ + [0.50 * math.cos(azimuth), -0.50 * math.sin(azimuth)], + [0.68 * math.cos(azimuth), 0.68 * math.sin(azimuth)], + ] + ) + unit_samples = torch.tensor([[0.0, 1.0], [1.0, 1.0]]) + + minimum_y_separation = torch.tensor([0.129, 0.150]) + target_xy = target_xy_behind_source( + source_xy, + target_center=(0.50, -0.18), + target_half_range=(0.08, 0.30), + minimum_y_separation=minimum_y_separation, + unit_samples=unit_samples, + ) + + assert bool(torch.all(target_xy[:, 0] >= 0.42)) + assert bool(torch.all(target_xy[:, 0] <= 0.58)) + assert bool(torch.all(target_xy[:, 1] >= -0.48)) + assert bool(torch.all(target_xy[:, 1] <= 0.12)) + assert bool(torch.all(source_xy[:, 1] - target_xy[:, 1] >= minimum_y_separation - 1.0e-6)) + + +def test_target_randomization_rejects_a_source_position_without_feasible_separation(): + source_xy = torch.tensor([[0.50, -0.40]]) + + with pytest.raises(ValueError, match="No target y-position"): + target_xy_behind_source( + source_xy, + target_center=(0.50, -0.18), + target_half_range=(0.08, 0.30), + minimum_y_separation=0.129, + unit_samples=torch.tensor([[0.5, 0.5]]), + ) + + +def test_finalize_builds_scene_assets_without_mutating_the_caller(): + original = FrankaPourEnvCfg() + original.scene.num_envs = 8 + + resolved = original.finalize() + + assert original.scene.source_cup is None + assert original.scene.target_cup is None + assert original.scene.media is None + assert isinstance(resolved.scene.source_cup, RigidObjectCfg) + assert isinstance(resolved.scene.target_cup, RigidObjectCfg) + assert isinstance(resolved.scene.media, MPMObjectCfg) + assert resolved is not original + assert resolved.scene is not original.scene + assert resolved.sim.physics.solver_cfg.scene_cfg is resolved.scene + assert _media_capacity(resolved) == 8 * 512 + + +def test_finalize_preserves_environment_spacing_for_large_sparse_batches(): + cfg = FrankaPourEnvCfg() + original_spacing = cfg.scene.env_spacing + cfg.scene.num_envs = 3068 + + resolved = cfg.finalize() + + assert cfg.scene.env_spacing == pytest.approx(original_spacing) + assert resolved.scene.env_spacing == pytest.approx(original_spacing) + assert FrankaPourEnvCfg_PLAY().finalize().scene.env_spacing == pytest.approx(original_spacing) + + +def test_finalize_propagates_final_cup_overrides_to_fresh_assets(): + original = FrankaPourEnvCfg() + original.source_cup_inner_width = 0.041 + original.target_cup_inner_depth = 0.153 + original.cup_grasp_box_half = (0.0275, 0.028, 0.0595) + original.gripper_preload_pos = 0.017 + original.cup_grasp_box_friction = 2.4 + original.target_cup_friction = 1.3 + original.cup_mass = 0.071 + original.cup_reset_pos = (0.43, 0.02, 0.01) + original.target_cup_reset_pos = (0.47, -0.21, 0.02) + original.arm_home = (0.1, 0.2, 0.3, -2.4, 0.5, 3.2, 0.7) + original.success_dwell_time_s = 0.10 + # This test intentionally moves the nominal cup independently of the task's base-centred + # polar workspace. Use the supported rectangular fallback so an unrelated asset override + # does not need to redefine the polar sector as well. + original.curriculum_randomized_source_radius_range = None + original.curriculum_randomized_source_position_range = (0.0, 0.0) + original.curriculum_randomized_carry_position_range = (0.0, 0.0) + + first = original.finalize() + first_source = first.scene.source_cup + first_target = first.scene.target_cup + assert isinstance(first_source.spawn, CubeBowlSpawnerCfg) + assert isinstance(first_target.spawn, CubeBowlSpawnerCfg) + assert first_source.spawn.inner_width == pytest.approx(0.041) + assert first_target.spawn.inner_depth == pytest.approx(0.153) + assert first_source.spawn.grasp_proxy_half_extents == pytest.approx((0.0275, 0.028, 0.0595)) + assert first_source.spawn.physics_material.static_friction == pytest.approx(2.4) + assert first_source.spawn.physics_material.dynamic_friction == pytest.approx(2.4) + assert first_target.spawn.physics_material.static_friction == pytest.approx(1.3) + assert first_target.spawn.physics_material.dynamic_friction == pytest.approx(1.3) + assert first_source.spawn.mass_props.mass == pytest.approx(0.071) + assert first_source.init_state.pos == pytest.approx((0.43, 0.02, 0.01)) + assert first_source.init_state.rot == (0.0, 0.0, 0.0, 1.0) + assert first_target.init_state.pos == pytest.approx((0.47, -0.21, 0.02)) + assert first_target.init_state.rot == (0.0, 0.0, 0.0, 1.0) + assert [first.scene.robot.init_state.joint_pos[f"panda_joint{i}"] for i in range(1, 8)] == pytest.approx( + original.arm_home + ) + assert first.terminations.success.params["dwell_time_s"] == pytest.approx(0.10) + assert first.actions.gripper_action.close_position == pytest.approx(0.014) + assert first.rewards.task_progress.params["grasp_preload_position"] == pytest.approx(0.017) + expected_contact_command = 0.028 - first.actions.gripper_action.contact_min_deflection + assert first.rewards.task_progress.params["max_gripper_command"] == pytest.approx(expected_contact_command) + assert first.rewards.delivered.params["max_gripper_command"] == pytest.approx(expected_contact_command) + assert first.terminations.success.params["max_gripper_command"] == pytest.approx(expected_contact_command) + + first.source_cup_inner_width = 0.049 + first.cup_grasp_box_half = (0.0315, 0.028, 0.0595) + second = first.finalize() + + assert second.scene.source_cup is not first.scene.source_cup + assert second.scene.target_cup is not first.scene.target_cup + assert second.scene.media is not first.scene.media + assert first.scene.source_cup.spawn.inner_width == pytest.approx(0.041) + assert second.scene.source_cup.spawn.inner_width == pytest.approx(0.049) + + +def test_resolved_cups_have_backend_neutral_rigid_properties(): + resolved = FrankaPourEnvCfg().finalize() + source = resolved.scene.source_cup + target = resolved.scene.target_cup + + assert source.prim_path == "{ENV_REGEX_NS}/SourceCup" + assert isinstance(source.spawn.rigid_props, UsdPhysicsRigidBodyCfg) + assert source.spawn.rigid_props.rigid_body_enabled is True + assert source.spawn.rigid_props.kinematic_enabled is False + assert isinstance(source.spawn.collision_props, UsdPhysicsCollisionCfg) + assert source.spawn.collision_props.collision_enabled is True + assert isinstance(source.spawn.mass_props, MassCfg) + assert isinstance(source.spawn.physics_material, RigidBodyMaterialBaseCfg) + assert source.spawn.grasp_proxy_half_extents == resolved.cup_grasp_box_half + + assert target.prim_path == "{ENV_REGEX_NS}/TargetCup" + assert isinstance(target.spawn.rigid_props, UsdPhysicsRigidBodyCfg) + assert target.spawn.rigid_props.rigid_body_enabled is True + assert target.spawn.rigid_props.kinematic_enabled is True + assert target.spawn.grasp_proxy_half_extents is None + assert isinstance(target.spawn.physics_material, RigidBodyMaterialBaseCfg) + + +def test_robot_authors_mujoco_gravity_compensation_with_joint_fragment(): + fragments = FrankaPourEnvCfg().scene.robot.spawn.joint_drive_props + + assert isinstance(fragments, list) + assert any(isinstance(fragment, MujocoJointCfg) and fragment.actuatorgravcomp is True for fragment in fragments) + + +def test_robot_uses_authored_pour_asset_actuator_defaults_without_mutating_global_franka_preset(): + robot = FrankaPourEnvCfg().finalize().scene.robot + + assert robot.spawn.usd_path == pour_env_cfg.FRANKA_POUR_ROBOT_USD_PATH + assert robot.spawn.usd_path == ( + "omniverse://isaac-dev.ov.nvidia.com/Isaac/IsaacLab/Robots/FrankaEmika/franka_panda.usda" + ) + for actuator_cfg in robot.actuators.values(): + assert actuator_cfg.effort_limit_sim is None + assert actuator_cfg.velocity_limit_sim is None + assert actuator_cfg.stiffness is None + assert actuator_cfg.damping is None + assert actuator_cfg.armature is None + + assert pour_env_cfg.FRANKA_PANDA_CFG.spawn.usd_path.endswith("/panda_instanceable.usd") + assert pour_env_cfg.FRANKA_PANDA_CFG.actuators["panda_shoulder"].stiffness == pytest.approx(80.0) + + +def test_robot_enables_only_authored_arm_collision_proxies_and_self_collision(): + robot = FrankaPourEnvCfg().scene.robot + + assert robot.spawn.func is pour_env_cfg.spawn_franka_with_arm_collisions + assert { + "link0_c", + "link1_c", + "link2_c", + "link3_c", + "link4_c", + "link5_c0", + "link5_c1", + "link5_c2", + "link6_c", + "link7_c", + } == pour_env_cfg.FRANKA_POUR_ARM_COLLISION_PROXIES + assert robot.spawn.articulation_props.enabled_self_collisions is True + + +def test_public_fixed_size_tuple_annotations_are_specific(): + expected = { + "arm_home": "tuple[float, float, float, float, float, float, float]", + "curriculum_pour_arm_q": "tuple[float, float, float, float, float, float, float]", + "curriculum_carry_arm_q": "tuple[float, float, float, float, float, float, float]", + "tcp_offset_pos": "tuple[float, float, float]", + "tcp_offset_rot": "tuple[float, float, float, float]", + "cup_grasp_box_half": "tuple[float, float, float]", + "cup_grasp_tcp_quat_c": "tuple[float, float, float, float]", + "cup_reset_pos": "tuple[float, float, float]", + "target_cup_reset_pos": "tuple[float, float, float]", + "curriculum_randomized_source_radius_range": "tuple[float, float] | None", + "curriculum_randomized_reset_tcp_offset_lower": "tuple[float, float, float] | None", + "curriculum_randomized_reset_tcp_offset_upper": "tuple[float, float, float] | None", + "particle_workspace_lower_bound": "tuple[float, float, float]", + "particle_workspace_upper_bound": "tuple[float, float, float]", + } + + assert {name: FrankaPourEnvCfg.__annotations__[name] for name in expected} == expected + + +def test_cfg_routes_each_body_to_exactly_one_solver(): + cfg = FrankaPourEnvCfg().finalize() + solver = cfg.sim.physics.solver_cfg + entries = {entry.name: entry for entry in solver.entries} + assert set(entries) == {"arm", "media"} + assert isinstance(solver, CouplerProxyCfg) + + arm = entries["arm"] + media = entries["media"] + assert arm.solver_cfg.integrator == "implicitfast" + assert arm.solver_cfg.use_mujoco_contacts is False + assert solver.scene_cfg is cfg.scene + assert cfg.sim.physics.collision_cfg.soft_contact_max == 0 + assert arm.include_static_shapes is True + assert arm.bodies == [SceneEntityCfg("robot"), SceneEntityCfg("source_cup"), SceneEntityCfg("target_cup")] + assert cfg.sim.physics.num_substeps == cfg.physics_substeps == 1 + assert arm.substeps == cfg.rigid_entry_substeps == 4 + assert media.substeps == cfg.mpm_entry_substeps == 2 + assert cfg.sim.physics.num_substeps * media.substeps == 2 + assert media.all_particles is True + assert media.bodies == [r".*/SpillFloor$"] + assert media.include_static_shapes is False + assert media.in_place is True + assert media.solver_cfg.grid_type == "sparse" + assert media.solver_cfg.grid_padding == 0 + assert media.solver_cfg.max_active_cell_count == 2 * 512 + assert media.solver_cfg.separate_worlds is True + assert media.solver_cfg.solver == "jacobi" + assert media.solver_cfg.warmstart_mode == "none" + assert media.solver_cfg.max_iterations == 24 + assert cfg.sim.physics.use_cuda_graph is True + + proxies = solver.proxies + assert len(proxies) == 1 + assert proxies[0].source == "arm" and proxies[0].destination == "media" + assert proxies[0].bodies == [SceneEntityCfg("source_cup"), SceneEntityCfg("target_cup")] + assert proxies[0].collision_pipeline is not None + assert proxies[0].collision_pipeline(None) is None + assert proxies[0].mass_scale == pytest.approx(cfg.proxy_mass_scale) + assert cfg.proxy_mass_scale == pytest.approx(100.0) + assert solver.iterations == cfg.proxy_iterations + assert not hasattr(pour_env_cfg, "CUP_LABEL_PATTERN") + assert all( + "Cup$" not in selector for entry in solver.entries for selector in entry.bodies if isinstance(selector, str) + ) + + +def test_finalize_propagates_post_init_solver_overrides(): + cfg = FrankaPourEnvCfg() + cfg.mpm_iterations = 48 + cfg.voxel_size = 0.02 + cfg.physics_substeps = 3 + cfg.rigid_entry_substeps = 4 + cfg.mpm_entry_substeps = 2 + cfg.use_cuda_graph = False + cfg.proxy_iterations = 4 + cfg.proxy_mass_scale = 0.25 + + # Match Hydra's override timing: the nested tree still contains values authored during + # ``__post_init__`` until finalization resolves the public top-level controls. + assert _media_entry(cfg).solver_cfg.max_iterations == 24 + assert _media_entry(cfg).solver_cfg.voxel_size == pytest.approx(0.01) + + resolved = cfg.finalize() + coupled_cfg = resolved.sim.physics.solver_cfg + entries = {entry.name: entry for entry in coupled_cfg.entries} + proxy = coupled_cfg.proxies[0] + + assert entries["media"].solver_cfg.max_iterations == 48 + assert entries["media"].solver_cfg.voxel_size == pytest.approx(0.02) + assert entries["arm"].substeps == 4 + assert entries["media"].substeps == 2 + assert resolved.sim.physics.num_substeps == 3 + assert resolved.sim.physics.use_cuda_graph is False + assert isinstance(coupled_cfg, CouplerProxyCfg) + assert coupled_cfg.iterations == 4 + assert proxy.mass_scale == pytest.approx(0.25) + + +@pytest.mark.parametrize( + "field,value,error_type", + [ + ("physics_substeps", 0, ValueError), + ("rigid_entry_substeps", True, ValueError), + ("mpm_entry_substeps", -1, ValueError), + ("mpm_iterations", 0, ValueError), + ("proxy_iterations", 0, ValueError), + ("voxel_size", 0.0, ValueError), + ("proxy_mass_scale", float("inf"), ValueError), + ("use_cuda_graph", 1, TypeError), + ], +) +def test_solver_controls_reject_invalid_overrides(field, value, error_type): + cfg = FrankaPourEnvCfg() + setattr(cfg, field, value) + + with pytest.raises(error_type, match=field): + cfg.finalize() + + +def test_visible_source_geometry_matches_grasp_proxy_and_fits_gripper(): + cfg = FrankaPourEnvCfg() + outer_width = cfg.source_cup_inner_width + 2.0 * cfg.source_cup_wall_thickness + outer_depth = cfg.source_cup_inner_depth + 2.0 * cfg.source_cup_wall_thickness + outer_height = cfg.source_cup_cavity_depth + cfg.source_cup_bottom_thickness + assert (outer_width, outer_depth, outer_height) == pytest.approx((0.056, 0.056, 0.119)) + assert cfg.cup_grasp_box_half == pytest.approx((outer_width / 2.0, outer_depth / 2.0, outer_height / 2.0)) + assert outer_depth < 2.0 * cfg.gripper_open_pos + assert cfg.grasp_contact_ke >= 1.0e5 + assert cfg.grasp_contact_kd >= 5.0e2 + assert cfg.cup_grasp_box_friction >= 2.0 + # The horizontal fingertip TCP sits well below the rim, so the complete finger pads engage the + # side wall and the glass rotates around an upper-middle pivot rather than its top edge. + assert cfg.source_cup_bottom_thickness < cfg.cup_grasp_height < outer_height + assert cfg.cup_grasp_height == pytest.approx(0.083) + assert outer_height - cfg.cup_grasp_height == pytest.approx(0.036) + assert 0.65 * outer_height < cfg.cup_grasp_height < 0.75 * outer_height + assert cfg.media_fill_frac * cfg.source_cup_cavity_depth == pytest.approx(0.70 * 0.027) + assert cfg.cup_grasp_tcp_quat_c == pytest.approx((0.0, math.sqrt(0.5), 0.0, math.sqrt(0.5))) + qx, qy, qz, qw = cfg.cup_grasp_tcp_quat_c + tool_axis_c = ( + 2.0 * (qx * qz + qy * qw), + 2.0 * (qy * qz - qx * qw), + 1.0 - 2.0 * (qx * qx + qy * qy), + ) + assert tool_axis_c == pytest.approx((1.0, 0.0, 0.0)) + jaw_axis_z = 2.0 * (qy * qz + qx * qw) + assert jaw_axis_z == pytest.approx(0.0) + + +@pytest.mark.parametrize( + "field,value", + [ + ("source_cup_cavity_depth", 0.0), + ("source_cup_wall_thickness", float("inf")), + ("cup_grasp_box_half", (0.028, 0.028, 0.018)), + ("cup_grasp_height", 0.009), + ("cup_grasp_height", 0.119), + ("cup_grasp_tcp_quat_c", (0.0, 0.0, 0.0)), + ("cup_grasp_tcp_quat_c", (0.0, 0.0, 0.0, 1.1)), + ("media_fill_frac", 0.0), + ("media_fill_frac", 1.01), + ], +) +def test_source_cup_rejects_inconsistent_geometry(field, value): + cfg = FrankaPourEnvCfg() + setattr(cfg, field, value) + + with pytest.raises(ValueError, match=field): + cfg.finalize() + + +def test_scene_cups_use_narrow_solver_only_builder_hook(): + source_path = Path(franka_pour.__file__).with_name("pour_env.py") + tree = ast.parse(source_path.read_text(encoding="utf-8"), filename=str(source_path)) + function_names = {node.name for node in ast.walk(tree) if isinstance(node, ast.FunctionDef)} + assert {"_build_custom_proto", "_add_cup_body", "_add_target_cup_bodies"}.isdisjoint(function_names) + + hook = next( + node + for node in ast.walk(tree) + if isinstance(node, ast.FunctionDef) and node.name == "_add_pour_world_to_builder" + ) + hook_calls = { + node.func.attr for node in ast.walk(hook) if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) + } + assert {"_find_world_body", "_add_particle_collider", "_add_rigid_collider"} <= hook_calls + + target_bridge = next( + node + for node in ast.walk(tree) + if isinstance(node, ast.FunctionDef) and node.name == "_add_kinematic_rigid_object_articulation" + ) + bridge_calls = { + node.func.attr + for node in ast.walk(target_bridge) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) + } + assert {"add_joint_free", "add_articulation"} <= bridge_calls + + +def test_randomized_ik_branch_ranking_prefers_open_joint6_complete_paths_and_falls_back(): + """Open reset branches must be preferred before ranking complete receiver paths.""" + source_path = Path(franka_pour.__file__).with_name("pour_env.py") + tree = ast.parse(source_path.read_text(encoding="utf-8"), filename=str(source_path)) + build_bank = next( + node + for node in ast.walk(tree) + if isinstance(node, ast.FunctionDef) and node.name == "_build_randomized_reset_bank" + ) + row_has_open_branch = next( + node + for node in ast.walk(build_bank) + if isinstance(node, ast.Assign) + and any(isinstance(target, ast.Name) and target.id == "row_has_open_reset_branch" for target in node.targets) + ) + open_path_filter = next( + node + for node in ast.walk(build_bank) + if isinstance(node, ast.AugAssign) + and isinstance(node.target, ast.Name) + and node.target.id == "complete_path_valid" + and node.lineno > row_has_open_branch.lineno + ) + collision_path_indices = next( + node + for node in ast.walk(build_bank) + if isinstance(node, ast.Assign) + and any(isinstance(target, ast.Name) and target.id == "collision_path_indices" for target in node.targets) + ) + global_path_indices = next( + node + for node in ast.walk(build_bank) + if isinstance(node, ast.Assign) + and any(isinstance(target, ast.Name) and target.id == "global_path_indices" for target in node.targets) + ) + + assert ast.unparse(row_has_open_branch.value) == ( + "(complete_path_valid & open_reset_branch[:, :, None, None]).any(dim=(1, 2, 3), keepdim=True)" + ) + assert ast.unparse(open_path_filter.value) == ("~row_has_open_reset_branch | open_reset_branch[:, :, None, None]") + assert isinstance(global_path_indices.value, ast.Attribute) + assert ast.unparse(global_path_indices.value.value.args[0]) == "flat_complete_path_score" + assert ast.unparse(collision_path_indices.value) == ( + "torch.cat((deterministic_source_path_indices, global_path_indices), dim=-1)" + ) + assert row_has_open_branch.lineno < open_path_filter.lineno < collision_path_indices.lineno + + complete_path_valid = torch.tensor( + [ + [[[True, False]], [[True, True]], [[False, False]]], + [[[True, True]], [[False, False]], [[True, False]]], + ] + ) + open_reset_branch = torch.tensor([[False, True, True], [False, True, False]]) + row_has_open_reset_branch = (complete_path_valid & open_reset_branch[:, :, None, None]).any( + dim=(1, 2, 3), keepdim=True + ) + ranking_mask = ~row_has_open_reset_branch | open_reset_branch[:, :, None, None] + + # Row 0 ranks only complete paths rooted at its available open reset branch. Row 1 has no + # complete open branch, so both valid folded reset branches remain available as fallbacks. + assert (complete_path_valid & ranking_mask).flatten(start_dim=1).sum(dim=-1).tolist() == [2, 3] + assert (complete_path_valid & ranking_mask)[0, 0].sum() == 0 + assert (complete_path_valid & ranking_mask)[0, 1].sum() == 2 + + +def test_invalid_initial_ik_rows_are_sanitized_then_filtered_from_reset_pools(): + """An infeasible reset row must not poison batched continuation IK or become sampleable.""" + source_path = Path(franka_pour.__file__).with_name("pour_env.py") + tree = ast.parse(source_path.read_text(encoding="utf-8"), filename=str(source_path)) + build_bank = next( + node + for node in ast.walk(tree) + if isinstance(node, ast.FunctionDef) and node.name == "_build_randomized_reset_bank" + ) + + def first_assignment(name: str) -> ast.Assign: + return min( + ( + node + for node in ast.walk(build_bank) + if isinstance(node, ast.Assign) + and any(isinstance(target, ast.Name) and target.id == name for target in node.targets) + ), + key=lambda node: node.lineno, + ) + + reset_has_candidate = first_assignment("reset_has_candidate") + fallback_indices = first_assignment("fallback_indices") + trajectory_q = first_assignment("trajectory_q") + randomized_collision_free = first_assignment("randomized_collision_free") + assert ast.unparse(reset_has_candidate.value) == "candidate_valid.any(dim=-1)" + assert ast.unparse(fallback_indices.value) == ( + "torch.where(reset_has_candidate, valid_fallback_indices, margin_fallback_indices)" + ) + assert ast.unparse(trajectory_q.value) == ( + "torch.where(trajectory_valid.unsqueeze(-1), expanded_q, fallback_q).clone()" + ) + assert "row_has_collision_free_path[randomized_rows]" in ast.unparse(randomized_collision_free.value) + + candidate_valid = torch.tensor([[False, True, False], [False, False, False]]) + expanded_margin = torch.tensor([[0.1, 0.2, 0.3], [float("nan"), -0.4, 0.5]]) + expanded_q = torch.arange(6, dtype=torch.float32).reshape(2, 3, 1) + reset_has_candidate_t = candidate_valid.any(dim=-1) + valid_fallback_indices = torch.argmax(candidate_valid.to(dtype=torch.int32), dim=-1) + margin_fallback_indices = torch.argmax(torch.nan_to_num(expanded_margin, nan=-torch.inf), dim=-1) + fallback_indices_t = torch.where( + reset_has_candidate_t, + valid_fallback_indices, + margin_fallback_indices, + ) + fallback_q = expanded_q[torch.arange(2), fallback_indices_t].unsqueeze(1) + sanitized = torch.where(candidate_valid.unsqueeze(-1), expanded_q, fallback_q) + + assert fallback_indices_t.tolist() == [1, 2] + torch.testing.assert_close(sanitized[:, :, 0], torch.tensor([[1.0, 1.0, 1.0], [5.0, 5.0, 5.0]])) + # Sanitization supplies finite continuation seeds without changing validity; the second row + # therefore remains ineligible for the eventual reset-index pool. + assert candidate_valid.any(dim=-1).tolist() == [True, False] + + +def test_randomized_collision_screen_covers_full_path_with_authored_finger_states(): + """Collision validation must include pour/tilt with the fingers preloaded around the glass.""" + source_path = Path(franka_pour.__file__).with_name("pour_env.py") + tree = ast.parse(source_path.read_text(encoding="utf-8"), filename=str(source_path)) + build_bank = next( + node + for node in ast.walk(tree) + if isinstance(node, ast.FunctionDef) and node.name == "_build_randomized_reset_bank" + ) + collision_call = next( + node + for node in ast.walk(build_bank) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "_collision_free_ik_candidates" + ) + source_candidates_assignment = next( + node + for node in ast.walk(build_bank) + if isinstance(node, ast.Assign) + and any(isinstance(target, ast.Name) and target.id == "source_collision_candidates" for target in node.targets) + ) + collision_waypoints = collision_call.args[1] + assert isinstance(collision_waypoints, ast.Tuple) + assert isinstance(collision_waypoints.elts[0], ast.Starred) + assert ast.unparse(collision_waypoints.elts[0].value) == "source_collision_candidates" + assert isinstance(source_candidates_assignment.value, ast.Tuple) + source_waypoints = source_candidates_assignment.value.elts + assert all( + isinstance(call, ast.Call) and isinstance(call.func, ast.Name) and call.func.id == "gather_source_waypoint" + for call in source_waypoints + ) + assert [ast.unparse(call.args[0]) for call in source_waypoints] == [ + "expanded_q", + "pregrasp_candidates", + "midgrasp_candidates", + "grasp_candidates", + "carry_candidates", + ] + assert [ast.unparse(call.args[1]) for call in source_waypoints] == [ + "float(self.cfg.gripper_open_pos)", + "float(self.cfg.gripper_open_pos)", + "float(self.cfg.gripper_open_pos)", + "float(self.cfg.gripper_open_pos)", + "float(self.cfg.gripper_preload_pos)", + ] + assert [ast.unparse(waypoint) for waypoint in collision_waypoints.elts[1:]] == [ + "pour_collision_candidates", + "tilt_collision_candidates", + ] + collision_candidate_source = ast.unparse(build_bank) + assert ( + "pour_collision_candidates[:, :, finger_coordinate_ids] = float(self.cfg.gripper_preload_pos)" + in collision_candidate_source + ) + assert ( + "tilt_collision_candidates[:, :, finger_coordinate_ids] = float(self.cfg.gripper_preload_pos)" + in collision_candidate_source + ) + + +def test_randomized_reset_bank_builds_ik_targets_in_environment_local_frame(): + """The reset bank must not inherit the clone layout's batch-dependent world offset.""" + source_path = Path(franka_pour.__file__).with_name("pour_env.py") + tree = ast.parse(source_path.read_text(encoding="utf-8"), filename=str(source_path)) + build_bank = next( + node + for node in ast.walk(tree) + if isinstance(node, ast.FunctionDef) and node.name == "_build_randomized_reset_bank" + ) + + assignments = { + target.id: node.value + for node in ast.walk(build_bank) + if isinstance(node, ast.Assign) + for target in node.targets + if isinstance(target, ast.Name) + } + assert ast.unparse(assignments["prototype_origin"]) == "-self.env_origins[0]" + assert ast.unparse(assignments["target_positions_w"]) == "tcp_positions" + assert ast.unparse(assignments["pregrasp_target_positions_w"]) == "pregrasp_tcp_positions" + assert ast.unparse(assignments["grasp_target_positions_w"]) == "grasp_tcp_positions" + + # One fresh centered builder feeds collision validation and another feeds the local IK model. + local_builder_calls = [ + node + for node in ast.walk(build_bank) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "local_prototype_builder" + ] + assert len(local_builder_calls) == 2 + env_origin_references = [ + node + for node in ast.walk(build_bank) + if isinstance(node, ast.Subscript) and ast.unparse(node) == "self.env_origins[0]" + ] + assert len(env_origin_references) == 1 + + collision_validation_call = next( + node + for node in ast.walk(build_bank) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "_collision_free_ik_candidates" + ) + assert ast.unparse(collision_validation_call.args[0]) == "prototype_builder" + + +def test_polar_source_workspace_defaults_cover_a_broad_base_centered_sector(): + """The preset's polar domain must be broad and fit its conservative Cartesian bounds.""" + cfg = FrankaPourEnvCfg() + + assert cfg.curriculum_randomized_source_radius_range == pytest.approx((0.40, 0.78)) + assert cfg.curriculum_randomized_source_azimuth_range == pytest.approx(math.radians(35.0)) + assert cfg.curriculum_randomized_source_xy_correlation == pytest.approx(0.0) + + minimum_radius, maximum_radius = cfg.curriculum_randomized_source_radius_range + azimuth = cfg.curriculum_randomized_source_azimuth_range + required_half_range = ( + max( + abs(minimum_radius * math.cos(azimuth) - cfg.cup_reset_pos[0]), + abs(maximum_radius - cfg.cup_reset_pos[0]), + ), + maximum_radius * math.sin(azimuth), + ) + assert all( + configured >= required + for configured, required in zip( + cfg.curriculum_randomized_source_position_range, + required_half_range, + strict=True, + ) + ) + + +def test_builder_hook_limits_lookups_to_current_world_tail(): + from isaaclab_tasks.contrib.franka_pour.pour_env import FrankaPourEnv + + builder = SimpleNamespace( + body_world=[None, None, 0, 0, 1, 1, 1], + shape_world=[None, 0, 0, 1, 1], + ) + + assert FrankaPourEnv._current_world_range(builder, "body", 1) == range(4, 7) + assert FrankaPourEnv._current_world_range(builder, "shape", 1) == range(3, 5) + with pytest.raises(RuntimeError, match="no body entries for open world 2"): + FrankaPourEnv._current_world_range(builder, "body", 2) + + +def test_task_has_receiving_cup_safe_actions_and_success_threshold(): + cfg = FrankaPourEnvCfg() + assert cfg.target_cup_reset_pos != cfg.cup_reset_pos + assert isinstance(cfg.actions.arm_action, mdp.RelativeJointPositionActionCfg) + assert cfg.actions.arm_action.joint_names == [f"panda_joint{i}" for i in range(1, 8)] + assert cfg.actions.arm_action.preserve_order is True + assert cfg.actions.arm_action.scale == pytest.approx(0.08) + assert cfg.actions.arm_action.use_zero_offset is True + assert not hasattr(cfg.actions.arm_action, "waypoint_phases") + # Keep a meaningful particle margin below the demonstrated trajectory's lower tail while the + # one-time delivery reward still drives transfer beyond the success threshold. + assert cfg.curriculum_target_frac[-1] == pytest.approx(0.30) + assert cfg.episode_length_s == pytest.approx(5.0) + + +def test_training_curriculum_closes_by_default_and_penalizes_irrecoverable_spill(): + cfg = FrankaPourEnvCfg() + assert type(cfg.actions.gripper_action).__name__ == "CurriculumGripperPositionActionCfg" + assert cfg.actions.gripper_action.scale == pytest.approx(0.016) + assert cfg.actions.gripper_action.alpha == pytest.approx(0.2) + assert cfg.actions.gripper_action.use_incremental_target is False + assert cfg.actions.gripper_action.default_position == pytest.approx(cfg.gripper_preload_pos) + assert cfg.actions.gripper_action.neutral_position == pytest.approx(cfg.gripper_open_pos) + assert cfg.actions.gripper_action.limit_to_preload is False + assert cfg.actions.gripper_action.force_open_before_phase_stage == -1 + assert not hasattr(cfg.actions.gripper_action, "hold_offset_through_stage") + assert cfg.actions.gripper_action.close_position == pytest.approx(0.021) + assert cfg.actions.gripper_action.open_position == pytest.approx(cfg.gripper_open_pos) + reward_names = {name for name, term_cfg in vars(cfg.rewards).items() if isinstance(term_cfg, RewardTermCfg)} + assert reward_names == { + "task_progress", + "approach_progress", + "grasp_lift_progress", + "delivered", + "success", + "spill", + "failure", + "action_rate", + "action_magnitude", + } + assert cfg.rewards.task_progress.func is mdp.PourTaskProgress + assert cfg.rewards.task_progress.weight == pytest.approx(5.0) + assert cfg.rewards.task_progress.params["grasp_reach_std"] == pytest.approx(0.015) + assert cfg.rewards.task_progress.params["grasp_preload_position"] == pytest.approx(0.024) + assert cfg.rewards.task_progress.params["source_offset_xy"] == pytest.approx(cfg.pour_source_offset_xy) + assert cfg.rewards.task_progress.params["target_tilt"] == pytest.approx(math.radians(140.0)) + minimum_drain_tilt = math.atan2(cfg.source_cup_cavity_depth, 0.5 * cfg.source_cup_inner_depth) + assert cfg.rewards.task_progress.params["target_tilt"] > minimum_drain_tilt + assert cfg.rewards.task_progress.params["target_tilt"] < math.pi + assert cfg.rewards.task_progress.params["pour_direction_xy"] == pytest.approx((0.0, -1.0)) + assert cfg.rewards.task_progress.params["source_mouth_height"] == pytest.approx( + cfg.source_cup_bottom_thickness + cfg.source_cup_cavity_depth + ) + assert cfg.rewards.task_progress.params["alignment_radius"] == pytest.approx(0.15) + assert cfg.rewards.task_progress.params["active_through_stage"] == cfg.curriculum_stage_names.index("carry") + assert cfg.rewards.task_progress.params["min_lift_height"] == pytest.approx(0.05) + assert cfg.rewards.task_progress.params["max_tcp_distance"] == pytest.approx(0.018) + assert cfg.rewards.task_progress.params["max_gripper_width_error"] == pytest.approx(0.006) + assert cfg.rewards.task_progress.params["max_gripper_command"] == pytest.approx( + cfg._resolved_success_max_gripper_command() + ) + assert not hasattr(cfg.rewards, "reach") + assert not hasattr(cfg.rewards, "grasp") + assert not hasattr(cfg.rewards, "lift") + assert not hasattr(cfg.rewards, "align") + assert not hasattr(cfg.rewards, "tilt") + assert cfg.rewards.approach_progress.func is mdp.ApproachProgress + assert cfg.rewards.approach_progress.weight == pytest.approx(8.0) + assert cfg.rewards.approach_progress.params["position_std"] == pytest.approx(0.20) + assert cfg.rewards.approach_progress.params["orientation_std"] == pytest.approx(0.75) + assert cfg.rewards.approach_progress.params["open_hand_fraction"] == pytest.approx(0.35) + assert cfg.rewards.approach_progress.params["active_from_stage"] == cfg.curriculum_stage_names.index("approach_1") + assert cfg.rewards.grasp_lift_progress.func is mdp.GraspLiftProgress + assert cfg.rewards.grasp_lift_progress.weight == pytest.approx(10.0) + assert cfg.rewards.grasp_lift_progress.params["target_height"] == pytest.approx(0.10) + assert cfg.rewards.grasp_lift_progress.params["grasp_reach_std"] == pytest.approx(0.025) + assert cfg.rewards.grasp_lift_progress.params["grasp_fraction"] == pytest.approx(0.40) + assert cfg.rewards.grasp_lift_progress.params["active_from_stage"] == cfg.curriculum_stage_names.index("near_carry") + assert cfg.rewards.delivered.func is mdp.HeldDeliveryProgress + assert cfg.rewards.delivered.params["min_lift_height"] == pytest.approx(0.05) + assert cfg.rewards.delivered.params["max_tcp_distance"] == pytest.approx(0.018) + assert cfg.rewards.delivered.params["max_gripper_width_error"] == pytest.approx(0.006) + assert cfg.rewards.delivered.params["max_gripper_command"] == pytest.approx( + cfg._resolved_success_max_gripper_command() + ) + assert cfg.rewards.spill.func is mdp.NewlySpilledParticles + assert cfg.rewards.task_progress.weight < cfg.rewards.success.weight + assert cfg.rewards.spill.weight == pytest.approx(-30.0) + assert cfg.rewards.failure.func is mdp.terminal_failure + assert cfg.rewards.failure.weight == pytest.approx(-35.0) + assert cfg.terminations.spill.func is mdp.excessive_spill + assert cfg.terminations.extreme_rigid_state.func is mdp.extreme_rigid_state + assert cfg.terminations.lost_grasp.func is mdp.lost_lifted_grasp + assert cfg.terminations.lost_grasp.params["dwell_time_s"] == pytest.approx(0.05) + assert cfg.terminations.success.func is mdp.stable_pour_success + assert cfg.terminations.success.params["dwell_time_s"] == pytest.approx(0.15) + assert cfg.terminations.success.params["min_lift_height"] == pytest.approx(0.05) + assert cfg.terminations.success.params["max_tcp_distance"] == pytest.approx(0.018) + assert cfg.terminations.success.params["max_gripper_width_error"] == pytest.approx(0.006) + assert cfg.terminations.success.params["max_gripper_command"] == pytest.approx( + cfg._resolved_success_max_gripper_command() + ) + termination_names = [ + name for name, term_cfg in vars(cfg.terminations).items() if isinstance(term_cfg, TerminationTermCfg) + ] + assert termination_names[-2:] == ["success", "time_out"] + assert cfg.terminations.time_out.func is mdp.unsuccessful_time_out + assert cfg.max_spill_fraction == pytest.approx(0.10) + assert cfg.spill_table_height == pytest.approx(0.0) + assert cfg.state_bound_joint_position_margin == pytest.approx(0.05) + assert cfg.state_bound_max_joint_velocity == pytest.approx(20.0) + assert cfg.state_bound_max_cup_linear_velocity == pytest.approx(10.0) + assert cfg.state_bound_max_cup_angular_velocity == pytest.approx(50.0) + assert abs(cfg.rewards.action_rate.weight) <= 0.002 + assert cfg.rewards.action_magnitude.func is mdp.action_l2 + assert cfg.rewards.action_magnitude.weight == pytest.approx(-0.05) + + agent = FrankaPourPPORunnerCfg() + assert agent.class_name == "OnPolicyRunner" + assert agent.save_interval == 50 + assert agent.actor.distribution_cfg.class_name == "GaussianDistribution" + assert agent.actor.distribution_cfg.init_std == pytest.approx(0.1) + assert agent.actor.distribution_cfg.std_type == "log" + assert agent.actor.obs_normalization is False + assert agent.critic.obs_normalization is False + assert agent.clip_actions == pytest.approx(1.0) + assert agent.algorithm.entropy_coef == pytest.approx(1.0e-3) + assert agent.algorithm.num_learning_epochs == 5 + assert agent.algorithm.clip_param == pytest.approx(0.2) + assert agent.algorithm.learning_rate == pytest.approx(1.0e-4) + assert agent.algorithm.max_grad_norm == pytest.approx(1.0) + assert agent.algorithm.schedule == "fixed" + assert cfg.rewards.task_progress.params["discount_factor"] == pytest.approx(agent.algorithm.gamma) + assert cfg.rewards.approach_progress.params["discount_factor"] == pytest.approx(agent.algorithm.gamma) + assert cfg.rewards.grasp_lift_progress.params["discount_factor"] == pytest.approx(agent.algorithm.gamma) + assert agent.obs_groups == {"actor": ["policy"], "critic": ["policy", "privileged"]} + assert agent.logger == "wandb" + assert agent.wandb_project == "franka-pour-mpm" + + +def test_backward_curriculum_config_is_complete_and_play_uses_randomized_task(): + cfg = FrankaPourEnvCfg() + stage_count = len(cfg.curriculum_stage_names) + + assert cfg.is_finite_horizon is True + assert isinstance(cfg.curriculum.stage, CurriculumTermCfg) + assert cfg.curriculum.stage.func is mdp.PourCurriculum + assert not hasattr(cfg, "reset_mixture_probabilities") + assert cfg.curriculum_stage_names == ( + "drain", + "deep_tilt", + "tilt", + "pour", + "near_carry", + "mid_carry", + "carry", + "grasp", + "approach_1", + "approach_2", + "approach_3", + "approach_4", + "approach_5", + "approach_6", + "full", + "randomized", + ) + assert stage_count == 16 + assert len(cfg.curriculum_target_frac) == stage_count + assert cfg.curriculum_target_frac == pytest.approx( + ( + 0.05, + 0.08, + 0.10, + 0.15, + 0.15, + 0.18, + 0.20, + 0.30, + 0.30, + 0.30, + 0.30, + 0.30, + 0.30, + 0.30, + 0.30, + 0.30, + ) + ) + assert list(cfg.curriculum_target_frac) == sorted(cfg.curriculum_target_frac) + assert cfg.curriculum_start_stage == 0 + assert cfg.curriculum_freeze is False + assert cfg.curriculum_min_resets_per_stage == 4096 + assert cfg.curriculum_min_reset_cohorts_per_stage == pytest.approx(8.0) + assert cfg.curriculum_previous_stage_replay_fraction == pytest.approx(0.1) + assert cfg.curriculum_frontier_entry_replay_fraction == pytest.approx(0.5) + assert cfg.curriculum_transport_reset_fractions == pytest.approx((1.0 / 3.0, 2.0 / 3.0)) + assert cfg.curriculum_grasp_approach_fractions == pytest.approx((0.75, 0.50, 0.375, 0.25, 0.125, 0.0)) + transport_arm_q = cfg._curriculum_transport_arm_configs() + assert len(transport_arm_q) == 2 + for fraction, arm_q in zip(cfg.curriculum_transport_reset_fractions, transport_arm_q, strict=True): + assert arm_q == pytest.approx( + tuple( + (1.0 - fraction) * pour_q + fraction * carry_q + for pour_q, carry_q in zip(cfg.curriculum_pour_arm_q, cfg.curriculum_carry_arm_q, strict=True) + ) + ) + assert cfg.curriculum_randomization_extent_levels == pytest.approx( + (0.0, 0.05, 0.10, 0.20, 0.35, 0.50, 0.70, 0.85, 1.0) + ) + assert cfg.curriculum_independent_arm_fraction_levels == pytest.approx( + (0.0, 0.0, 0.0, 0.0, 0.10, 0.25, 0.50, 0.75, 1.0) + ) + assert cfg.curriculum_independent_target_fraction_levels == pytest.approx( + (0.0, 0.0, 0.0, 0.10, 0.25, 0.50, 0.70, 0.85, 1.0) + ) + assert stage_count - 1 + len(cfg.curriculum_randomization_extent_levels) == 24 + assert cfg.curriculum_randomization_promotion_threshold == pytest.approx(0.65) + assert cfg.curriculum_randomization_start_level == 0 + # Do not expose a generic stage identifier. The actor does receive behaviorally relevant + # finite-state variables and the active delivery goal so the finite-horizon MDP is Markov. + assert not hasattr(cfg.observations.policy, "curriculum_context") + assert not hasattr(cfg.observations.policy, "arm_reference_phase") + assert not hasattr(cfg.observations.policy, "arm_reference_error") + assert not hasattr(cfg.observations.policy, "trajectory_status") + assert cfg.observations.policy.time_remaining.func is mdp.time_remaining_obs + assert cfg.observations.policy.pour_target_fraction.func is mdp.pour_target_fraction_obs + assert not hasattr(cfg.observations.policy, "success_dwell") + assert not hasattr(cfg.observations.policy, "lost_grasp_dwell") + assert cfg.observations.policy.target_pose.func is mdp.target_pose_obs + assert not hasattr(cfg.observations.policy, "particle_fractions") + assert not hasattr(cfg.observations.policy, "particle_transfer") + assert cfg.observations.policy.arm_q.scale == pytest.approx(0.3) + assert cfg.observations.policy.arm_qd.scale == pytest.approx(0.05) + assert cfg.observations.policy.tcp_to_grasp_position_c.func is mdp.tcp_to_grasp_position_c_obs + assert cfg.observations.policy.tcp_to_grasp_position_c.scale == pytest.approx(10.0) + assert cfg.observations.policy.grasp_to_tcp_quat.func is mdp.grasp_to_tcp_quat_obs + assert cfg.observations.policy.target_position_c.func is mdp.target_position_c_obs + assert cfg.observations.policy.target_position_c.scale == pytest.approx(5.0) + assert cfg.observations.policy.finger_position.func is mdp.finger_position_obs + assert cfg.observations.policy.finger_position.scale == pytest.approx(25.0) + assert cfg.observations.policy.finger_velocity.func is mdp.finger_velocity_obs + assert cfg.observations.policy.finger_velocity.scale == pytest.approx(5.0) + assert cfg.observations.policy.last_action.scale == pytest.approx(0.2) + assert cfg.observations.policy.gripper_contact.func is mdp.gripper_contact_obs + assert cfg.observations.policy.gripper_contact.scale == pytest.approx(250.0) + assert cfg.observations.privileged.success_dwell.func is mdp.success_dwell_obs + assert cfg.observations.privileged.lost_grasp_dwell.func is mdp.lost_grasp_dwell_obs + assert cfg.observations.privileged.cup_velocity.func is mdp.cup_velocity_obs + assert cfg.observations.privileged.particle_fractions.func is mdp.particle_fractions_obs + assert cfg.observations.privileged.particle_transfer.func is mdp.particle_transfer_obs + assert cfg.observations.privileged.held_delivery_history.func is mdp.held_delivery_history_obs + assert cfg.actions.gripper_action.contact_min_deflection == pytest.approx(0.001) + assert cfg.actions.gripper_action.contact_max_velocity == pytest.approx(0.05) + + agent = FrankaPourPPORunnerCfg() + episode_steps = round(cfg.episode_length_s / (cfg.sim.dt * cfg.decimation)) + documented_training_env_count = 512 + available_resets = documented_training_env_count * agent.max_iterations * agent.num_steps_per_env // episode_steps + resets_per_frontier = max( + cfg.curriculum_min_resets_per_stage, + math.ceil(cfg.curriculum_min_reset_cohorts_per_stage * documented_training_env_count), + ) + required_resets = resets_per_frontier * (stage_count - 1 + len(cfg.curriculum_randomization_extent_levels)) + assert required_resets <= available_resets + + assert cfg.curriculum_randomized_source_position_range == pytest.approx((0.30, 0.45)) + assert cfg.curriculum_randomized_source_radius_range == pytest.approx((0.40, 0.78)) + assert cfg.curriculum_randomized_source_azimuth_range == pytest.approx(math.radians(35.0)) + assert cfg.curriculum_randomized_source_xy_correlation == pytest.approx(0.0) + assert cfg.curriculum_randomized_carry_position_range == pytest.approx((0.03, 0.03)) + assert cfg.curriculum_randomized_source_yaw_range == pytest.approx(math.radians(30.0)) + assert cfg.curriculum_randomized_target_center_xy == pytest.approx((0.50, -0.18)) + assert cfg.curriculum_randomized_target_position_range == pytest.approx((0.15, 0.44)) + assert cfg.curriculum_randomized_cup_clearance == pytest.approx(0.04) + assert cfg.curriculum_grasp_descent_overshoot == pytest.approx(0.0) + assert cfg.curriculum_randomized_reset_tcp_standoff == pytest.approx((-0.12, 0.0, 0.0)) + assert cfg.curriculum_randomized_reset_tcp_jitter == pytest.approx((0.02, 0.03, 0.0)) + assert cfg.curriculum_randomized_reset_tcp_offset_lower == pytest.approx((-0.16, -0.16, 0.0)) + assert cfg.curriculum_randomized_reset_tcp_offset_upper == pytest.approx((0.02, 0.16, 0.25)) + assert cfg.curriculum_randomized_reset_tcp_rotation_angle_range == pytest.approx( + (math.radians(20.0), math.radians(60.0)) + ) + assert cfg.curriculum_randomized_reset_tcp_min_grasp_distance == pytest.approx(0.09) + assert cfg.curriculum_randomized_pour_clearance == pytest.approx(0.01) + assert cfg.curriculum_randomized_reset_ik_grid_size == 7 + assert cfg.curriculum_randomized_reset_ik_samples_per_source == 11 + assert cfg.curriculum_randomized_min_source_cell_fraction == pytest.approx(0.2) + assert cfg.curriculum_randomized_min_reset_variants_per_source == 2 + assert cfg.curriculum_randomized_reset_joint6_max == pytest.approx(3.75) + + arm_configs = ( + cfg.curriculum_drain_arm_q, + cfg.curriculum_deep_tilt_arm_q, + cfg.curriculum_tilt_arm_q, + cfg.curriculum_pour_arm_q, + cfg.curriculum_carry_arm_q, + cfg.arm_home, + cfg.arm_home, + cfg.arm_home, + ) + for arm_q in arm_configs: + assert len(arm_q) == 7 + for position, (lower, upper) in zip(arm_q, pour_env_cfg.PANDA_ARM_JOINT_LIMITS, strict=True): + assert lower <= position <= upper + + for preset in (FrankaPourEnvCfg_PLAY(), FrankaPourEnvCfg_TELEOP()): + assert preset.curriculum_start_stage == stage_count - 1 + assert preset.curriculum_randomization_start_level == len(preset.curriculum_randomization_extent_levels) - 1 + assert preset.curriculum_freeze is True + + +def test_reset_dataset_config_uses_adaptive_sampling_and_frozen_evaluation(): + cfg = FrankaPourEnvCfg_RESET_DATASET() + eval_cfg = FrankaPourEnvCfg_RESET_DATASET_EVAL() + agent = FrankaPourResetDatasetPPORunnerCfg() + + assert cfg.curriculum.reset_dataset.func is mdp.PourResetDatasetCurriculum + assert cfg.reset_dataset_path == "datasets/franka_pour/reset_dataset.pt" + assert cfg.reset_dataset_content_sha256 is None + assert cfg.reset_dataset_top_grasp_count is None + for legacy_field in ( + "reset_mixture_probabilities", + "reset_mixture_near_object_open_phase_probabilities", + "reset_mixture_near_object_preloaded_probability", + "reset_mixture_statistics_window_size", + "reset_mixture_validated_cache_path", + "reset_mixture_validation_steps", + ): + assert not hasattr(cfg, legacy_field) + assert cfg.reset_dataset_sampler.target_success_rate == pytest.approx(0.50) + assert cfg.reset_dataset_sampler.temperature == pytest.approx(0.10) + assert cfg.reset_dataset_sampler.history_capacity == 32 + assert cfg.reset_dataset_sampler.prior_strength == pytest.approx(4.0) + assert cfg.reset_dataset_sampler.initial_frontier_size == 128 + assert cfg.reset_dataset_sampler.probe_size == 256 + assert cfg.reset_dataset_sampler.probe_fraction == pytest.approx(0.10) + assert cfg.reset_dataset_sampler.replay_fraction == pytest.approx(0.10) + assert cfg.reset_dataset_sampler.frontier_evidence == pytest.approx(2.0) + assert cfg.pour_target_frac == pytest.approx(0.30) + assert isinstance(cfg.actions.arm_action, mdp.RelativeJointPositionActionCfg) + assert cfg.actions.arm_action.joint_names == [f"panda_joint{i}" for i in range(1, 8)] + assert cfg.actions.arm_action.preserve_order is True + assert cfg.actions.arm_action.scale == pytest.approx(0.015) + assert cfg.actions.arm_action.use_zero_offset is True + assert cfg.actions.gripper_action.use_incremental_target is False + assert cfg.actions.gripper_action.binary_threshold == pytest.approx(0.0) + assert cfg.actions.gripper_action.alpha == pytest.approx(1.0 - 0.8 ** (1.0 / 3.0)) + assert eval_cfg.curriculum_freeze is True + assert eval_cfg.reset_dataset_path == cfg.reset_dataset_path + assert eval_cfg.reset_dataset_top_grasp_count is None + reward_names = {name for name, term_cfg in vars(cfg.rewards).items() if isinstance(term_cfg, RewardTermCfg)} + assert reward_names == { + "reach", + "goal_distance", + "success", + "action_magnitude", + "action_rate", + "joint_velocity", + "failure", + } + assert cfg.rewards.reach.func is mdp.tcp_cup_distance_tanh + assert cfg.rewards.reach.weight == pytest.approx(0.1) + assert cfg.rewards.reach.params["std"] == pytest.approx(0.3) + assert cfg.rewards.goal_distance.func is mdp.media_target_distance_tanh + assert cfg.rewards.goal_distance.weight == pytest.approx(0.1) + assert cfg.rewards.goal_distance.params["std"] == pytest.approx(0.2) + assert cfg.rewards.success.func is mdp.pour_success_bonus + assert cfg.rewards.success.weight == pytest.approx(1.0) + assert cfg.rewards.success.params == {} + assert cfg.rewards.action_magnitude.func is mdp.action_l2 + assert cfg.rewards.action_magnitude.weight == pytest.approx(-1.0e-4) + assert cfg.rewards.action_rate.func is mdp.action_rate_l2 + assert cfg.rewards.action_rate.weight == pytest.approx(-1.0e-3) + assert cfg.rewards.joint_velocity.func is mdp.finite_joint_velocity_l2 + assert cfg.rewards.joint_velocity.weight == pytest.approx(-1.0e-2) + assert cfg.rewards.failure.func is mdp.terminal_failure + assert cfg.rewards.failure.weight == pytest.approx(-1.0) + assert cfg.rewards.failure.params == {"include_time_out": False} + assert cfg.terminations.success.func is mdp.immediate_pour_success + assert cfg.terminations.success.params == {} + assert cfg.terminations.lost_grasp.params["terminate"] is False + assert cfg.terminations.spill.params["terminate"] is True + assert cfg.max_spill_fraction == pytest.approx(0.30) + assert cfg.terminations.time_out.func is mdp.unsuccessful_time_out + assert cfg.terminations.time_out.time_out is True + assert cfg.episode_length_s == pytest.approx(7.0) + assert cfg.is_finite_horizon is False + assert cfg.decimation == 4 + assert cfg.sim.render_interval == cfg.decimation + assert cfg.observations.policy.history_length == 13 + policy_step_s = cfg.sim.dt * cfg.decimation + assert 1.0 / policy_step_s == pytest.approx(30.0) + assert agent.num_steps_per_env * policy_step_s == pytest.approx(32.0 / 30.0) + assert round(cfg.episode_length_s / policy_step_s) == 210 + assert not any( + "stage" in parameter + for term_cfg in vars(cfg.rewards).values() + if isinstance(term_cfg, RewardTermCfg) + for parameter in (term_cfg.params or {}) + ) + assert vars(eval_cfg.rewards) == vars(cfg.rewards) + finalized = cfg.finalize() + assert isinstance(finalized.rewards, pour_env_cfg.ResetDatasetRewardsCfg) + assert isinstance(finalized.actions.arm_action, mdp.RelativeJointPositionActionCfg) + assert finalized.actions.arm_action.scale == pytest.approx(0.015) + assert finalized.actions.gripper_action.binary_threshold == pytest.approx(0.0) + + overridden = FrankaPourEnvCfg_RESET_DATASET() + overridden.actions.arm_action.scale = 0.015 + overridden.decimation = 6 + overridden = overridden.finalize() + assert overridden.actions.arm_action.scale == pytest.approx(0.015) + assert overridden.sim.render_interval == 6 + assert agent.experiment_name == "franka_pour_reset_dataset_joint_rel" + assert agent.run_name == "reset_dataset_joint_rel" + + +def test_reset_dataset_play_preserves_policy_abi_with_captured_sparse_grid(): + eval_cfg = FrankaPourEnvCfg_RESET_DATASET_EVAL().finalize() + play_cfg = FrankaPourEnvCfg_RESET_DATASET_PLAY().finalize() + + assert play_cfg.scene.num_envs == 1 + assert play_cfg.curriculum_freeze is True + assert play_cfg.reset_dataset_path == eval_cfg.reset_dataset_path + assert play_cfg.reset_dataset_content_sha256 == eval_cfg.reset_dataset_content_sha256 + assert play_cfg.reset_dataset_top_grasp_count == eval_cfg.reset_dataset_top_grasp_count + assert play_cfg.actions.to_dict() == eval_cfg.actions.to_dict() + assert play_cfg.observations.policy.to_dict() == eval_cfg.observations.policy.to_dict() + assert play_cfg.episode_length_s == pytest.approx(eval_cfg.episode_length_s) + assert _media_entry(eval_cfg).solver_cfg.grid_type == "sparse" + assert eval_cfg.sim.physics.use_cuda_graph is True + play_solver_cfg = _media_entry(play_cfg).solver_cfg + assert play_solver_cfg.grid_type == "sparse" + assert play_solver_cfg.grid_padding == 0 + assert play_solver_cfg.max_active_cell_count == 512 + assert play_solver_cfg.separate_worlds is True + assert play_cfg.sim.physics.use_cuda_graph is True + assert play_cfg.terminations.success.func is mdp.immediate_pour_success + assert play_cfg.terminations.success.params == {} + assert play_cfg.viewer.lookat == pytest.approx(eval_cfg.viewer.lookat) + assert play_cfg.viewer.eye == pytest.approx((0.9, 0.65, 0.5)) + eval_distance = math.dist(eval_cfg.viewer.eye, eval_cfg.viewer.lookat) + play_distance = math.dist(play_cfg.viewer.eye, play_cfg.viewer.lookat) + assert eval_distance - play_distance == pytest.approx(1.0, abs=0.02) + assert play_cfg.sim.default_visualizer_cfg.eye == pytest.approx(play_cfg.viewer.eye) + assert play_cfg.sim.default_visualizer_cfg.lookat == pytest.approx(play_cfg.viewer.lookat) + + +def test_reset_dataset_semantics_do_not_change_traditional_curriculum_defaults(): + reset_dataset = FrankaPourEnvCfg_RESET_DATASET() + reverse = FrankaPourEnvCfg() + + assert reset_dataset.rewards.success.func is mdp.pour_success_bonus + assert reset_dataset.rewards.failure.params == {"include_time_out": False} + assert reset_dataset.terminations.success.func is mdp.immediate_pour_success + assert reset_dataset.terminations.success.params == {} + assert reset_dataset.terminations.lost_grasp.params["terminate"] is False + assert reset_dataset.terminations.spill.params["terminate"] is True + assert reset_dataset.max_spill_fraction == pytest.approx(0.30) + assert reset_dataset.terminations.time_out.func is mdp.unsuccessful_time_out + assert isinstance(reset_dataset.actions.arm_action, mdp.RelativeJointPositionActionCfg) + assert reset_dataset.actions.arm_action.scale == pytest.approx(0.015) + assert reset_dataset.actions.gripper_action.use_incremental_target is False + assert reset_dataset.actions.gripper_action.binary_threshold == pytest.approx(0.0) + + assert reverse.rewards.success.func is mdp.pour_success_bonus + assert reverse.rewards.failure.func is mdp.terminal_failure + assert reverse.rewards.failure.params == {} + assert reverse.terminations.success.func is mdp.stable_pour_success + assert "terminate" not in reverse.terminations.lost_grasp.params + assert "terminate" not in reverse.terminations.spill.params + assert reverse.terminations.time_out.func is mdp.unsuccessful_time_out + assert reverse.actions.arm_action.scale == pytest.approx(0.08) + assert reverse.actions.gripper_action.use_incremental_target is False + assert reverse.actions.gripper_action.binary_threshold is None + assert reverse.episode_length_s == pytest.approx(5.0) + assert reverse.is_finite_horizon is True + + +def test_reset_dataset_ppo_is_calibrated_without_changing_traditional_curriculum(): + reset_dataset = FrankaPourResetDatasetPPORunnerCfg() + reverse = FrankaPourPPORunnerCfg() + + assert reset_dataset.num_steps_per_env == 32 + assert reset_dataset.save_interval == 25 + assert reset_dataset.resume is False + assert reset_dataset.actor.hidden_dims == [512, 256, 128, 64] + assert reset_dataset.critic.hidden_dims == [512, 256, 128, 64] + assert reset_dataset.actor.obs_normalization is False + assert reset_dataset.critic.obs_normalization is False + assert reset_dataset.actor.distribution_cfg.class_name == "HeteroscedasticGaussianDistribution" + assert reset_dataset.actor.distribution_cfg.init_std == pytest.approx(0.25) + assert reset_dataset.actor.distribution_cfg.std_range == pytest.approx((0.05, 0.75)) + assert reset_dataset.actor.distribution_cfg.std_type == "log" + assert reset_dataset.algorithm.entropy_coef == pytest.approx(1.0e-3) + assert reset_dataset.algorithm.gamma == pytest.approx(0.99 ** (1.0 / 3.0)) + assert reset_dataset.algorithm.lam == pytest.approx(0.95 ** (1.0 / 3.0)) + assert reset_dataset.algorithm.num_learning_epochs == 5 + assert reset_dataset.algorithm.num_mini_batches == 4 + assert reset_dataset.algorithm.learning_rate == pytest.approx(1.0e-4) + assert reset_dataset.algorithm.schedule == "fixed" + + assert reverse.num_steps_per_env == 32 + assert reverse.actor.hidden_dims == [256, 128, 64] + assert reverse.critic.hidden_dims == [256, 128, 64] + assert reverse.actor.obs_normalization is False + assert reverse.critic.obs_normalization is False + assert reverse.actor.distribution_cfg.init_std == pytest.approx(0.1) + assert reverse.algorithm.entropy_coef == pytest.approx(1.0e-3) + assert reverse.algorithm.gamma == pytest.approx(0.99) + assert reverse.algorithm.lam == pytest.approx(0.95) + + +@pytest.mark.parametrize( + "field, value, error_type, message", + [ + ("reset_dataset_path", "", ValueError, "reset_dataset_path must be nonempty"), + ("reset_dataset_path", None, TypeError, "reset_dataset_path must be a string"), + ("reset_dataset_content_sha256", "a" * 63, ValueError, "lowercase SHA-256"), + ("reset_dataset_content_sha256", "A" * 64, ValueError, "lowercase SHA-256"), + ], +) +def test_reset_dataset_rejects_invalid_configuration(field, value, error_type, message): + cfg = FrankaPourEnvCfg_RESET_DATASET() + setattr(cfg, field, value) + + with pytest.raises(error_type, match=message): + cfg.finalize() + + +@pytest.mark.parametrize( + "field, value", + [ + ("target_success_rate", 0.0), + ("target_success_rate", 1.0), + ("target_success_rate", float("nan")), + ("temperature", 0.0), + ("temperature", float("nan")), + ("history_capacity", 0), + ("history_capacity", 0.5), + ("history_capacity", False), + ("prior_strength", 0.0), + ("prior_strength", float("inf")), + ("initial_frontier_size", 0), + ("initial_frontier_size", 0.5), + ("probe_size", -1), + ("probe_size", 0.5), + ("probe_fraction", 1.0), + ("probe_fraction", float("nan")), + ("replay_fraction", 1.0), + ("replay_fraction", float("nan")), + ("frontier_evidence", 0.0), + ("frontier_evidence", float("inf")), + ], +) +def test_reset_dataset_rejects_invalid_sampler_configuration(field, value): + cfg = FrankaPourEnvCfg_RESET_DATASET() + setattr(cfg.reset_dataset_sampler, field, value) + + with pytest.raises(ValueError, match=field): + cfg.finalize() + + +def test_reset_dataset_accepts_custom_path_and_content_hash(): + cfg = FrankaPourEnvCfg_RESET_DATASET() + cfg.reset_dataset_path = "datasets/custom/reset_dataset.pt" + cfg.reset_dataset_content_sha256 = "a" * 64 + + resolved = cfg.finalize() + + assert resolved.reset_dataset_path == "datasets/custom/reset_dataset.pt" + assert resolved.reset_dataset_content_sha256 == "a" * 64 + + +def test_reset_dataset_top_grasp_filter_is_frozen_playback_only(): + cfg = FrankaPourEnvCfg_RESET_DATASET( + reset_dataset_top_grasp_count=64, + curriculum_freeze=True, + ).finalize() + assert cfg.reset_dataset_top_grasp_count == 64 + + with pytest.raises(ValueError, match="positive integer"): + FrankaPourEnvCfg_RESET_DATASET( + reset_dataset_top_grasp_count=0, + curriculum_freeze=True, + ).finalize() + with pytest.raises(ValueError, match="curriculum_freeze=True"): + FrankaPourEnvCfg_RESET_DATASET(reset_dataset_top_grasp_count=64).finalize() + + +@pytest.mark.parametrize("threshold", [False, float("nan"), -1.0, 1.0]) +def test_reset_dataset_rejects_invalid_binary_gripper_threshold(threshold): + cfg = FrankaPourEnvCfg_RESET_DATASET() + cfg.actions.gripper_action.binary_threshold = threshold + + with pytest.raises(ValueError, match="binary_threshold"): + cfg.finalize() + + +def test_reset_dataset_rejects_binary_incremental_gripper_combination(): + cfg = FrankaPourEnvCfg_RESET_DATASET() + cfg.actions.gripper_action.use_incremental_target = True + + with pytest.raises(ValueError, match="mutually exclusive"): + cfg.finalize() + + +@pytest.mark.parametrize( + "term_name, parameter, value", + [ + ("reach", "std", 0.0), + ("goal_distance", "std", float("nan")), + ("joint_velocity", "max_velocity", 0.0), + ], +) +def test_reset_dataset_rejects_invalid_general_reward(term_name, parameter, value): + cfg = FrankaPourEnvCfg_RESET_DATASET() + getattr(cfg.rewards, term_name).params[parameter] = value + + with pytest.raises(ValueError, match=term_name): + cfg.finalize() + + +def test_finalize_resynchronizes_overridden_gripper_action_bound(): + cfg = FrankaPourEnvCfg() + cfg.gripper_open_pos = 0.035 + cfg.success_max_tcp_distance = 0.02 + + resolved = cfg.finalize() + + assert resolved.actions.gripper_action.open_position == pytest.approx(0.035) + assert resolved.actions.gripper_action.default_position == pytest.approx(resolved.gripper_preload_pos) + assert resolved.actions.gripper_action.scale == pytest.approx(0.035 - resolved.gripper_preload_pos) + assert resolved.scene.robot.init_state.joint_pos["panda_finger_joint.*"] == pytest.approx(0.035) + expected_contact_command = resolved.cup_grasp_box_half[1] - resolved.actions.gripper_action.contact_min_deflection + assert resolved.rewards.task_progress.params["max_gripper_command"] == pytest.approx(expected_contact_command) + assert resolved.rewards.delivered.params["max_gripper_command"] == pytest.approx(expected_contact_command) + assert resolved.terminations.lost_grasp.params["max_gripper_command"] == pytest.approx(expected_contact_command) + assert resolved.terminations.success.params["max_gripper_command"] == pytest.approx(expected_contact_command) + + +@pytest.mark.parametrize("command", [0.023, 0.028]) +def test_finalize_rejects_gripper_success_commands_without_guaranteed_preload(command): + cfg = FrankaPourEnvCfg() + cfg.success_max_gripper_command = command + + with pytest.raises(ValueError, match="success_max_gripper_command"): + cfg.finalize() + + +def test_backward_curriculum_rejects_invalid_stage_and_arm_overrides(): + invalid_stage = FrankaPourEnvCfg() + invalid_stage.curriculum_start_stage = len(invalid_stage.curriculum_stage_names) + with pytest.raises(ValueError, match="curriculum_start_stage"): + invalid_stage.finalize() + + invalid_arm = FrankaPourEnvCfg() + invalid_arm.curriculum_pour_arm_q = (10.0,) * 7 + with pytest.raises(ValueError, match="outside"): + invalid_arm.finalize() + + invalid_replay = FrankaPourEnvCfg() + invalid_replay.curriculum_previous_stage_replay_fraction = 1.0 + with pytest.raises(ValueError, match="curriculum_previous_stage_replay_fraction"): + invalid_replay.finalize() + + invalid_entry_replay = FrankaPourEnvCfg() + invalid_entry_replay.curriculum_frontier_entry_replay_fraction = 0.05 + with pytest.raises(ValueError, match="curriculum_frontier_entry_replay_fraction"): + invalid_entry_replay.finalize() + + invalid_cohort_count = FrankaPourEnvCfg() + invalid_cohort_count.curriculum_min_reset_cohorts_per_stage = -1.0 + with pytest.raises(ValueError, match="curriculum_min_reset_cohorts_per_stage"): + invalid_cohort_count.finalize() + + invalid_transport = FrankaPourEnvCfg() + invalid_transport.curriculum_transport_reset_fractions = (0.75, 0.25) + with pytest.raises(ValueError, match="curriculum_transport_reset_fractions"): + invalid_transport.finalize() + + invalid_grasp_approach = FrankaPourEnvCfg() + invalid_grasp_approach.curriculum_grasp_approach_fractions = (0.75, 0.50, 0.70, 0.25, 0.125, 0.0) + with pytest.raises(ValueError, match="curriculum_grasp_approach_fractions"): + invalid_grasp_approach.finalize() + + unscreened_grasp_approach = FrankaPourEnvCfg() + unscreened_grasp_approach.curriculum_grasp_approach_fractions = (0.70, 0.50, 0.375, 0.25, 0.125, 0.0) + with pytest.raises(ValueError, match="eighth-segment"): + unscreened_grasp_approach.finalize() + + invalid_randomization_threshold = FrankaPourEnvCfg() + invalid_randomization_threshold.curriculum_randomization_promotion_threshold = 0.0 + with pytest.raises(ValueError, match="curriculum_randomization_promotion_threshold"): + invalid_randomization_threshold.finalize() + + invalid_randomization_threshold.curriculum_randomization_promotion_threshold = 0.9 + with pytest.raises(ValueError, match="curriculum_randomization_promotion_threshold"): + invalid_randomization_threshold.finalize() + + invalid_randomization_start = FrankaPourEnvCfg() + invalid_randomization_start.curriculum_randomization_start_level = len( + invalid_randomization_start.curriculum_randomization_extent_levels + ) + with pytest.raises(ValueError, match="curriculum_randomization_start_level"): + invalid_randomization_start.finalize() + + +@pytest.mark.parametrize("field,value", [("scale", 0.0), ("use_zero_offset", False), ("joint_names", ["panda_joint1"])]) +def test_backward_curriculum_rejects_invalid_relative_arm_actions(field, value): + cfg = FrankaPourEnvCfg() + setattr(cfg.actions.arm_action, field, value) + + with pytest.raises(ValueError, match="Arm|Panda|relative"): + cfg.finalize() + + +@pytest.mark.parametrize( + "levels", + [ + (), + (-0.1, 1.0), + (0.5, float("inf"), 1.0), + (0.5, 0.5, 1.0), + (0.75, 0.5, 1.0), + (0.5, 1.0), + (0.5, 0.9), + (0.5, 1.1), + ], +) +def test_randomized_curriculum_rejects_invalid_randomization_extent_levels(levels): + cfg = FrankaPourEnvCfg() + cfg.curriculum_randomization_extent_levels = levels + + with pytest.raises(ValueError, match="curriculum_randomization_extent_levels"): + cfg.finalize() + + +@pytest.mark.parametrize( + "fractions", + [ + (0.1, 0.35, 0.75, 1.0), + (0.0, 0.35, 0.75), + (0.0, 0.35, 1.1, 1.0), + (0.0, 0.75, 0.35, 1.0), + (0.0, 0.35, 0.75, 0.9), + ], +) +def test_randomized_curriculum_rejects_invalid_independent_reset_fractions(fractions): + cfg = FrankaPourEnvCfg() + cfg.curriculum_independent_arm_fraction_levels = fractions + + with pytest.raises(ValueError, match="curriculum_independent_arm_fraction_levels"): + cfg.finalize() + + +@pytest.mark.parametrize( + "parameter,value", + [ + ("target_tilt", 0.0), + ("target_tilt", math.pi), + ("pour_direction_xy", (0.0, 0.0)), + ("pour_direction_xy", (0.0,)), + ("alignment_radius", 0.0), + ("active_through_stage", 8), + ("discount_factor", 0.0), + ], +) +def test_tilt_curriculum_rejects_invalid_configuration(parameter, value): + cfg = FrankaPourEnvCfg() + cfg.rewards.task_progress.params[parameter] = value + + with pytest.raises(ValueError, match=parameter): + cfg.finalize() + + +@pytest.mark.parametrize( + "field,value", + [ + ("curriculum_randomized_source_position_range", (-0.1, 0.1)), + ("curriculum_randomized_source_radius_range", (0.68, 0.38)), + ("curriculum_randomized_source_azimuth_range", 0.0), + ("curriculum_randomized_source_xy_correlation", -0.01), + ("curriculum_randomized_source_xy_correlation", 1.0), + ("curriculum_randomized_source_xy_correlation", float("inf")), + ("curriculum_randomized_carry_position_range", (0.31, 0.10)), + ("curriculum_randomized_source_yaw_range", math.pi / 2.0), + ("curriculum_randomized_target_position_range", (0.05, float("inf"))), + ("curriculum_randomized_cup_clearance", -0.01), + ("curriculum_randomized_reset_tcp_standoff", (0.0, 0.05)), + ("curriculum_randomized_reset_tcp_jitter", (0.01, -0.01, 0.01)), + ("curriculum_randomized_reset_tcp_rotation_angle_range", (-0.1, 0.2)), + ("curriculum_randomized_reset_tcp_rotation_angle_range", (0.3, 0.2)), + ("curriculum_randomized_reset_tcp_rotation_angle_range", (0.0, math.pi)), + ("curriculum_randomized_reset_tcp_min_grasp_distance", 0.0), + ("curriculum_randomized_reset_joint6_max", 0.0), + ("curriculum_randomized_reset_joint6_max", float("inf")), + ("curriculum_randomized_min_source_cell_fraction", 0.0), + ("curriculum_randomized_min_source_cell_fraction", 1.01), + ("curriculum_randomized_min_source_cell_fraction", float("inf")), + ("curriculum_randomized_min_reset_variants_per_source", 0), + ("curriculum_randomized_min_reset_variants_per_source", 12), + ("curriculum_grasp_descent_overshoot", -0.001), + ("curriculum_randomized_reset_ik_grid_size", 1), + ("curriculum_randomized_reset_ik_samples_per_source", 1), + ("curriculum_randomized_reset_ik_samples_per_source", 4), + ("curriculum_randomized_reset_ik_iterations", 0), + ], +) +def test_randomized_curriculum_rejects_invalid_configuration(field, value): + cfg = FrankaPourEnvCfg() + setattr(cfg, field, value) + + with pytest.raises(ValueError, match=field): + cfg.finalize() + + +def test_randomized_curriculum_rejects_ranges_without_collision_free_cup_placement(): + cfg = FrankaPourEnvCfg() + cfg.curriculum_randomized_source_radius_range = None + cfg.curriculum_randomized_source_position_range = (0.20, 0.50) + + with pytest.raises(ValueError, match="no collision-free target y-position"): + cfg.finalize() + + +def test_randomized_curriculum_rejects_standoff_not_opposite_tool_axis(): + cfg = FrankaPourEnvCfg() + cfg.curriculum_randomized_reset_tcp_standoff = (0.0, 0.0, 0.12) + + with pytest.raises(ValueError, match="antiparallel"): + cfg.finalize() + + +def test_randomized_curriculum_rejects_vertical_tool_orientation(): + cfg = FrankaPourEnvCfg() + cfg.cup_grasp_tcp_quat_c = (0.0, 0.0, 0.0, 1.0) + + with pytest.raises(ValueError, match="parallel to the table"): + cfg.finalize() + + +def test_randomized_curriculum_rejects_vertical_jaw_axis_with_horizontal_tool_axis(): + cfg = FrankaPourEnvCfg() + # Tool +Z still points along cup +X, but a 90-degree tool roll makes Panda jaw +Y vertical. + cfg.cup_grasp_tcp_quat_c = (0.5, 0.5, 0.5, 0.5) + + with pytest.raises(ValueError, match="jaw axis parallel to the table"): + cfg.finalize() + + +def test_legacy_symmetric_tcp_jitter_rejects_missing_table_clearance(): + cfg = FrankaPourEnvCfg() + cfg.curriculum_randomized_reset_tcp_offset_lower = None + cfg.curriculum_randomized_reset_tcp_offset_upper = None + cfg.curriculum_randomized_reset_tcp_jitter = (0.02, 0.04, 0.082) + + with pytest.raises(ValueError, match="above the table"): + cfg.finalize() + + +def test_asymmetric_tcp_offset_rejects_below_pregrasp_start(): + cfg = FrankaPourEnvCfg() + cfg.curriculum_randomized_reset_tcp_offset_lower = (-0.20, -0.10, -0.001) + + with pytest.raises(ValueError, match="must not place the initial TCP below"): + cfg.finalize() + + +def test_randomized_curriculum_rejects_tcp_box_below_minimum_grasp_distance(): + cfg = FrankaPourEnvCfg() + cfg.curriculum_randomized_reset_tcp_min_grasp_distance = 0.13 + + with pytest.raises(ValueError, match="cannot guarantee curriculum_randomized_reset_tcp_min_grasp_distance"): + cfg.finalize() + + +@pytest.mark.parametrize( + "field,value", + [ + ("success_dwell_time_s", 0.0), + ("success_min_lift_height", 0.0), + ("success_max_tcp_distance", float("inf")), + ("success_max_gripper_width_error", 0.0), + ("state_bound_joint_position_margin", -0.1), + ("state_bound_max_joint_velocity", 0.0), + ("state_bound_max_cup_linear_velocity", float("inf")), + ("state_bound_max_cup_angular_velocity", 0.0), + ("curriculum_randomized_pour_clearance", -0.001), + ], +) +def test_success_and_state_bounds_reject_invalid_configuration(field, value): + cfg = FrankaPourEnvCfg() + setattr(cfg, field, value) + + with pytest.raises(ValueError, match=field): + cfg.finalize() + + +def test_finger_actuator_is_kept_and_teleop_disables_timeout(): + cfg = FrankaPourEnvCfg() + hand = cfg.scene.robot.actuators["panda_hand"] + assert hand.joint_names_expr == ["panda_finger_joint.*"] + assert hand.stiffness is None + assert hand.damping is None + assert hand.armature is None + assert FrankaPourEnvCfg_PLAY().scene.num_envs == 1 + teleop_cfg = FrankaPourEnvCfg_TELEOP() + assert teleop_cfg.terminations.time_out is None + assert teleop_cfg.actions.arm_action.clip == { + joint_name: limits + for joint_name, limits in zip( + teleop_cfg.actions.arm_action.joint_names, + pour_env_cfg.PANDA_ARM_JOINT_LIMITS, + strict=True, + ) + } + assert isinstance(teleop_cfg.finalize().actions.arm_action, mdp.CurriculumJointPositionActionCfg) + + +def _media_capacity(cfg): + return _media_entry(cfg).solver_cfg.max_active_cell_count + + +def _media_entry(cfg): + return next(entry for entry in cfg.sim.physics.solver_cfg.entries if entry.name == "media") + + +@pytest.mark.parametrize("num_envs", [1, 4, 8, 64, 200, 1024, 2048]) +def test_sparse_training_reserves_capturable_isolated_grid_capacity(num_envs): + cfg = FrankaPourEnvCfg() + cfg.scene.num_envs = num_envs + resolved = cfg.finalize() + + assert _media_entry(resolved).solver_cfg.grid_type == "sparse" + assert _media_capacity(resolved) == 512 * num_envs + solver_cfg = _media_entry(resolved).solver_cfg + assert solver_cfg.separate_worlds is True + assert solver_cfg.max_lower_node_count == max(32, 16 * num_envs) + assert solver_cfg.max_upper_node_count == max(32, (num_envs + 1) // 2) + assert resolved.scene.env_spacing == pytest.approx(cfg.scene.env_spacing) + assert resolved.sim.physics.use_cuda_graph is True + + +def test_play_uses_sparse_grid_capacity_and_honors_exact_override(): + play_cfg = FrankaPourEnvCfg_PLAY() + + assert _media_entry(play_cfg).solver_cfg.grid_type == "sparse" + assert _resolve_mpm_cell_cap(play_cfg) == 512 + play_cfg.mpm_cell_cap_override = 23456 + assert _resolve_mpm_cell_cap(play_cfg) == 23456 + + +def test_sparse_play_capacity_rejects_nonpositive_alignment(): + cfg = FrankaPourEnvCfg_PLAY() + cfg.mpm_cell_capacity_alignment = 0 + + with pytest.raises(ValueError, match="alignment must be positive"): + _resolve_mpm_cell_cap(cfg) + + +def test_mpm_uses_captured_sparse_training_and_play_configs(): + cfg = FrankaPourEnvCfg() + cfg.scene.num_envs = 200 + play_cfg = FrankaPourEnvCfg_PLAY().finalize() + + solver_cfg = _media_entry(cfg.finalize()).solver_cfg + play_solver_cfg = _media_entry(play_cfg).solver_cfg + + # PIC27 bounds collider nodes by particle samples while Q1 keeps both solves compact. + assert solver_cfg.velocity_basis == "Q1" + assert solver_cfg.collider_basis == "pic27" + assert solver_cfg.collider_velocity_mode == "forward" + assert solver_cfg.project_outside_colliders is False + assert solver_cfg.grid_type == "sparse" + assert solver_cfg.max_active_cell_count == 200 * 512 + assert solver_cfg.separate_worlds is True + assert play_solver_cfg.collider_basis == "pic27" + assert play_solver_cfg.grid_type == "sparse" + assert play_solver_cfg.grid_padding == 0 + assert play_solver_cfg.max_active_cell_count == 512 + assert play_solver_cfg.separate_worlds is True + assert play_cfg.sim.physics.use_cuda_graph is True + + +def test_coarse_voxel_resolution_finalizes_without_manual_hierarchy_overrides(): + cfg = FrankaPourEnvCfg(voxel_size=0.03) + cfg.scene.num_envs = 1 + + assert cfg.voxel_size == pytest.approx(0.03) + particles, _ = cup_cavity_lattice(cfg) + assert particles.shape[0] > 0 + + solver_cfg = _media_entry(cfg.finalize()).solver_cfg + + assert solver_cfg.voxel_size == pytest.approx(0.03) + assert solver_cfg.grid_type == "sparse" + assert solver_cfg.max_active_cell_count == 512 + + +def test_particle_workspace_bounds_contain_every_curriculum_reset(): + cfg = FrankaPourEnvCfg() + assert cfg.terminations.particle_out_of_bounds.func is mdp.particle_out_of_bounds + assert cfg.particle_max_velocity == pytest.approx(10.0) + lower = torch.tensor(cfg.particle_workspace_lower_bound) + upper = torch.tensor(cfg.particle_workspace_upper_bound) + local_particles = torch.from_numpy(cup_cavity_lattice(cfg)[0]) + + source_range = torch.tensor((*cfg.curriculum_randomized_source_position_range, 0.0)) + source_center = torch.tensor(cfg.cup_reset_pos) + for signed_range in (-source_range, source_range): + particles = local_particles + source_center + signed_range + assert bool(torch.all(particles >= lower)) + assert bool(torch.all(particles <= upper)) + + invalid = FrankaPourEnvCfg() + invalid.particle_workspace_lower_bound = (1.0, 0.0, 0.0) + invalid.particle_workspace_upper_bound = (0.0, 1.0, 1.0) + with pytest.raises(ValueError, match="particle_workspace"): + invalid.finalize() + + randomized_outside = FrankaPourEnvCfg() + randomized_outside.particle_workspace_upper_bound = (0.55, 1.0, 1.5) + with pytest.raises(ValueError, match="randomized source media"): + randomized_outside.finalize() + + invalid_velocity = FrankaPourEnvCfg() + invalid_velocity.particle_max_velocity = 0.0 + with pytest.raises(ValueError, match="particle_max_velocity"): + invalid_velocity.finalize() + + invalid_spill_fraction = FrankaPourEnvCfg() + invalid_spill_fraction.max_spill_fraction = 1.0 + with pytest.raises(ValueError, match="max_spill_fraction"): + invalid_spill_fraction.finalize() + + invalid_spill_height = FrankaPourEnvCfg() + invalid_spill_height.spill_table_height = float("inf") + with pytest.raises(ValueError, match="spill_table_height"): + invalid_spill_height.finalize() + + invalid_count_margin = FrankaPourEnvCfg() + invalid_count_margin.particle_count_margin = -0.001 + with pytest.raises(ValueError, match="particle_count_margin"): + invalid_count_margin.finalize() + + +def test_capacity_resolver_reads_fixed_grid_type_from_media_entry_without_mutation(): + cfg = FrankaPourEnvCfg() + media = _media_entry(cfg) + media.solver_cfg.grid_type = "fixed" + media.solver_cfg.max_active_cell_count = 120000 + arm = next(entry for entry in cfg.sim.physics.solver_cfg.entries if entry.name == "arm") + arm.solver_cfg.max_active_cell_count = 777 + + resolved_capacity = _resolve_mpm_cell_cap(cfg) + + assert resolved_capacity == 120000 + assert _media_capacity(cfg) == 120000 + assert arm.solver_cfg.max_active_cell_count == 777 + + +def test_media_selector_includes_spill_floor_without_unrelated_shapes(): + cfg = FrankaPourEnvCfg().finalize() + media = _media_entry(cfg) + model = SimpleNamespace( + body_label=[ + "/World/envs/env_0/TargetCup", + "/World/envs/env_0/SpillFloor", + "/World/envs/env_0/Robot/panda_hand", + ], + shape_count=4, + shape_body=torch.tensor([0, 1, 2, -1], dtype=torch.int32), + particle_count=3, + ) + + resolved = NewtonCoupler._resolve_entry(model, media, cfg.scene) + + assert resolved.bodies == [1] + assert resolved.shapes == [1] + assert resolved.particles == [0, 1, 2] + + +@pytest.mark.parametrize( + "task_id, cfg_name, runner_name", + [ + ("Isaac-Pour-Franka-v0", "FrankaPourEnvCfg", "FrankaPourPPORunnerCfg"), + ("Isaac-Pour-Franka-Play-v0", "FrankaPourEnvCfg_PLAY", "FrankaPourPPORunnerCfg"), + ("Isaac-Pour-Franka-Teleop-v0", "FrankaPourEnvCfg_TELEOP", None), + ( + "Isaac-Pour-Franka-Reset-Dataset-v0", + "FrankaPourEnvCfg_RESET_DATASET", + "FrankaPourResetDatasetPPORunnerCfg", + ), + ( + "Isaac-Pour-Franka-Reset-Dataset-Eval-v0", + "FrankaPourEnvCfg_RESET_DATASET_EVAL", + "FrankaPourResetDatasetPPORunnerCfg", + ), + ( + "Isaac-Pour-Franka-Reset-Dataset-Play-v0", + "FrankaPourEnvCfg_RESET_DATASET_PLAY", + "FrankaPourResetDatasetPPORunnerCfg", + ), + ( + "Isaac-Pour-Franka-Reset-Mixture-v0", + "FrankaPourEnvCfg_RESET_MIXTURE", + "FrankaPourResetMixturePPORunnerCfg", + ), + ( + "Isaac-Pour-Franka-Reset-Mixture-Eval-v0", + "FrankaPourEnvCfg_RESET_MIXTURE_EVAL", + "FrankaPourResetMixturePPORunnerCfg", + ), + ( + "Isaac-Pour-Franka-Reset-Mixture-Play-v0", + "FrankaPourEnvCfg_RESET_MIXTURE_PLAY", + "FrankaPourResetMixturePPORunnerCfg", + ), + ], +) +def test_gym_registration_exactly_matches_task_entry_points(task_id, cfg_name, runner_name): + spec = gym.spec(task_id) + assert spec.entry_point == "isaaclab_tasks.contrib.franka_pour.pour_env:FrankaPourEnv" + assert spec.disable_env_checker is True + expected_kwargs = { + "env_cfg_entry_point": f"isaaclab_tasks.contrib.franka_pour.pour_env_cfg:{cfg_name}", + } + if runner_name is not None: + expected_kwargs["rsl_rl_cfg_entry_point"] = ( + f"isaaclab_tasks.contrib.franka_pour.config.franka.agents.rsl_rl_ppo_cfg:{runner_name}" + ) + assert spec.kwargs == expected_kwargs + + +def test_deprecated_reset_mixture_names_alias_reset_dataset_implementation(): + assert FrankaPourResetMixturePPORunnerCfg is FrankaPourResetDatasetPPORunnerCfg + assert FrankaPourEnvCfg_RESET_MIXTURE is FrankaPourEnvCfg_RESET_DATASET + assert FrankaPourEnvCfg_RESET_MIXTURE_EVAL is FrankaPourEnvCfg_RESET_DATASET_EVAL + assert FrankaPourEnvCfg_RESET_MIXTURE_PLAY is FrankaPourEnvCfg_RESET_DATASET_PLAY + assert mdp.PourResetMixture is mdp.PourResetDatasetCurriculum + + +def test_task_source_does_not_traverse_private_solver_state(): + source_path = Path(franka_pour.__file__).with_name("pour_env.py") + tree = ast.parse(source_path.read_text(encoding="utf-8"), filename=str(source_path)) + private_manager_attrs = { + node.attr + for node in ast.walk(tree) + if isinstance(node, ast.Attribute) + and node.attr.startswith("_") + and isinstance(node.value, ast.Name) + and node.value.id in {"NewtonManager", "NewtonCoupler"} + } + assert private_manager_attrs == set() + + +def test_tcp_pose_uses_public_robot_pose_data_and_configured_offset(): + source_path = Path(franka_pour.__file__).with_name("pour_env.py") + tree = ast.parse(source_path.read_text(encoding="utf-8"), filename=str(source_path)) + tcp_pose = next(node for node in ast.walk(tree) if isinstance(node, ast.FunctionDef) and node.name == "tcp_pose_e") + + attributes = {node.attr for node in ast.walk(tcp_pose) if isinstance(node, ast.Attribute)} + assert "get_term" not in attributes + assert "_compute_frame_pose" not in attributes + assert {"body_link_pose_w", "root_link_pose_w"} <= attributes + assert {"subtract_frame_transforms", "combine_frame_transforms"} <= attributes + + setup = next( + node for node in ast.walk(tree) if isinstance(node, ast.FunctionDef) and node.name == "_setup_after_physics" + ) + setup_source = ast.unparse(setup) + assert "cfg.tcp_body_name" in setup_source + assert "cfg.tcp_offset_pos" in setup_source + assert "cfg.tcp_offset_rot" in setup_source + assert "cfg.actions.arm_action.body_name" not in setup_source + assert "cfg.actions.arm_action.body_offset" not in setup_source + + +def test_media_refill_is_batched_on_device_without_host_readback(): + source_path = Path(franka_pour.__file__).with_name("pour_env.py") + tree = ast.parse(source_path.read_text(encoding="utf-8"), filename=str(source_path)) + sample_media = next( + node for node in ast.walk(tree) if isinstance(node, ast.FunctionDef) and node.name == "_sample_cup_media" + ) + + attributes = {node.attr for node in ast.walk(sample_media) if isinstance(node, ast.Attribute)} + assert {"cpu", "numpy", "tolist"}.isdisjoint(attributes) + assert "quat_apply" in attributes + + +def test_task_reset_uses_public_asset_writers_and_manager_lifecycle(): + source_path = Path(franka_pour.__file__).with_name("pour_env.py") + tree = ast.parse(source_path.read_text(encoding="utf-8"), filename=str(source_path)) + + function_names = {node.name for node in ast.walk(tree) if isinstance(node, ast.FunctionDef)} + assert "_reset_mpm_particle_state" not in function_names + + reset = next( + node for node in ast.walk(tree) if isinstance(node, ast.FunctionDef) and node.name == "reset_pour_scene" + ) + call_names = { + node.func.attr + for node in ast.walk(reset) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) + } + assert { + "write_joint_position_to_sim_index", + "write_joint_velocity_to_sim_index", + "write_root_pose_to_sim_index", + "write_root_velocity_to_sim_index", + "write_particle_pos_to_sim_index", + "write_particle_velocity_to_sim_index", + } <= call_names + assert "reset_solver_state" in call_names + assert {"eval_fk", "get_state_0", "get_state_1"}.isdisjoint(call_names) + + reset_attributes = {node.attr for node in ast.walk(reset) if isinstance(node, ast.Attribute)} + assert "body_link_pose_w" in reset_attributes + + manager_forward_calls = [ + node + for node in ast.walk(reset) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr in {"forward", "forward_pending"} + ] + assert manager_forward_calls == [] diff --git a/source/isaaclab_tasks/test/contrib/test_franka_pour_mdp.py b/source/isaaclab_tasks/test/contrib/test_franka_pour_mdp.py new file mode 100644 index 000000000000..12504fa1bdc6 --- /dev/null +++ b/source/isaaclab_tasks/test/contrib/test_franka_pour_mdp.py @@ -0,0 +1,1754 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Unit tests for the Franka pour reward terms (no simulator).""" + +import math +import warnings +from types import SimpleNamespace + +import pytest +import torch + +from isaaclab.managers import RewardTermCfg + +import isaaclab_tasks.contrib.franka_pour.mdp as mdp_api +from isaaclab_tasks.contrib.franka_pour.mdp import observations, rewards, terminations + + +class FakeActionManager: + def __init__(self, num_envs: int): + self.action = torch.zeros((num_envs, 8)) + self._terms = { + "gripper_action": SimpleNamespace( + commanded_position=torch.full((num_envs, 1), 0.04), + bilateral_preload=torch.ones(num_envs, dtype=torch.bool), + bilateral_contact=torch.ones(num_envs, dtype=torch.bool), + contact_deflection=torch.zeros((num_envs, 2)), + ) + } + + def get_term(self, name: str): + return self._terms[name] + + +def test_terminal_failure_pulses_once_for_overlapping_failure_predicates(): + termination_manager = FakeTerminationManager(4) + termination_manager.terminated[:] = torch.tensor([True, True, False, False]) + termination_manager.time_outs[:] = torch.tensor([False, False, True, False]) + termination_manager._terms["success"][:] = torch.tensor([False, True, False, False]) + env = SimpleNamespace(step_dt=0.02, termination_manager=termination_manager) + + penalty = rewards.terminal_failure(env) + + torch.testing.assert_close(penalty, torch.tensor([50.0, 0.0, 50.0, 0.0])) + + +def test_general_reach_reward_is_bounded_monotonic_and_finite(): + env = SimpleNamespace( + tcp_pos_e=lambda: torch.tensor(((0.0, 0.0, 0.0), (0.5, 0.0, 0.0), (float("nan"), 0.0, 0.0))), + cup_grasp_point_e=lambda: torch.zeros((3, 3)), + ) + + quality = rewards.tcp_cup_distance_tanh(env, std=1.0) + + assert quality[0] == 1.0 + assert 0.0 < quality[1] < quality[0] + assert quality[2] == 0.0 + with pytest.raises(ValueError, match="std must be finite and positive"): + rewards.tcp_cup_distance_tanh(env, std=0.0) + + +def test_general_media_goal_reward_measures_distance_to_target_cavity(): + identity = torch.tensor((0.0, 0.0, 0.0, 1.0)) + target_pose = torch.cat((torch.zeros((3, 3)), identity.repeat(3, 1)), dim=-1) + particles = torch.tensor( + ( + ((0.0, 0.0, 0.05), (0.02, -0.02, 0.08)), + ((0.20, 0.0, 0.05), (0.20, 0.0, 0.05)), + ((float("inf"), 0.0, 0.05), (float("inf"), 0.0, 0.05)), + ) + ) + env = SimpleNamespace( + cfg=SimpleNamespace(particle_count_margin=0.0), + _source_inner_hi_t=torch.tensor((0.05, 0.05, 0.10)), + _target_inner_lo_t=torch.tensor((-0.05, -0.05, 0.01)), + _target_inner_hi_t=torch.tensor((0.05, 0.05, 0.10)), + cup_pose_e=lambda: target_pose, + target_pose_e=lambda: target_pose, + particle_pos_e=lambda: particles, + particle_region_masks=lambda: tuple(torch.zeros((3, 2), dtype=torch.bool) for _ in range(3)), + ) + + quality = rewards.media_target_distance_tanh(env, std=1.0) + + assert quality[0] == 1.0 + assert 0.0 < quality[1] < quality[0] + assert quality[2] == 0.0 + with pytest.raises(ValueError, match="std must be finite and positive"): + rewards.media_target_distance_tanh(env, std=float("nan")) + + +def test_general_media_goal_reward_requires_release_from_nested_source(): + identity = torch.tensor((0.0, 0.0, 0.0, 1.0)) + pose = torch.cat((torch.zeros(3), identity)).unsqueeze(0).repeat(2, 1) + particles = torch.tensor((((0.0, 0.0, 0.05),), ((0.0, 0.0, 0.05),))) + in_source = torch.tensor(((True,), (False,))) + env = SimpleNamespace( + cfg=SimpleNamespace(particle_count_margin=0.0), + _source_inner_hi_t=torch.tensor((0.05, 0.05, 0.10)), + _target_inner_lo_t=torch.tensor((-0.10, -0.10, 0.0)), + _target_inner_hi_t=torch.tensor((0.10, 0.10, 0.15)), + cup_pose_e=lambda: pose, + target_pose_e=lambda: pose, + particle_pos_e=lambda: particles, + particle_region_masks=lambda: ( + in_source, + ~in_source, + torch.zeros_like(in_source), + ), + ) + + quality = rewards.media_target_distance_tanh(env, std=0.1) + + assert quality[0] == pytest.approx(1.0 - math.tanh(0.5)) + assert quality[0] < quality[1] + assert quality[1] == 1.0 + + +def test_general_media_goal_reward_zeroes_spilled_particles(): + identity = torch.tensor((0.0, 0.0, 0.0, 1.0)) + pose = torch.cat((torch.zeros(3), identity)).unsqueeze(0) + particles = torch.tensor((((0.0, 0.0, 0.05), (0.0, 0.0, 0.05)),)) + false = torch.zeros((1, 2), dtype=torch.bool) + env = SimpleNamespace( + cfg=SimpleNamespace(particle_count_margin=0.0), + _source_inner_hi_t=torch.tensor((0.05, 0.05, 0.10)), + _target_inner_lo_t=torch.tensor((-0.10, -0.10, 0.0)), + _target_inner_hi_t=torch.tensor((0.10, 0.10, 0.15)), + cup_pose_e=lambda: pose, + target_pose_e=lambda: pose, + particle_pos_e=lambda: particles, + particle_region_masks=lambda: (false, false, torch.tensor(((False, True),))), + ) + + quality = rewards.media_target_distance_tanh(env, std=0.1) + + assert quality[0] == pytest.approx(0.5) + + +def test_general_media_goal_reward_uses_rotated_target_frame(): + half_sqrt = math.sqrt(0.5) + target_pose = torch.tensor(((0.0, 0.0, 0.0, 0.0, 0.0, half_sqrt, half_sqrt),)) + identity_pose = torch.tensor(((0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0),)) + particles = torch.tensor((((0.0, 0.03, 0.05),),)) + false = torch.zeros((1, 1), dtype=torch.bool) + env = SimpleNamespace( + cfg=SimpleNamespace(particle_count_margin=0.0), + _source_inner_hi_t=torch.tensor((0.05, 0.05, 0.10)), + _target_inner_lo_t=torch.tensor((-0.02, -0.10, 0.0)), + _target_inner_hi_t=torch.tensor((0.02, 0.10, 0.10)), + cup_pose_e=lambda: identity_pose, + target_pose_e=lambda: target_pose, + particle_pos_e=lambda: particles, + particle_region_masks=lambda: (false, false, false), + ) + + quality = rewards.media_target_distance_tanh(env, std=0.1) + + assert quality[0] == pytest.approx(1.0 - math.tanh(0.1)) + + +def test_general_distance_rewards_are_independent_of_reset_stage(): + identity = torch.tensor((0.0, 0.0, 0.0, 1.0)) + env = SimpleNamespace( + cfg=SimpleNamespace(particle_count_margin=0.0), + curriculum_stage=torch.tensor((15, 7, 6, 2)), + _source_inner_hi_t=torch.tensor((0.05, 0.05, 0.10)), + _target_inner_lo_t=torch.tensor((-0.05, -0.05, 0.01)), + _target_inner_hi_t=torch.tensor((0.05, 0.05, 0.10)), + tcp_pos_e=lambda: torch.full((4, 3), 0.1), + cup_grasp_point_e=lambda: torch.zeros((4, 3)), + cup_pose_e=lambda: torch.cat((torch.zeros((4, 3)), identity.repeat(4, 1)), dim=-1), + target_pose_e=lambda: torch.cat((torch.zeros((4, 3)), identity.repeat(4, 1)), dim=-1), + particle_pos_e=lambda: torch.full((4, 2, 3), 0.2), + particle_region_masks=lambda: tuple(torch.zeros((4, 2), dtype=torch.bool) for _ in range(3)), + ) + + reach = rewards.tcp_cup_distance_tanh(env) + goal_distance = rewards.media_target_distance_tanh(env) + + torch.testing.assert_close(reach, reach[0].expand_as(reach)) + torch.testing.assert_close(goal_distance, goal_distance[0].expand_as(goal_distance)) + + +def test_general_joint_velocity_penalty_is_finite_on_invalid_terminal_state(): + velocity = torch.tensor(((1.0, 2.0), (float("nan"), float("inf")), (100.0, -100.0))) + env = SimpleNamespace( + scene={"robot": SimpleNamespace(data=SimpleNamespace(joint_vel=SimpleNamespace(torch=velocity)))} + ) + asset_cfg = SimpleNamespace(name="robot", joint_ids=[0, 1]) + + penalty = rewards.finite_joint_velocity_l2(env, asset_cfg=asset_cfg, max_velocity=20.0) + + torch.testing.assert_close(penalty, torch.tensor((5.0, 800.0, 800.0))) + assert torch.isfinite(penalty).all() + + +class FakeTerminationManager: + def __init__(self, num_envs: int): + self.terminated = torch.zeros(num_envs, dtype=torch.bool) + self.time_outs = torch.zeros(num_envs, dtype=torch.bool) + self._terms = {"success": torch.zeros(num_envs, dtype=torch.bool)} + + @property + def active_terms(self) -> list[str]: + return list(self._terms) + + @property + def dones(self) -> torch.Tensor: + return self.terminated | self.time_outs + + def get_term(self, name: str) -> torch.Tensor: + return self._terms[name] + + +class FakeHeldDeliveryTracker: + """Small in-memory implementation of the task's held-delivery interface.""" + + def _init_held_delivery_tracker(self) -> None: + self.common_step_counter = 0 + self._target_entry_seen = torch.zeros((self.num_envs, self.num_particles), dtype=torch.bool) + self._held_delivered = torch.zeros_like(self._target_entry_seen) + self._held_delivery_tracker_step = -1 + + def update_held_delivery_tracker(self, held_pour: torch.Tensor) -> None: + if held_pour.shape != (self.num_envs,): + raise ValueError + if self._held_delivery_tracker_step == self.common_step_counter: + return + in_target = self.particles_in_target_mask() + first_entry = in_target & ~self._target_entry_seen + self._held_delivered |= first_entry & held_pour.unsqueeze(-1) + self._target_entry_seen |= in_target + self._held_delivery_tracker_step = self.common_step_counter + + def held_delivered_mask(self) -> torch.Tensor: + return self._held_delivered + + def current_held_delivered_mask(self) -> torch.Tensor: + return self._held_delivered & self.particles_in_target_mask() + + +class FakeEnv(FakeHeldDeliveryTracker): + """Minimal vectorized stand-in exposing the interface consumed by pure reward terms.""" + + def __init__(self): + self.num_envs = 4 + self.num_particles = 1000 + self.device = "cpu" + self.step_dt = 1.0 / 60.0 + self.pour_target_frac = torch.full((self.num_envs,), 0.9) + self.curriculum_stage = torch.full((self.num_envs,), 2, dtype=torch.long) + self.cfg = type( + "Cfg", + (), + { + "curriculum_stage_names": ("pour", "carry", "lift", "full"), + "max_spill_fraction": 0.10, + }, + )() + self.episode_succeeded = torch.zeros(self.num_envs, dtype=torch.bool) + self.ep_max_target_frac = torch.zeros(self.num_envs) + self._success_dwell_count = torch.zeros(self.num_envs, dtype=torch.long) + self._lost_grasp_dwell_count = torch.zeros(self.num_envs, dtype=torch.long) + self._lifted_grasp_seen = torch.zeros(self.num_envs, dtype=torch.bool) + self.termination_manager = FakeTerminationManager(self.num_envs) + self.gripper_open_width = 0.08 + self.gripper_grasp_width = 0.060 + self.cup_reset_height = 0.0 + self.action_manager = FakeActionManager(self.num_envs) + self._gripper_command = self.action_manager.get_term("gripper_action").commanded_position[:, 0] + + self._tcp = torch.tensor([[0.50, 0.00, 0.032], [0.50, 0.00, 0.032], [0.50, 0.00, 0.032], [0.20, 0.00, 0.032]]) + self._tcp_quat = torch.tensor([[0.0, 0.0, 0.0, 1.0]]).repeat(self.num_envs, 1) + self._grasp = torch.tensor([[0.50, 0.00, 0.032]]).repeat(self.num_envs, 1) + self._cup = torch.tensor( + [ + [0.50, 0.00, 0.00, 0.0, 0.0, 0.0, 1.0], + [0.50, 0.00, 0.12, 0.0, 0.0, 0.0, 1.0], + [0.50, -0.17, 0.12, 0.7071068, 0.0, 0.0, 0.7071068], + [0.20, 0.00, 0.00, 0.0, 0.0, 0.0, 1.0], + ] + ) + self._target = torch.tensor([[0.50, -0.18, 0.00, 0.0, 0.0, 0.0, 1.0]]).repeat(self.num_envs, 1) + self._width = torch.tensor([0.08, 0.060, 0.060, 0.060]) + self._src = torch.tensor([1000.0, 750.0, 400.0, 0.0]) + self._tgt = torch.tensor([0.0, 250.0, 500.0, 950.0]) + self._spill = torch.tensor([0.0, 0.0, 100.0, 50.0]) + self._arm_q = torch.zeros((self.num_envs, 7)) + self._init_held_delivery_tracker() + + def tcp_pos_e(self): + return self._tcp + + def tcp_pose_e(self): + return torch.cat((self._tcp, self._tcp_quat), dim=-1) + + def desired_grasp_tcp_quat_c(self): + return torch.tensor([[0.0, 0.0, 0.0, 1.0]]).repeat(self.num_envs, 1) + + def cup_grasp_point_e(self): + return self._grasp + + def cup_pose_e(self): + return self._cup + + def target_pose_e(self): + return self._target + + def gripper_width(self): + return self._width + + def arm_joint_pos(self): + return self._arm_q + + def count_in_target(self): + return self._tgt + + def particles_in_target_mask(self): + particle_ids = torch.arange(self.num_particles).unsqueeze(0) + return particle_ids < self._tgt.to(dtype=torch.long).unsqueeze(-1) + + def count_in_source(self): + return self._src + + def count_spilled(self): + return self._spill + + def spilled_fraction(self): + return self._spill / self.num_particles + + +def test_legacy_reward_api_remains_available_with_deprecation_warning(): + legacy_names = ( + "reach_cup", + "grasp_cup", + "lift_cup", + "lift_command_progress", + "align_cup_over_target", + "align_command_progress", + "tilt_over_target", + "tilt_command_progress", + ) + assert all(hasattr(mdp_api, name) for name in legacy_names) + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + result = mdp_api.reach_cup(FakeEnv()) + + assert result.shape == (4,) + assert any(item.category is DeprecationWarning for item in caught) + + +class FakeDeliveryEnv(FakeHeldDeliveryTracker): + """Four-particle environment used to exercise entry and reward idempotence.""" + + def __init__(self, num_envs: int = 2, num_particles: int = 4): + self.num_envs = num_envs + self.num_particles = num_particles + self.device = "cpu" + self.step_dt = 0.02 + self.pour_target_frac = torch.ones(num_envs) + self.termination_manager = FakeTerminationManager(num_envs) + self._target_mask = torch.zeros((num_envs, num_particles), dtype=torch.bool) + self._init_held_delivery_tracker() + + def particles_in_target_mask(self) -> torch.Tensor: + return self._target_mask + + +def _set_simple_env_held_state(env, held: torch.Tensor | None = None) -> None: + """Attach the source-grasp interface required by held-delivery shaping.""" + if held is None: + held = torch.ones(env.num_envs, dtype=torch.bool) + cup_pose = torch.zeros((env.num_envs, 7)) + cup_pose[:, 2] = 0.06 + tcp = torch.zeros((env.num_envs, 3)) + command = torch.where(held, 0.0, 0.04).unsqueeze(-1) + env.cup_reset_height = 0.0 + env.gripper_grasp_width = 0.06 + env.cup_pose_e = lambda: cup_pose + env.tcp_pos_e = lambda: tcp + env.cup_grasp_point_e = lambda: tcp + env.gripper_width = lambda: torch.full((env.num_envs,), 0.06) + env.action_manager = SimpleNamespace( + get_term=lambda name: SimpleNamespace(commanded_position=command), + ) + + +def test_particle_fractions_spill_and_success(): + env = FakeEnv() + assert torch.allclose(rewards.particles_in_target(env), torch.tensor([0.0, 0.25, 0.5, 0.95])) + assert torch.allclose(rewards.particles_in_source(env), torch.tensor([1.0, 0.75, 0.4, 0.0])) + assert torch.allclose(rewards.spilled_particles(env), torch.tensor([0.0, 0.0, 0.1, 0.05]), atol=1e-6) + fractions = observations.particle_fractions_obs(env) + torch.testing.assert_close( + fractions, + torch.tensor( + [ + [1.0, 0.0, 0.0, 0.0], + [0.75, 0.25, 0.0, 0.0], + [0.4, 0.5, 0.1, 0.0], + [0.0, 0.95, 0.05, 0.0], + ] + ), + ) + + +def test_stable_success_requires_consecutive_valid_steps_and_rejects_failures(): + env = FakeEnv() + env._cup[:, 2] = 0.06 + env._tcp[:] = env._grasp + env._width[:] = env.gripper_grasp_width + env._gripper_command[:] = 0.0 + env._tgt[:] = torch.tensor([950.0, 950.0, 0.0, 950.0]) + first = terminations.stable_pour_success(env, dwell_time_s=2.0 * env.step_dt) + env.common_step_counter += 1 + env._tgt[:] = torch.tensor([950.0, 0.0, 950.0, 950.0]) + env.termination_manager.terminated[3] = True + second = terminations.stable_pour_success(env, dwell_time_s=2.0 * env.step_dt) + + assert first.tolist() == [False, False, False, False] + assert second.tolist() == [True, False, False, False] + assert env.episode_succeeded.tolist() == [True, False, False, False] + torch.testing.assert_close(env.ep_max_target_frac, torch.tensor([0.95, 0.95, 0.95, 0.95])) + + +def test_immediate_pour_success_uses_only_current_target_fraction_and_failure_precedence(): + env = FakeEnv() + env.pour_target_frac[:] = 0.30 + env._tgt[:] = torch.tensor([299.0, 300.0, 950.0, 950.0]) + env.termination_manager.terminated[3] = True + + success = terminations.immediate_pour_success(env) + + assert success.tolist() == [False, True, True, False] + assert env.episode_succeeded.tolist() == [False, True, True, False] + assert env._success_dwell_count.tolist() == [0, 1, 1, 0] + torch.testing.assert_close(env.ep_max_target_frac, torch.tensor([0.299, 0.300, 0.950, 0.950])) + + +def test_stable_success_counter_resets_selected_worlds_only(): + env = FakeEnv() + env._cup[:, 2] = 0.06 + env._tcp[:] = env._grasp + env._width[:] = env.gripper_grasp_width + env._gripper_command[:] = 0.0 + env._tgt[:] = 950.0 + assert not bool(torch.any(terminations.stable_pour_success(env, dwell_time_s=2.0 * env.step_dt))) + env._success_dwell_count[0] = 0 + env.common_step_counter += 1 + success = terminations.stable_pour_success(env, dwell_time_s=2.0 * env.step_dt) + + assert success.tolist() == [False, True, True, True] + + +def test_stable_success_is_directly_reusable_as_a_state_reward(): + env = FakeEnv() + env._cup[:, 2] = 0.06 + env._tcp[:] = env._grasp + env._width[:] = env.gripper_grasp_width + env._gripper_command[:] = 0.0 + env._tgt[:] = 950.0 + + first = terminations.stable_pour_success(env, dwell_time_s=2.0 * env.step_dt) + env.common_step_counter += 1 + second = terminations.stable_pour_success(env, dwell_time_s=2.0 * env.step_dt) + + assert not bool(torch.any(first)) + assert bool(torch.all(second)) + assert env._success_dwell_count.tolist() == [2, 2, 2, 2] + assert env.episode_succeeded.tolist() == [True, True, True, True] + + +def test_nonterminating_success_context_tracks_state_and_remains_replay_discoverable(): + env = FakeEnv() + env._cup[:, 2] = 0.06 + env._tcp[:] = env._grasp + env._width[:] = env.gripper_grasp_width + env._gripper_command[:] = 0.0 + env._tgt[:] = 950.0 + + first = terminations.nonterminating_stable_pour_success(env, dwell_time_s=2.0 * env.step_dt) + env.common_step_counter += 1 + second = terminations.nonterminating_stable_pour_success(env, dwell_time_s=2.0 * env.step_dt) + + assert not bool(torch.any(first)) + assert not bool(torch.any(second)) + assert env.episode_succeeded.tolist() == [True, True, True, True] + assert rewards.sustained_pour_success(env, dwell_time_s=2.0 * env.step_dt).tolist() == [1.0, 1.0, 1.0, 1.0] + + # Record/replay tools remove the managed term before invoking its configured callable. + env.termination_manager._terms.pop("success") + replay_success = terminations.nonterminating_stable_pour_success(env, dwell_time_s=2.0 * env.step_dt) + assert bool(torch.all(replay_success)) + + +def test_sustained_success_reward_is_current_unit_state_not_terminal_pulse(): + env = FakeEnv() + env._success_dwell_count[:] = torch.tensor([0, 1, 2, 3]) + + success = rewards.sustained_pour_success(env, dwell_time_s=2.0 * env.step_dt) + + assert success.tolist() == [0.0, 0.0, 1.0, 1.0] + assert float(success.max()) == 1.0 + + +def test_unsuccessful_timeout_excludes_same_step_success(): + env = SimpleNamespace( + episode_length_buf=torch.tensor([10, 10, 9, 10]), + max_episode_length=10, + episode_succeeded=torch.tensor([True, False, False, True]), + ) + + timed_out = terminations.unsuccessful_time_out(env) + + assert timed_out.tolist() == [False, True, False, False] + + +def test_stable_success_requires_a_preloaded_held_and_lifted_source(): + env = FakeEnv() + env._tgt[:] = 950.0 + env._cup[:, 2] = 0.06 + env._tcp[:] = env._grasp + env._width[:] = env.gripper_grasp_width + env._gripper_command[:] = 0.0 + + env._cup[0, 2] = 0.0 + env._tcp[1] = env._grasp[1] + torch.tensor([0.04, 0.0, 0.0]) + env._gripper_command[2] = 0.04 + + success = terminations.stable_pour_success(env, dwell_time_s=env.step_dt) + + assert success.tolist() == [False, False, False, True] + + +def test_stable_success_rejects_a_geometrically_plausible_unilateral_grasp(): + env = FakeEnv() + env._tgt[:] = 950.0 + env._cup[:, 2] = 0.06 + env._tcp[:] = env._grasp + env._width[:] = env.gripper_grasp_width + env._gripper_command[:] = 0.025 + env.action_manager.get_term("gripper_action").bilateral_contact[0] = False + + success = terminations.stable_pour_success(env, dwell_time_s=env.step_dt) + + assert success.tolist() == [False, True, True, True] + + +def test_gripper_contact_observation_preserves_per_finger_asymmetry(): + env = FakeEnv() + env.action_manager.get_term("gripper_action").contact_deflection[:] = torch.tensor( + [[0.002, 0.002], [0.003, 0.0002], [0.0, 0.0], [0.001, 0.004]] + ) + + contact = observations.gripper_contact_obs(env) + + torch.testing.assert_close( + contact, + torch.tensor([[0.002, 0.002], [0.003, 0.0002], [0.0, 0.0], [0.001, 0.004]]), + ) + + +def test_policy_grasp_geometry_is_cup_relative_and_quaternion_sign_invariant(): + half_sqrt = math.sqrt(0.5) + cup_pose = torch.tensor( + [ + [1.0, 2.0, 0.0, 0.0, 0.0, 0.0, 1.0], + [1.0, 2.0, 0.0, 0.0, 0.0, half_sqrt, half_sqrt], + ] + ) + tcp_pose = torch.tensor( + [ + [1.0, 2.0, 0.03, 0.0, 0.0, 0.0, -1.0], + [1.02, 2.01, 0.03, 0.0, 0.0, -half_sqrt, -half_sqrt], + ] + ) + target_pose = torch.tensor( + [ + [1.2, 2.1, 0.0, 0.0, 0.0, 0.0, 1.0], + [0.9, 2.2, 0.0, 0.0, 0.0, 0.0, 1.0], + ] + ) + env = SimpleNamespace( + cfg=SimpleNamespace(cup_grasp_height=0.03), + cup_pose_e=lambda: cup_pose, + tcp_pose_e=lambda: tcp_pose, + tcp_pos_e=lambda: tcp_pose[:, :3], + target_pose_e=lambda: target_pose, + desired_grasp_tcp_quat_c=lambda: torch.tensor([[0.0, 0.0, 0.0, 1.0]]).repeat(2, 1), + ) + + torch.testing.assert_close( + observations.tcp_to_grasp_position_c_obs(env), + torch.tensor([[0.0, 0.0, 0.0], [-0.01, 0.02, 0.0]]), + rtol=0.0, + atol=1.0e-6, + ) + torch.testing.assert_close( + observations.grasp_to_tcp_quat_obs(env), + torch.tensor([[0.0, 0.0, 0.0, 1.0], [0.0, 0.0, 0.0, 1.0]]), + rtol=0.0, + atol=1.0e-6, + ) + torch.testing.assert_close( + observations.target_position_c_obs(env), + torch.tensor([[0.2, 0.1, 0.0], [0.2, 0.1, 0.0]]), + rtol=0.0, + atol=1.0e-6, + ) + torch.testing.assert_close( + observations.tcp_pose_obs(env)[:, 3:7], + torch.tensor([[0.0, 0.0, 0.0, 1.0], [0.0, 0.0, half_sqrt, half_sqrt]]), + rtol=0.0, + atol=1.0e-6, + ) + + +def test_individual_finger_and_held_delivery_observations_preserve_state(): + finger_position = torch.tensor([[0.01, 0.02], [0.03, 0.04]]) + finger_velocity = torch.tensor([[0.1, -0.2], [0.3, -0.4]]) + held = torch.tensor([[True, False, True, False], [False, False, True, False]]) + env = SimpleNamespace( + num_particles=4, + finger_joint_pos=lambda: finger_position, + finger_joint_vel=lambda: finger_velocity, + held_delivered_mask=lambda: held, + ) + + torch.testing.assert_close(observations.finger_position_obs(env), finger_position) + torch.testing.assert_close(observations.finger_velocity_obs(env), finger_velocity) + torch.testing.assert_close(observations.held_delivery_history_obs(env), torch.tensor([[0.5], [0.25]])) + + +def test_grasp_command_comparison_tolerates_float32_filter_roundoff(): + env = FakeEnv() + env._cup[:, 2] = 0.06 + env._tcp[:] = env._grasp + env._width[:] = env.gripper_grasp_width + env._gripper_command[:] = 0.025 + 5.0e-7 + + _, preloaded, lifted = terminations.source_grasp_milestones( + env, + min_lift_height=0.05, + max_tcp_distance=0.018, + max_gripper_width_error=0.006, + max_gripper_command=0.025, + ) + + assert bool(torch.all(preloaded)) + assert bool(torch.all(lifted)) + env._gripper_command[0] = 0.025 + 2.0e-6 + _, preloaded, _ = terminations.source_grasp_milestones( + env, + min_lift_height=0.05, + max_tcp_distance=0.018, + max_gripper_width_error=0.006, + max_gripper_command=0.025, + ) + assert not bool(preloaded[0]) + + +def test_stable_success_cannot_recover_particles_first_delivered_while_unheld(): + env = FakeEnv() + env.pour_target_frac[:] = 0.4 + env._tgt[:] = 950.0 + env._cup[:, 2] = 0.0 + env._tcp[:] = env._grasp + env._width[:] = env.gripper_grasp_width + env._gripper_command[:] = 0.0 + + assert not bool(torch.any(terminations.stable_pour_success(env, dwell_time_s=env.step_dt))) + assert not bool(torch.any(env.held_delivered_mask())) + + # Re-grasping and lifting the now-empty source cannot retroactively validate the first entry. + env.common_step_counter += 1 + env._cup[:, 2] = 0.06 + assert not bool(torch.any(terminations.stable_pour_success(env, dwell_time_s=env.step_dt))) + assert not bool(torch.any(env.current_held_delivered_mask())) + + +def test_success_reward_mirrors_cached_success_terminal(): + env = FakeEnv() + env._tgt[:] = 0.0 + env.termination_manager.terminated[:] = torch.tensor([True, True, False, False]) + env.termination_manager._terms["success"][:] = torch.tensor([True, False, False, False]) + + success = rewards.pour_success_bonus(env) * env.step_dt + + assert success.tolist() == [1.0, 0.0, 0.0, 0.0] + + +def test_success_reward_tolerates_recording_tools_removing_success_termination(): + env = FakeEnv() + env.termination_manager._terms.pop("success") + + success = rewards.pour_success_bonus(env) * env.step_dt + + assert success.tolist() == [0.0, 0.0, 0.0, 0.0] + + +def test_terminal_failure_tolerates_recording_tools_removing_success_termination(): + env = FakeEnv() + env.termination_manager._terms.pop("success") + env.termination_manager.terminated[:] = torch.tensor([True, False, True, False]) + + failure = rewards.terminal_failure(env) * env.step_dt + + assert failure.tolist() == [1.0, 0.0, 1.0, 0.0] + + +def test_terminal_failure_can_exclude_ordinary_fixed_horizon_timeout(): + env = FakeEnv() + env.termination_manager.terminated[:] = torch.tensor([True, False, True, False]) + env.termination_manager.time_outs[:] = torch.tensor([False, True, True, False]) + env.termination_manager._terms["success"][:] = torch.tensor([False, False, True, False]) + + failure = rewards.terminal_failure(env, include_time_out=False) * env.step_dt + + assert failure.tolist() == [1.0, 0.0, 0.0, 0.0] + + +def test_held_delivery_tracker_records_only_first_entries_and_is_step_idempotent(): + from isaaclab_tasks.contrib.franka_pour.pour_env import FrankaPourEnv + + env = FakeDeliveryEnv() + env._target_mask[:] = torch.tensor([[True, False, False, False], [True, False, False, False]]) + + FrankaPourEnv.update_held_delivery_tracker(env, torch.tensor([True, False])) + assert FrankaPourEnv.held_delivered_mask(env).tolist() == [ + [True, False, False, False], + [False, False, False, False], + ] + + # A second consumer in the same manager step cannot reinterpret a changed view as a new entry. + env._target_mask[:, 1] = True + FrankaPourEnv.update_held_delivery_tracker(env, torch.ones(2, dtype=torch.bool)) + assert FrankaPourEnv.held_delivered_mask(env).tolist() == [ + [True, False, False, False], + [False, False, False, False], + ] + + env.common_step_counter += 1 + FrankaPourEnv.update_held_delivery_tracker(env, torch.ones(2, dtype=torch.bool)) + assert FrankaPourEnv.held_delivered_mask(env).tolist() == [ + [True, True, False, False], + [False, True, False, False], + ] + + # An earlier unheld entry can qualify after the particle leaves and validly re-enters. + env._target_mask[1, 0] = False + env.common_step_counter += 1 + FrankaPourEnv.update_held_delivery_tracker(env, torch.ones(2, dtype=torch.bool)) + env._target_mask[1, 0] = True + env.common_step_counter += 1 + FrankaPourEnv.update_held_delivery_tracker(env, torch.ones(2, dtype=torch.bool)) + assert bool(FrankaPourEnv.held_delivered_mask(env)[1, 0]) + + env._target_mask[0, 0] = False + assert bool(FrankaPourEnv.held_delivered_mask(env)[0, 0]) + assert not bool(FrankaPourEnv.current_held_delivered_mask(env)[0, 0]) + + +def test_held_delivery_progress_is_signed_capped_and_resets_selectively(): + env = FakeDeliveryEnv() + _set_simple_env_held_state(env) + env.pour_target_frac[:] = torch.tensor([0.5, 0.25]) + env._target_mask[:] = torch.tensor([[True, True, True, False], [True, True, False, False]]) + term = rewards.HeldDeliveryProgress(SimpleNamespace(), env) + + # Credit stops at each environment's active success threshold. + torch.testing.assert_close(term(env) * env.step_dt, torch.tensor([0.5, 0.25])) + torch.testing.assert_close(term(env) * env.step_dt, torch.zeros(2)) + + env.common_step_counter += 1 + env._target_mask[:] = torch.tensor([[False, True, False, False], [True, True, True, True]]) + torch.testing.assert_close(term(env) * env.step_dt, torch.tensor([-0.25, 0.0])) + + env.common_step_counter += 1 + env._target_mask[:] = torch.tensor([[True, True, True, False], [False, False, False, False]]) + torch.testing.assert_close(term(env) * env.step_dt, torch.tensor([0.25, -0.25])) + + term.reset(torch.tensor([0])) + torch.testing.assert_close(term(env) * env.step_dt, torch.tensor([0.5, 0.0])) + + +def test_held_delivery_progress_claws_back_timeout_credit_and_preserves_success_credit(): + env = FakeDeliveryEnv() + _set_simple_env_held_state(env) + env._target_mask[:] = torch.tensor([[True, True, False, False], [True, True, False, False]]) + term = rewards.HeldDeliveryProgress(SimpleNamespace(), env) + + torch.testing.assert_close(term(env) * env.step_dt, torch.tensor([0.5, 0.5])) + + env.common_step_counter += 1 + env.termination_manager.time_outs[0] = True + env.termination_manager.terminated[1] = True + env.termination_manager._terms["success"][1] = True + + torch.testing.assert_close(term(env) * env.step_dt, torch.tensor([-0.5, 0.0])) + torch.testing.assert_close(term._previous_credit, torch.tensor([0.0, 0.5])) + + +def test_held_delivery_progress_does_not_pay_new_credit_on_failed_terminal_step(): + env = FakeDeliveryEnv() + _set_simple_env_held_state(env) + env._target_mask[:] = torch.tensor([[True, True, False, False], [True, True, False, False]]) + env.termination_manager.terminated[0] = True + env.termination_manager._terms["success"][1] = True + env.termination_manager.terminated[1] = True + term = rewards.HeldDeliveryProgress(SimpleNamespace(), env) + + torch.testing.assert_close(term(env) * env.step_dt, torch.tensor([0.0, 0.5])) + + +def test_new_delivery_never_rewards_unheld_first_entry_after_regrasp(): + env = FakeDeliveryEnv() + _set_simple_env_held_state(env, torch.tensor([True, False])) + env._target_mask[:] = torch.tensor([[True, True, False, False]] * 2) + term = rewards.NewlyDeliveredParticles(SimpleNamespace(), env) + + torch.testing.assert_close(term(env) * env.step_dt, torch.tensor([0.5, 0.0])) + # First entry is consumed even when invalid, so grabbing the source afterward cannot recover it. + env.action_manager.get_term("gripper_action").commanded_position[1] = 0.0 + env.common_step_counter += 1 + torch.testing.assert_close(term(env) * env.step_dt, torch.zeros(2)) + + +def test_new_spill_penalizes_each_particle_at_most_once_and_resets_selectively(): + env = SimpleNamespace( + num_envs=2, + num_particles=4, + device="cpu", + step_dt=0.02, + ) + spill_mask = torch.tensor([[True, True, False, False], [True, True, False, False]]) + env.particles_spilled_mask = lambda: spill_mask + term = rewards.NewlySpilledParticles(SimpleNamespace(), env) + + torch.testing.assert_close(term(env) * env.step_dt, torch.tensor([0.5, 0.5])) + torch.testing.assert_close(term(env) * env.step_dt, torch.zeros(2)) + + spill_mask = torch.tensor([[False, True, False, False], [False, True, False, False]]) + torch.testing.assert_close(term(env) * env.step_dt, torch.zeros(2)) + + spill_mask = torch.tensor([[True, True, False, False], [False, False, True, False]]) + torch.testing.assert_close(term(env) * env.step_dt, torch.tensor([0.0, 0.25])) + + term.reset(torch.tensor([0])) + torch.testing.assert_close(term(env) * env.step_dt, torch.tensor([0.5, 0.0])) + + +def test_spill_mask_requires_table_contact_outside_both_cups(): + points = torch.zeros((2, 4, 3)) + points[:, :, 2] = torch.tensor([0.003, 0.0031, 0.0, -0.01]) + in_source = torch.tensor([[False, False, True, False], [False, False, False, False]]) + in_target = torch.tensor([[False, False, False, True], [False, False, False, False]]) + + spilled = terminations._spilled_particle_mask(points, in_source, in_target, max_height=0.003) + + assert spilled.tolist() == [[True, False, False, False], [True, False, True, True]] + + +def test_delivered_particle_mask_excludes_particles_still_inside_source_cup(): + in_source = torch.tensor([[True, True, False, False], [True, False, False, True]]) + in_target = torch.tensor([[True, False, True, False], [True, True, False, True]]) + + delivered = terminations._delivered_particle_mask(in_source, in_target) + + assert delivered.tolist() == [[False, False, True, False], [False, True, False, False]] + assert not bool(torch.any(delivered & in_source)) + + +def test_particle_region_masks_use_exclusive_target_membership(): + from isaaclab_tasks.contrib.franka_pour.pour_env import FrankaPourEnv + + in_source = torch.tensor([[True, True, False, False]]) + in_target_region = torch.tensor([[True, False, True, False]]) + region_results = iter((in_source, in_target_region)) + points = torch.zeros((1, 4, 3)) + env = SimpleNamespace( + _particle_region_cache=None, + _particle_region_cache_step=-1, + common_step_counter=4, + particle_pos_e=lambda: points, + cup_pose_e=lambda: torch.empty((1, 7)), + target_pose_e=lambda: torch.empty((1, 7)), + _source_inner_lo_t=torch.empty(3), + _source_inner_hi_t=torch.empty(3), + _target_inner_lo_t=torch.empty(3), + _target_inner_hi_t=torch.empty(3), + _points_inside_cup=lambda *_: next(region_results), + cfg=SimpleNamespace(spill_table_height=0.0, particle_count_margin=0.003), + ) + + source, target, spilled = FrankaPourEnv._particle_region_masks(env) + + assert source.tolist() == [[True, True, False, False]] + assert target.tolist() == [[False, False, True, False]] + assert spilled.tolist() == [[False, False, False, True]] + + +def test_excessive_spill_is_strictly_greater_than_ten_percent(): + env = FakeEnv() + env.num_particles = 10 + env._spill = torch.tensor([0.0, 1.0, 2.0, 1.0]) + + excessive = terminations.excessive_spill(env) + + assert excessive.tolist() == [False, False, True, False] + + +def test_reset_dataset_spill_threshold_triggers_on_particle_74_of_245(): + env = FakeEnv() + env.num_particles = 245 + env.cfg.max_spill_fraction = 0.30 + env._spill = torch.tensor([0.0, 73.0, 74.0, 245.0]) + + excessive = terminations.excessive_spill(env) + + assert excessive.tolist() == [False, False, True, True] + + +def test_excessive_spill_can_monitor_without_terminating(): + env = FakeEnv() + env._spill = torch.tensor([0.0, 100.0, 101.0, 1000.0]) + + reported = terminations.excessive_spill(env, terminate=False) + + assert not bool(torch.any(reported)) + + +def test_stable_success_rejects_excessive_spill_without_spill_termination(): + env = FakeEnv() + env.pour_target_frac[:] = 0.9 + env._cup[:, 2] = 0.06 + env._tcp[:] = env._grasp + env._width[:] = env.gripper_grasp_width + env._gripper_command[:] = 0.0 + env._tgt[:] = 950.0 + env._spill[:] = torch.tensor([0.0, 100.0, 101.0, 1000.0]) + + success = terminations.stable_pour_success(env, dwell_time_s=env.step_dt) + + assert success.tolist() == [True, True, False, False] + + +def test_lost_grasp_requires_consecutive_loss_after_a_demonstrated_lift(): + env = FakeEnv() + env.cfg.success_min_lift_height = 0.05 + env._cup[:, 2] = 0.06 + env._tcp[:] = env._grasp + env._width[:] = env.gripper_grasp_width + env._gripper_command[:] = 0.0 + env._lifted_grasp_seen[:] = torch.tensor([True, True, False, False]) + env._tcp[1, 0] += 0.04 + + first = terminations.lost_lifted_grasp(env, dwell_time_s=3.0 * env.step_dt) + second = terminations.lost_lifted_grasp(env, dwell_time_s=3.0 * env.step_dt) + env._tcp[1] = env._grasp[1] + recovered = terminations.lost_lifted_grasp(env, dwell_time_s=3.0 * env.step_dt) + env._tcp[1, 0] += 0.04 + terminations.lost_lifted_grasp(env, dwell_time_s=3.0 * env.step_dt) + terminations.lost_lifted_grasp(env, dwell_time_s=3.0 * env.step_dt) + third = terminations.lost_lifted_grasp(env, dwell_time_s=3.0 * env.step_dt) + + assert first.tolist() == [False, False, False, False] + assert second.tolist() == [False, False, False, False] + assert recovered.tolist() == [False, False, False, False] + assert third.tolist() == [False, True, False, False] + + +def test_lost_grasp_can_monitor_without_terminating(): + env = FakeEnv() + env.cfg.success_min_lift_height = 0.05 + env._cup[:, 2] = 0.06 + env._tcp[:] = env._grasp + env._width[:] = env.gripper_grasp_width + env._gripper_command[:] = 0.0 + terminations.lost_lifted_grasp(env, dwell_time_s=env.step_dt, terminate=False) + + env._tcp[:, 0] += 0.1 + reported = terminations.lost_lifted_grasp(env, dwell_time_s=env.step_dt, terminate=False) + + assert not bool(torch.any(reported)) + assert env._lifted_grasp_seen.tolist() == [True, True, True, True] + assert env._lost_grasp_dwell_count.tolist() == [1, 1, 1, 1] + + +def test_cached_grasp_drop_terminates_and_produces_failure_pulse(): + env = FakeEnv() + env.cfg.success_min_lift_height = 0.05 + env._cup[:, 2] = 0.06 + env._width[:] = env.gripper_grasp_width + env._gripper_command[:] = 0.0 + # Model a grasping cache row followed by non-grasping rows. All cups are geometrically lost, + # but only the demonstrated grasp is eligible for dropped-cup termination. + env._lifted_grasp_seen[:] = torch.tensor([True, False, False, False]) + env._tcp[:] = env._grasp + torch.tensor((0.1, 0.0, 0.0)) + + dropped = terminations.lost_lifted_grasp(env, dwell_time_s=env.step_dt) + env.termination_manager.terminated[:] = dropped + failure_pulse = rewards.terminal_failure(env, include_time_out=False) * env.step_dt + + assert dropped.tolist() == [True, False, False, False] + torch.testing.assert_close(failure_pulse, torch.tensor([1.0, 0.0, 0.0, 0.0])) + + +def test_trajectory_phase_and_applied_reference_error_are_observable(): + arm_q = torch.arange(14, dtype=torch.float32).reshape(2, 7) + reference_target = arm_q + 0.25 + terms = { + "arm_action": SimpleNamespace( + reference_phase=torch.tensor([0.2, 0.7]), + reference_target=reference_target, + processed_actions=reference_target + 0.5, + reference_error=reference_target + 0.5 - arm_q, + ) + } + env = SimpleNamespace( + num_envs=2, + device="cpu", + action_manager=SimpleNamespace(get_term=lambda name: terms[name]), + arm_joint_pos=lambda: arm_q, + ) + + torch.testing.assert_close(observations.arm_reference_phase_obs(env), torch.tensor([[0.2], [0.7]])) + torch.testing.assert_close(observations.arm_reference_error_obs(env), torch.full((2, 7), 0.75)) + + +def test_trajectory_status_exposes_every_stateful_gate_and_dwell(): + arm_status = torch.tensor([[1.0, 0.0, 0.0, 0.0, 0.5, 0.25], [1.0, 1.0, 1.0, 0.0, 1.0, 1.0]]) + capture_status = torch.tensor([[0.0, 0.4], [1.0, 0.0]]) + terms = { + "arm_action": SimpleNamespace(milestone_status=arm_status), + "gripper_action": SimpleNamespace(capture_status=capture_status), + } + env = SimpleNamespace( + num_envs=2, + device="cpu", + action_manager=SimpleNamespace(get_term=lambda name: terms[name]), + ) + + torch.testing.assert_close( + observations.trajectory_status_obs(env), + torch.cat((arm_status, capture_status), dim=-1), + ) + + +def test_time_and_failure_dwell_observations_make_finite_horizon_state_observable(): + env = SimpleNamespace( + episode_length_buf=torch.tensor([0, 5, 10]), + max_episode_length=10, + step_dt=0.02, + _lost_grasp_dwell_count=torch.tensor([0, 1, 3]), + pour_target_frac=torch.tensor([0.1, 0.2, 0.35]), + cfg=SimpleNamespace(lost_grasp_dwell_time_s=0.05), + ) + + torch.testing.assert_close( + observations.time_remaining_obs(env), + torch.tensor([[1.0], [0.5], [0.0]]), + ) + torch.testing.assert_close( + observations.lost_grasp_dwell_obs(env), + torch.tensor([[0.0], [1.0 / 3.0], [1.0]]), + ) + torch.testing.assert_close( + observations.pour_target_fraction_obs(env), + torch.tensor([[0.1], [0.2], [0.35]]), + ) + + +def test_particle_transfer_observation_reports_airborne_flow_and_handles_empty_stream(): + source = torch.tensor([[True, True, False, False], [True, True, True, True]]) + target = torch.tensor([[False, False, False, True], [False, False, False, False]]) + spilled = torch.zeros_like(source) + positions = torch.zeros((2, 4, 3)) + positions[0, 2] = torch.tensor([0.8, -0.18, 0.0]) + velocities = torch.zeros_like(positions) + velocities[0, 2] = torch.tensor([2.0, 0.0, 0.0]) + target_pose = torch.tensor([[0.5, -0.18, 0.0, 0.0, 0.0, 0.0, 1.0]]).repeat(2, 1) + env = SimpleNamespace( + num_particles=4, + particle_region_masks=lambda: (source, target, spilled), + particle_pos_e=lambda: positions, + particle_vel_e=lambda: velocities, + target_pose_e=lambda: target_pose, + ) + + summary = observations.particle_transfer_obs(env) + + torch.testing.assert_close(summary[0], torch.tensor([1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.25])) + torch.testing.assert_close(summary[1], torch.zeros(7)) + + +def test_task_progress_is_signed_hold_neutral_and_cycle_neutral(): + env = FakeEnv() + env.curriculum_stage.zero_() + env._cup[:, :3] = torch.tensor([0.50, 0.00, 0.00]) + env._tcp[:] = env._grasp + torch.tensor([0.0, 0.0, 0.05]) + env._width[:] = env.gripper_open_width + params = { + "target_height": 0.12, + "reach_std": 0.07, + "grasp_reach_std": 0.015, + "grasp_preload_position": 0.025, + "lift_height": 0.06, + "align_std": 0.12, + "source_offset_xy": (0.0, 0.05), + "target_tilt": math.radians(150.0), + "pour_direction_xy": (0.0, -1.0), + "source_mouth_height": 0.036, + "alignment_radius": 0.15, + "active_through_stage": 1, + "min_lift_height": 0.05, + "max_tcp_distance": 0.015, + "max_gripper_width_error": 0.012, + "max_gripper_command": 0.025, + "discount_factor": 0.99, + } + term = rewards.PourTaskProgress(RewardTermCfg(func=rewards.PourTaskProgress, weight=5.0, params=params), env) + term.reset() + + env._tcp[0] = env._grasp[0] + torch.tensor([0.0, 0.0, 0.02]) + first = term(env, **params) * env.step_dt + env._tcp[0] = env._grasp[0] + torch.tensor([0.0, 0.0, 0.05]) + second = term(env, **params) * env.step_dt + + assert first[0] > first[1] + cycle_return = first[0] + params["discount_factor"] * second[0] + hold_return = first[1] + params["discount_factor"] * second[1] + torch.testing.assert_close(cycle_return, hold_return, atol=1.0e-6, rtol=0.0) + + +def test_task_progress_closes_the_potential_on_timeout(): + params = { + "target_height": 0.12, + "reach_std": 0.07, + "grasp_reach_std": 0.015, + "grasp_preload_position": 0.025, + "lift_height": 0.06, + "align_std": 0.12, + "source_offset_xy": (0.0, 0.05), + "target_tilt": math.radians(150.0), + "pour_direction_xy": (0.0, -1.0), + "source_mouth_height": 0.036, + "alignment_radius": 0.15, + "active_through_stage": 1, + "min_lift_height": 0.05, + "max_tcp_distance": 0.015, + "max_gripper_width_error": 0.012, + "max_gripper_command": 0.025, + "discount_factor": 0.99, + } + timeout_env = FakeEnv() + terminated_env = FakeEnv() + timeout_term = rewards.PourTaskProgress( + RewardTermCfg(func=rewards.PourTaskProgress, weight=5.0, params=params), timeout_env + ) + terminated_term = rewards.PourTaskProgress( + RewardTermCfg(func=rewards.PourTaskProgress, weight=5.0, params=params), terminated_env + ) + timeout_term.reset() + terminated_term.reset() + timeout_env.termination_manager.time_outs[:] = True + terminated_env.termination_manager.terminated[:] = True + + timeout_progress = timeout_term(timeout_env, **params) + terminated_progress = terminated_term(terminated_env, **params) + + torch.testing.assert_close(timeout_progress, terminated_progress) + + +def test_approach_progress_preserves_reach_gradient_after_premature_closure(): + env = FakeEnv() + env.curriculum_stage[:] = 3 + env._cup[:] = torch.tensor([0.50, 0.00, 0.00, 0.0, 0.0, 0.0, 1.0]) + env._grasp[:] = torch.tensor([0.50, 0.00, 0.032]) + env._tcp[:] = env._grasp + torch.tensor([0.20, 0.0, 0.0]) + env._width[:] = env.gripper_open_width + env._gripper_command[:] = 0.04 + params = { + "position_std": 0.20, + "orientation_std": 0.75, + "open_hand_fraction": 0.35, + "active_from_stage": 3, + "discount_factor": 0.99, + } + term = rewards.ApproachProgress( + RewardTermCfg(func=rewards.ApproachProgress, weight=8.0, params=params), + env, + ) + term.reset() + + # Closing at stand-off must lose the coordination bonus, but must not erase the independent + # Cartesian reach gradient that lets the policy recover from this mistake. + env._width[0] = env.gripper_grasp_width + env._gripper_command[0] = 0.02 + premature_close = term(env, **params) * env.step_dt + env._tcp[0] = env._grasp[0] + torch.tensor([0.10, 0.0, 0.0]) + recover_reach = term(env, **params) * env.step_dt + + assert premature_close[0] < 0.0 + assert recover_reach[0] > 0.0 + torch.testing.assert_close(premature_close[1:], recover_reach[1:]) + + +def test_approach_progress_rewards_pose_alignment_without_a_reference_trajectory(): + env = FakeEnv() + env.curriculum_stage[:] = 3 + env._cup[:] = torch.tensor([0.50, 0.00, 0.00, 0.0, 0.0, 0.0, 1.0]) + env._grasp[:] = torch.tensor([0.50, 0.00, 0.032]) + env._tcp[:] = env._grasp + torch.tensor([0.08, 0.0, 0.0]) + env._tcp_quat[0] = torch.tensor([0.0, 0.0, math.sin(0.5), math.cos(0.5)]) + params = { + "position_std": 0.20, + "orientation_std": 0.75, + "open_hand_fraction": 0.35, + "active_from_stage": 3, + "discount_factor": 0.99, + } + term = rewards.ApproachProgress( + RewardTermCfg(func=rewards.ApproachProgress, weight=8.0, params=params), + env, + ) + term.reset() + + env._tcp_quat[0] = torch.tensor([0.0, 0.0, 0.0, 1.0]) + align = term(env, **params) * env.step_dt + env._tcp[0] = env._grasp[0] + approach = term(env, **params) * env.step_dt + + assert align[0] > 0.0 + assert approach[0] > 0.0 + + +def test_grasp_lift_progress_requires_near_contact_then_rewards_lift(): + env = FakeEnv() + env.curriculum_stage[:] = 2 + env._cup[:] = torch.tensor([0.50, 0.00, 0.00, 0.0, 0.0, 0.0, 1.0]) + env._grasp[:] = torch.tensor([0.50, 0.00, 0.032]) + env._tcp[:] = env._grasp + torch.tensor([0.10, 0.0, 0.0]) + env._width[:] = env.gripper_open_width + env._gripper_command[:] = 0.04 + params = { + "target_height": 0.10, + "grasp_reach_std": 0.025, + "grasp_preload_position": 0.025, + "grasp_fraction": 0.40, + "active_from_stage": 2, + "discount_factor": 0.99, + } + term = rewards.GraspLiftProgress( + RewardTermCfg(func=rewards.GraspLiftProgress, weight=10.0, params=params), + env, + ) + term.reset() + + env._width[0] = env.gripper_grasp_width + env._gripper_command[0] = 0.025 + far_close = term(env, **params) * env.step_dt + env._tcp[0] = env._grasp[0] + contact = term(env, **params) * env.step_dt + env._cup[0, 2] = 0.10 + env._grasp[0, 2] += 0.10 + env._tcp[0, 2] += 0.10 + lift = term(env, **params) * env.step_dt + + assert far_close[0] <= 0.0 + assert contact[0] > 0.0 + assert lift[0] > 0.0 + + +def test_lift_progress_is_signed_bounded_and_cycle_neutral(): + env = FakeEnv() + env._width[:] = env.gripper_grasp_width + env._gripper_command[:] = 0.025 + env._tcp[:] = env._grasp + env._cup[:, :3] = torch.tensor([0.50, 0.00, 0.00]) + cfg = RewardTermCfg( + func=rewards.LiftProgress, + weight=5.0, + params={"target_height": 0.12, "reach_std": 0.10}, + ) + term = rewards.LiftProgress(cfg, env) + term.reset() + + torch.testing.assert_close(term(env) * env.step_dt, torch.zeros(env.num_envs)) + env._cup[0, 2] = 0.06 + forward = term(env) * env.step_dt + env._cup[0, 2] = 0.0 + reverse = term(env) * env.step_dt + + assert 0.0 < forward[0] <= 1.0 + assert -1.0 <= reverse[0] < 0.0 + torch.testing.assert_close(forward + reverse, torch.zeros(env.num_envs)) + + +def test_lift_progress_rewards_ordered_open_approach_grasp_and_lift(): + env = FakeEnv() + env._cup[:, :3] = torch.tensor([0.50, 0.00, 0.00]) + env._tcp[:] = env._grasp + torch.tensor([0.0, 0.0, 0.05]) + env._width[:] = env.gripper_open_width + cfg = RewardTermCfg( + func=rewards.LiftProgress, + weight=10.0, + params={ + "target_height": 0.12, + "reach_std": 0.07, + "grasp_reach_std": 0.015, + "approach_fraction": 0.2, + "grasp_fraction": 0.3, + }, + ) + term = rewards.LiftProgress(cfg, env) + term.reset() + + env._tcp[0] = env._grasp[0] + torch.tensor([0.0, 0.0, 0.025]) + approach = term(env) * env.step_dt + env._tcp[0] = env._grasp[0] + env._width[0] = env.gripper_grasp_width + env._gripper_command[0] = 0.025 + grasp = term(env) * env.step_dt + env._cup[0, 2] = 0.06 + lift = term(env) * env.step_dt + + env._cup[0, 2] = 0.0 + reverse_lift = term(env) * env.step_dt + env._width[0] = env.gripper_open_width + env._tcp[0] = env._grasp[0] + torch.tensor([0.0, 0.0, 0.025]) + reverse_grasp = term(env) * env.step_dt + env._tcp[0] = env._grasp[0] + torch.tensor([0.0, 0.0, 0.05]) + reverse_approach = term(env) * env.step_dt + + assert approach[0] > 0.0 + assert grasp[0] > 0.0 + assert lift[0] > 0.0 + torch.testing.assert_close( + approach + grasp + lift + reverse_lift + reverse_grasp + reverse_approach, + torch.zeros(env.num_envs), + atol=1.0e-6, + rtol=0.0, + ) + + +def test_lift_progress_penalizes_closing_empty_gripper_far_from_cup(): + env = FakeEnv() + env._cup[:, :3] = torch.tensor([0.50, 0.00, 0.00]) + env._tcp[:] = env._grasp + torch.tensor([0.0, 0.0, 0.08]) + env._width[:] = env.gripper_open_width + cfg = RewardTermCfg( + func=rewards.LiftProgress, + weight=10.0, + params={ + "target_height": 0.12, + "reach_std": 0.07, + "grasp_reach_std": 0.015, + "approach_fraction": 0.2, + "grasp_fraction": 0.3, + }, + ) + term = rewards.LiftProgress(cfg, env) + term.reset() + + # An empty hand can exactly match the nominal cup width without touching the cup. + env._width[0] = env.gripper_grasp_width + env._gripper_command[0] = 0.025 + premature_close = term(env) * env.step_dt + + assert premature_close[0] < 0.0 + + +def test_lift_progress_monotonically_rewards_near_contact_closure_and_rejects_overclose(): + env = FakeEnv() + env._cup[:, :3] = torch.tensor([0.50, 0.00, 0.00]) + env._tcp[:] = env._grasp + env._width[:] = env.gripper_open_width + cfg = RewardTermCfg( + func=rewards.LiftProgress, + weight=10.0, + params={ + "target_height": 0.12, + "reach_std": 0.07, + "grasp_reach_std": 0.015, + "approach_fraction": 0.2, + "grasp_fraction": 0.3, + }, + ) + term = rewards.LiftProgress(cfg, env) + term.reset() + + for width, command in ((0.075, 0.035), (0.070, 0.030), (0.065, 0.025), (env.gripper_grasp_width, 0.020)): + env._width[0] = width + env._gripper_command[0] = command + assert (term(env) * env.step_dt)[0] > 0.0 + + env._width[0] = 0.0 + assert (term(env) * env.step_dt)[0] < 0.0 + + +def test_lift_progress_does_not_treat_contact_compressed_open_fingers_as_a_grasp(): + env = FakeEnv() + env._cup[:, :3] = torch.tensor([0.50, 0.00, 0.00]) + env._tcp[:] = env._grasp + env._width[:] = env.gripper_grasp_width + env._gripper_command[:] = 0.04 + cfg = RewardTermCfg( + func=rewards.LiftProgress, + weight=10.0, + params={ + "target_height": 0.12, + "reach_std": 0.07, + "grasp_reach_std": 0.015, + "grasp_preload_position": 0.025, + "approach_fraction": 0.2, + "grasp_fraction": 0.3, + }, + ) + term = rewards.LiftProgress(cfg, env) + term.reset() + + env._gripper_command[0] = 0.025 + preload = term(env) * env.step_dt + + assert preload[0] > 0.0 + + +def test_alignment_progress_is_signed_bounded_and_release_repays_progress(): + env = FakeEnv() + env._cup[:, :3] = torch.tensor([0.50, 0.00, 0.00]) + env._tcp[:] = env._grasp + env._width[:] = env.gripper_grasp_width + env._gripper_command[:] = 0.025 + cfg = RewardTermCfg( + func=rewards.AlignProgress, + weight=5.0, + params={ + "lift_height": 0.06, + "std": 0.12, + "grasp_reach_std": 0.015, + "grasp_preload_position": 0.025, + }, + ) + term = rewards.AlignProgress(cfg, env) + term.reset() + + torch.testing.assert_close(term(env) * env.step_dt, torch.zeros(env.num_envs)) + env._cup[0, :3] = torch.tensor([0.50, -0.18, 0.12]) + env._grasp[0, :2] = torch.tensor([0.50, -0.18]) + env._tcp[0] = env._grasp[0] + forward = term(env) * env.step_dt + env._gripper_command[0] = 0.04 + release = term(env) * env.step_dt + + assert 0.0 < forward[0] <= 1.0 + assert -1.0 <= release[0] < 0.0 + torch.testing.assert_close(forward + release, torch.zeros(env.num_envs)) + + +def test_prerequisite_progress_is_gated_by_curriculum_stage(): + env = FakeEnv() + env.curriculum_stage[:] = torch.arange(env.num_envs) + env._cup[:, :3] = torch.tensor([0.50, 0.00, 0.00]) + env._tcp[:] = env._grasp + env._width[:] = env.gripper_grasp_width + env._gripper_command[:] = 0.025 + lift = rewards.LiftProgress(RewardTermCfg(func=rewards.LiftProgress, weight=5.0), env) + align = rewards.AlignProgress(RewardTermCfg(func=rewards.AlignProgress, weight=5.0), env) + lift.reset() + align.reset() + + env._cup[:, :3] = torch.tensor([0.50, -0.18, 0.12]) + env._grasp[:, :2] = torch.tensor([0.50, -0.18]) + env._tcp[:] = env._grasp + lift_progress = lift(env) * env.step_dt + align_progress = align(env) * env.step_dt + + torch.testing.assert_close(lift_progress[:2], torch.zeros(2)) + assert bool(torch.all(lift_progress[2:] > 0.0)) + assert align_progress[0] == 0.0 + assert bool(torch.all(align_progress[1:] > 0.0)) + + +def _prepare_tilt_progress_env() -> FakeEnv: + env = FakeEnv() + env.curriculum_stage[:] = torch.arange(env.num_envs) + env._cup[:, :3] = torch.tensor([0.50, -0.18, 0.06]) + env._cup[:, 3:7] = torch.tensor([0.0, 0.0, 0.0, 1.0]) + env._target[:, :3] = torch.tensor([0.50, -0.18, 0.00]) + env._tcp[:] = env._grasp + env._width[:] = env.gripper_grasp_width + env._gripper_command[:] = 0.0 + return env + + +def _tilt_progress_term(env: FakeEnv) -> rewards.PourTiltProgress: + cfg = RewardTermCfg( + func=rewards.PourTiltProgress, + weight=5.0, + params={ + "target_tilt": math.radians(150.0), + "pour_direction_xy": (0.0, -1.0), + "source_mouth_height": 0.0, + "alignment_radius": 0.10, + "active_through_stage": 1, + }, + ) + return rewards.PourTiltProgress(cfg, env) + + +def test_tilt_progress_is_directional_stage_gated_and_cycle_neutral(): + env = _prepare_tilt_progress_env() + term = _tilt_progress_term(env) + term.reset() + + target_tilt = math.radians(150.0) + partial_tilt = 0.55 + # The former 31.5-degree target must remain partial progress toward the physical drain angle. + env._cup[:, 3:7] = torch.tensor([math.sin(0.5 * partial_tilt), 0.0, 0.0, math.cos(0.5 * partial_tilt)]) + partial = term(env) * env.step_dt + env._cup[:, 3:7] = torch.tensor([math.sin(0.5 * target_tilt), 0.0, 0.0, math.cos(0.5 * target_tilt)]) + completion = term(env) * env.step_dt + env._cup[:, 3:7] = torch.tensor([0.0, 0.0, 0.0, 1.0]) + reverse = term(env) * env.step_dt + + partial_fraction = partial_tilt / target_tilt + torch.testing.assert_close( + partial, + torch.tensor([partial_fraction, partial_fraction, 0.0, 0.0]), + atol=1.0e-6, + rtol=0.0, + ) + torch.testing.assert_close( + partial + completion, + torch.tensor([1.0, 1.0, 0.0, 0.0]), + atol=1.0e-6, + rtol=0.0, + ) + torch.testing.assert_close(partial + completion + reverse, torch.zeros(env.num_envs), atol=1.0e-6, rtol=0.0) + + +def test_tilt_progress_rejects_wrong_direction_unaligned_and_unheld_motion(): + env = _prepare_tilt_progress_env() + env.curriculum_stage.zero_() + term = _tilt_progress_term(env) + term.reset() + + half_angle = 0.5 * 0.55 + env._cup[0, 3:7] = torch.tensor([-math.sin(half_angle), 0.0, 0.0, math.cos(half_angle)]) + env._cup[1, 3:7] = torch.tensor([0.0, math.sin(half_angle), 0.0, math.cos(half_angle)]) + env._cup[2, 3:7] = torch.tensor([math.sin(half_angle), 0.0, 0.0, math.cos(half_angle)]) + env._cup[2, 0] += 0.10 + env._cup[3, 3:7] = torch.tensor([math.sin(half_angle), 0.0, 0.0, math.cos(half_angle)]) + env._gripper_command[3] = 0.04 + + torch.testing.assert_close(term(env) * env.step_dt, torch.zeros(env.num_envs), atol=1.0e-6, rtol=0.0) + + +def test_tilt_progress_selective_reset_baselines_only_selected_worlds_and_release_repays(): + env = _prepare_tilt_progress_env() + env.curriculum_stage.zero_() + term = _tilt_progress_term(env) + term.reset() + + half_angle = 0.5 * 0.55 + env._cup[:2, 3:7] = torch.tensor([math.sin(half_angle), 0.0, 0.0, math.cos(half_angle)]) + term.reset(torch.tensor([0])) + progress = term(env) * env.step_dt + env._gripper_command[1] = 0.04 + release = term(env) * env.step_dt + + assert progress[0] == 0.0 + assert progress[1] > 0.0 + torch.testing.assert_close(progress[1] + release[1], torch.tensor(0.0), atol=1.0e-6, rtol=0.0) + torch.testing.assert_close(progress[2:], torch.zeros(2)) + + +def test_pour_reference_progress_tracks_validated_path_and_is_stage_gated(): + env = _prepare_tilt_progress_env() + env.curriculum_stage[:] = torch.tensor([0, 0, 0, 1]) + start_q = (0.0,) * 7 + target_q = (1.0, -1.0, 0.5, -0.5, 0.25, -0.25, 0.75) + cfg = RewardTermCfg( + func=rewards.PourReferenceProgress, + weight=10.0, + params={"start_q": start_q, "target_q": target_q, "active_stage": 0}, + ) + term = rewards.PourReferenceProgress(cfg, env) + term.reset() + + env._arm_q[:] = 0.5 * torch.tensor(target_q) + halfway = term(env, start_q=start_q, target_q=target_q) * env.step_dt + env._arm_q[:] = torch.tensor(target_q) + completion = term(env, start_q=start_q, target_q=target_q) * env.step_dt + env._arm_q.zero_() + reverse = term(env, start_q=start_q, target_q=target_q) * env.step_dt + + torch.testing.assert_close(halfway[:3], torch.full((3,), 0.5), atol=1.0e-6, rtol=0.0) + torch.testing.assert_close(completion[:3], torch.full((3,), 0.5), atol=1.0e-6, rtol=0.0) + torch.testing.assert_close((halfway + completion + reverse)[:3], torch.zeros(3), atol=1.0e-6, rtol=0.0) + torch.testing.assert_close(halfway[3:], torch.zeros(1)) + + +def test_progress_reset_baselines_selected_worlds_without_cross_world_history(): + env = FakeEnv() + env._width[:] = env.gripper_grasp_width + env._gripper_command[:] = 0.025 + env._tcp[:] = env._grasp + env._cup[:, :3] = torch.tensor([0.50, 0.00, 0.00]) + cfg = RewardTermCfg( + func=rewards.LiftProgress, + weight=5.0, + params={"target_height": 0.12, "reach_std": 0.10}, + ) + term = rewards.LiftProgress(cfg, env) + term.reset() + + env._cup[:2, 2] = 0.06 + term.reset(torch.tensor([0])) + progress = term(env) * env.step_dt + + assert progress[0] == 0.0 + assert progress[1] > 0.0 + torch.testing.assert_close(progress[2:], torch.zeros(2)) + + +def test_state_finite_rejects_raw_nonfinite_cup_and_robot_state(): + robot_joint_pos = torch.zeros((6, 7)) + robot_joint_vel = torch.zeros((6, 7)) + tcp_body_q = torch.tensor([[0.5, 0.0, 0.1, 0.0, 0.0, 0.0, 1.0]]).repeat(6, 1) + cup_body_q = tcp_body_q.clone() + cup_lin_vel = torch.zeros((6, 3)) + cup_ang_vel = torch.zeros((6, 3)) + particle_pos = torch.zeros((6, 16, 3)) + robot_joint_pos[1, 0] = float("nan") + robot_joint_vel[2, 0] = float("inf") + tcp_body_q[3, 0] = float("nan") + cup_lin_vel[4, 0] = float("inf") + particle_pos[5, 0, 0] = float("nan") + + finite = terminations._state_finite( + robot_joint_pos, + robot_joint_vel, + tcp_body_q, + cup_body_q, + cup_lin_vel, + cup_ang_vel, + particle_pos, + ) + + assert finite.tolist() == [True, False, False, False, False, False] + + +def test_rigid_state_bounds_reject_each_extreme_finite_observation_source(): + count = 8 + robot_joint_pos = torch.zeros((count, 9)) + robot_joint_vel = torch.zeros((count, 9)) + joint_pos_limits = torch.tensor([[[-1.0, 1.0]]]).repeat(count, 9, 1) + tcp_body_q = torch.tensor([[0.5, 0.0, 0.1, 0.0, 0.0, 0.0, 1.0]]).repeat(count, 1) + cup_body_q = tcp_body_q.clone() + cup_lin_vel = torch.zeros((count, 3)) + cup_ang_vel = torch.zeros((count, 3)) + env_origins = torch.zeros((count, 3)) + + robot_joint_pos[1, 0] = 1.051 + robot_joint_vel[2, 0] = 20.01 + tcp_body_q[3, 0] = 1.501 + cup_body_q[4, 2] = -0.501 + cup_lin_vel[5, 0] = 10.01 + cup_ang_vel[6, 0] = 50.01 + cup_body_q[7, 3:7] = 0.0 + + in_bounds = terminations._rigid_state_in_bounds( + robot_joint_pos, + robot_joint_vel, + joint_pos_limits, + tcp_body_q, + cup_body_q, + cup_lin_vel, + cup_ang_vel, + env_origins, + lower_bound=(-0.5, -1.0, -0.5), + upper_bound=(1.5, 1.0, 1.5), + joint_position_margin=0.05, + max_joint_velocity=20.0, + max_cup_linear_velocity=10.0, + max_cup_angular_velocity=50.0, + ) + + assert in_bounds.tolist() == [True, False, False, False, False, False, False, False] + + +def test_pose_conversion_returns_identity_for_nonfinite_or_degenerate_quaternions(): + from isaaclab_tasks.contrib.franka_pour.pour_env import FrankaPourEnv + + env = SimpleNamespace(env_origins=torch.zeros((4, 3))) + pose_w = torch.tensor( + [ + [0.5, 0.0, 0.1, 0.0, 0.0, 0.0, 2.0], + [0.5, 0.0, 0.1, float("nan"), 0.0, 0.0, 1.0], + [0.5, 0.0, 0.1, 0.0, float("inf"), 0.0, 1.0], + [0.5, 0.0, 0.1, 0.0, 0.0, 0.0, 0.0], + ] + ) + + pose_e = FrankaPourEnv._pose_w_to_e(env, pose_w) + + assert torch.isfinite(pose_e).all() + torch.testing.assert_close(pose_e[0, 3:7], torch.tensor([0.0, 0.0, 0.0, 1.0])) + torch.testing.assert_close( + pose_e[1:, 3:7], + torch.tensor([[0.0, 0.0, 0.0, 1.0]]).repeat(3, 1), + ) + + +def test_gripper_width_uses_open_width_for_nonfinite_joint_positions(): + from isaaclab_tasks.contrib.franka_pour.pour_env import FrankaPourEnv + + joint_pos = torch.tensor( + [ + [0.02, 0.03], + [float("nan"), 0.03], + [0.02, float("inf")], + ] + ) + env = SimpleNamespace( + _robot=SimpleNamespace(data=SimpleNamespace(joint_pos=SimpleNamespace(torch=joint_pos))), + _finger_joint_ids=[0, 1], + gripper_open_width=0.08, + ) + + width = FrankaPourEnv.gripper_width(env) + + assert torch.isfinite(width).all() + torch.testing.assert_close(width, torch.tensor([0.05, 0.08, 0.08])) + + +def test_particle_workspace_rejects_finite_outliers_per_environment(): + particle_pos_e = torch.zeros((3, 4, 3)) + particle_pos_e[0] = torch.tensor([0.5, 0.0, 0.2]) + particle_pos_e[1, 0] = torch.tensor([1.51, 0.0, 0.2]) + particle_pos_e[2, 0] = torch.tensor([0.5, 0.0, -0.51]) + + inside = terminations._particles_in_workspace( + particle_pos_e, + lower_bound=(-0.5, -1.0, -0.5), + upper_bound=(1.5, 1.0, 1.5), + ) + + assert inside.tolist() == [True, False, False] diff --git a/source/isaaclab_tasks/test/contrib/test_franka_pour_media_fill.py b/source/isaaclab_tasks/test/contrib/test_franka_pour_media_fill.py new file mode 100644 index 000000000000..866c6ca3a434 --- /dev/null +++ b/source/isaaclab_tasks/test/contrib/test_franka_pour_media_fill.py @@ -0,0 +1,93 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Unit tests for the Franka pour granular-media fill (no simulator).""" + +from types import SimpleNamespace + +import numpy as np + +from isaaclab_tasks.contrib.franka_pour.cube_bowl_mesh import cube_bowl_inner_bounds +from isaaclab_tasks.contrib.franka_pour.cup_media import cup_cavity_lattice, particle_mass_and_radius +from isaaclab_tasks.contrib.franka_pour.media_fill import cube_fill_points, expected_fill_count +from isaaclab_tasks.contrib.franka_pour.pour_env_cfg import FrankaPourEnvCfg + +LO, HI = cube_bowl_inner_bounds(0.037, 0.037, 0.045, 0.009) +CLR = max(0.003, 3 * 0.002) + + +def _min_neighbor_distance(pts: np.ndarray, sample: int = 200) -> float: + """Smallest nearest-neighbour distance over a subsample (numpy-only, no scipy).""" + idx = np.linspace(0, len(pts) - 1, min(sample, len(pts))).astype(int) + sub = pts[idx] + best = np.inf + for i in range(len(sub)): + d = np.linalg.norm(pts - sub[i], axis=1) + d[np.argmin(d)] = np.inf # drop the self-distance (0) + best = min(best, float(d.min())) + return best + + +def test_points_nonempty_and_inside_cavity(): + pts = cube_fill_points(LO, HI, spacing=0.003, fill_frac=1.0, jitter=0.0) + assert pts.dtype == np.float32 and pts.shape[1] == 3 and len(pts) > 200 + assert np.all(pts[:, 0] >= LO[0] + CLR - 1e-6) and np.all(pts[:, 0] <= HI[0] - CLR + 1e-6) + assert np.all(pts[:, 1] >= LO[1] + CLR - 1e-6) and np.all(pts[:, 1] <= HI[1] - CLR + 1e-6) + assert np.all(pts[:, 2] >= LO[2] + CLR - 1e-6) + + +def test_fill_frac_limits_height(): + full = cube_fill_points(LO, HI, spacing=0.003, fill_frac=1.0, jitter=0.0) + half = cube_fill_points(LO, HI, spacing=0.003, fill_frac=0.5, jitter=0.0) + assert float(half[:, 2].max()) < float(full[:, 2].max()) + assert len(half) < len(full) + + +def test_deterministic_seed(): + a = cube_fill_points(LO, HI, spacing=0.003, seed=7) + b = cube_fill_points(LO, HI, spacing=0.003, seed=7) + assert np.array_equal(a, b) + + +def test_no_overlap_min_spacing(): + pts = cube_fill_points(LO, HI, spacing=0.003, jitter=0.0) + assert _min_neighbor_distance(pts) > 0.5 * 0.003 + + +def test_expected_count_matches_actual(): + n = expected_fill_count(LO, HI, spacing=0.003, fill_frac=1.0) + pts = cube_fill_points(LO, HI, spacing=0.003, fill_frac=1.0, jitter=0.0) + assert n == len(pts) + + +def test_particle_mass_and_radius_represent_one_full_lattice_cell(): + """Implicit MPM derives particle volume as 8*r^3, so r must be half the lattice spacing.""" + cfg = SimpleNamespace( + voxel_size=0.006, + particles_per_cell=2.0, + media_material=SimpleNamespace(density=1500.0), + ) + mass, radius = particle_mass_and_radius(cfg) + spacing = cfg.voxel_size / cfg.particles_per_cell + represented_volume = 8.0 * radius**3 + + assert np.isclose(radius, 0.5 * spacing) + assert np.isclose(represented_volume, spacing**3) + assert np.isclose(mass / represented_volume, cfg.media_material.density) + + +def test_task_fill_fraction_means_represented_cavity_volume_not_inset_point_height(): + cfg = FrankaPourEnvCfg() + # Fill-fraction fidelity needs multiple lattice layers. The task's intentionally coarse + # rollout resolution is covered separately by the environment-config regression tests. + cfg.voxel_size = 0.006 + points, cell = cup_cavity_lattice(cfg) + cavity_volume = cfg.source_cup_inner_width * cfg.source_cup_inner_depth * cfg.source_cup_cavity_depth + represented_fill = len(points) * float(np.prod(cell)) / cavity_volume + points_per_layer = len(np.unique(np.round(points[:, :2], decimals=5), axis=0)) + one_layer_fraction = points_per_layer * float(np.prod(cell)) / cavity_volume + + # The safe wall/rim inset can leave the nearest higher layer infeasible in a shallow cup. + assert abs(represented_fill - cfg.media_fill_frac) <= one_layer_fraction diff --git a/source/isaaclab_tasks/test/contrib/test_franka_pour_reset_bridge.py b/source/isaaclab_tasks/test/contrib/test_franka_pour_reset_bridge.py new file mode 100644 index 000000000000..c9a9c0b6bd5f --- /dev/null +++ b/source/isaaclab_tasks/test/contrib/test_franka_pour_reset_bridge.py @@ -0,0 +1,151 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Tests for restoring Franka Pour reset-dataset rows into simulation assets.""" + +from types import SimpleNamespace + +import torch + +import isaaclab_tasks.contrib.franka_pour.pour_env as pour_env_module +from isaaclab_tasks.contrib.franka_pour.pour_env import FrankaPourEnv + + +class _RobotRecorder: + def __init__(self): + self.position_writes = [] + self.velocity_writes = [] + self.position_targets = [] + self.data = SimpleNamespace(body_link_pose_w=torch.empty(0)) + + def write_joint_position_to_sim_index(self, *, position, **_kwargs): + self.position_writes.append(position.clone()) + + def write_joint_velocity_to_sim_index(self, *, velocity, **_kwargs): + self.velocity_writes.append(velocity.clone()) + + def set_joint_position_target_index(self, *, target, **_kwargs): + self.position_targets.append(target.clone()) + + +class _RigidRecorder: + def __init__(self): + self.root_pose = None + self.root_velocity = None + + def write_root_pose_to_sim_index(self, *, root_pose, **_kwargs): + self.root_pose = root_pose.clone() + + def write_root_velocity_to_sim_index(self, *, root_velocity, **_kwargs): + self.root_velocity = root_velocity.clone() + + +class _MediaRecorder: + def __init__(self): + self.position = None + self.velocity = None + + def write_particle_pos_to_sim_index(self, position, **_kwargs): + self.position = position.clone() + + def write_particle_velocity_to_sim_index(self, velocity, **_kwargs): + self.velocity = velocity.clone() + + +def test_reset_dataset_restores_exact_state_and_clears_solver_history(monkeypatch): + robot = _RobotRecorder() + source_cup = _RigidRecorder() + target_cup = _RigidRecorder() + media = _MediaRecorder() + gripper_reset = [] + gripper_action = SimpleNamespace( + set_reset_position=lambda position, **_kwargs: gripper_reset.append(position.clone()) + ) + rows = torch.tensor((1, 0)) + identity = torch.tensor((0.0, 0.0, 0.0, 1.0)) + states = { + "category": torch.tensor((0, 1), dtype=torch.int8), + "arm_joint_position": torch.arange(14, dtype=torch.float32).reshape(2, 7), + "arm_joint_velocity": torch.arange(14, dtype=torch.float32).reshape(2, 7) * 0.01, + "finger_joint_position": torch.tensor(((0.01, 0.01), (0.028, 0.028))), + "finger_joint_velocity": torch.tensor(((0.1, 0.1), (0.2, 0.2))), + "finger_joint_target": torch.tensor(((0.012, 0.012), (0.024, 0.024))), + "source_root_pose": torch.stack( + ( + torch.cat((torch.tensor((0.4, 0.1, 0.0)), identity)), + torch.cat((torch.tensor((0.6, -0.2, 0.3)), identity)), + ) + ), + "source_root_velocity": torch.arange(12, dtype=torch.float32).reshape(2, 6) * 0.01, + "target_root_pose": torch.stack( + ( + torch.cat((torch.tensor((0.5, -0.2, 0.0)), identity)), + torch.cat((torch.tensor((0.7, 0.2, 0.0)), identity)), + ) + ), + "target_root_velocity": torch.arange(12, dtype=torch.float32).reshape(2, 6) * 0.02, + "particle_layout_id": torch.zeros(2, dtype=torch.int32), + } + solver_resets = [] + env = SimpleNamespace( + num_envs=2, + device="cpu", + reset_dataset_row_id=rows, + _reset_dataset_states=states, + _reset_dataset_particle_local_position=torch.tensor((((0.0, 0.0, 0.01), (0.01, 0.0, 0.02)),)), + _reset_dataset_particle_local_velocity=torch.zeros((1, 2, 3)), + _robot=robot, + _arm_joint_ids=torch.arange(7), + _finger_joint_ids=torch.tensor((7, 8)), + action_manager=SimpleNamespace( + get_term=lambda name: gripper_action if name == "gripper_action" else SimpleNamespace() + ), + env_origins=torch.tensor(((10.0, 0.0, 0.0), (20.0, 0.0, 0.0))), + _source_cup=source_cup, + _target_cup=target_cup, + _media=media, + _last_source_bank_index=torch.zeros(2, dtype=torch.long), + _last_arm_bank_index=torch.zeros(2, dtype=torch.long), + _last_target_bank_index=torch.zeros(2, dtype=torch.long), + _particle_region_cache=object(), + _particle_region_cache_step=1, + episode_succeeded=torch.ones(2, dtype=torch.bool), + ep_max_target_frac=torch.ones(2), + _success_dwell_count=torch.ones(2, dtype=torch.long), + _lost_grasp_dwell_count=torch.ones(2, dtype=torch.long), + _lifted_grasp_seen=torch.ones(2, dtype=torch.bool), + _target_entry_seen=torch.ones((2, 2), dtype=torch.bool), + _held_delivered=torch.ones((2, 2), dtype=torch.bool), + _held_delivery_tracker_step=1, + ) + monkeypatch.setattr( + pour_env_module.NewtonManager, + "reset_solver_state", + lambda **kwargs: solver_resets.append(kwargs), + ) + + FrankaPourEnv._reset_from_dataset(env, torch.arange(2), torch.ones(2, dtype=torch.bool)) + + torch.testing.assert_close(robot.position_writes[0], states["arm_joint_position"][rows]) + torch.testing.assert_close(robot.velocity_writes[0], states["arm_joint_velocity"][rows]) + torch.testing.assert_close(robot.position_writes[1], states["finger_joint_position"][rows]) + torch.testing.assert_close(robot.position_targets[1], states["finger_joint_target"][rows]) + torch.testing.assert_close(gripper_reset[0], states["finger_joint_target"][rows, :1]) + torch.testing.assert_close( + source_cup.root_pose[:, :3], + states["source_root_pose"][rows, :3] + env.env_origins, + ) + torch.testing.assert_close(target_cup.root_velocity, states["target_root_velocity"][rows]) + expected_particles = env._reset_dataset_particle_local_position.expand(2, -1, -1).clone() + expected_particles += source_cup.root_pose[:, None, :3] + torch.testing.assert_close(media.position, expected_particles) + torch.testing.assert_close(media.velocity, torch.zeros_like(media.velocity)) + assert len(solver_resets) == 1 + assert ( + solver_resets[0]["flags"] == pour_env_module.newton.StateFlags.BODY | pour_env_module.newton.StateFlags.PARTICLE + ) + assert env._particle_region_cache is None + assert not bool(env.episode_succeeded.any()) + assert env._lifted_grasp_seen.tolist() == [True, False] diff --git a/source/isaaclab_tasks/test/contrib/test_franka_pour_reset_dataset.py b/source/isaaclab_tasks/test/contrib/test_franka_pour_reset_dataset.py new file mode 100644 index 000000000000..83f6b6ad200e --- /dev/null +++ b/source/isaaclab_tasks/test/contrib/test_franka_pour_reset_dataset.py @@ -0,0 +1,175 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Unit tests for the Franka Pour reset-dataset curriculum adapter.""" + +from types import SimpleNamespace + +import pytest +import torch + +from isaaclab.managers import CurriculumTermCfg + +from isaaclab_tasks.contrib.franka_pour.mdp.reset_dataset import ( + PourResetDatasetCurriculum, + reset_dataset_difficulty, +) +from isaaclab_tasks.utils.adaptive_reset_sampler import AdaptiveResetSamplerCfg + + +def _task_contract() -> dict[str, object]: + return { + "arm_home": (0.0, 0.0), + "arm_joint_limits": torch.tensor(((-2.0, 2.0), (-2.0, 2.0))), + "source_region_center": (0.0, 0.0, 0.1), + "target_center_xy": (0.0, -0.2), + "tabletop_support_lower_xy": (-1.0, -1.0), + "tabletop_support_upper_xy": (1.0, 1.0), + "gripper_position_range": (0.0, 0.04), + } + + +def _states() -> dict[str, torch.Tensor]: + # Rows zero and one are grasping; rows two through five are increasingly hard reaching states. + return { + "category": torch.tensor((1, 1, 0, 0, 0, 0), dtype=torch.int8), + "objective": torch.tensor((1.0, 0.0, -1.0, -1.0, -1.0, -1.0)), + "arm_joint_position": torch.tensor(((0.0, 0.0), (0.0, 0.0), (0.0, 0.0), (0.4, 0.4), (1.0, 1.0), (1.8, 1.8))), + "source_root_pose": torch.tensor( + ( + (0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0), + (0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0), + (0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0), + (0.1, 0.1, 0.0, 0.0, 0.0, 0.0, 1.0), + (0.4, 0.4, 0.0, 0.0, 0.0, 0.0, 1.0), + (0.8, 0.8, 0.0, 0.0, 0.0, 0.0, 1.0), + ) + ), + "target_root_pose": torch.tensor( + ( + (0.0, -0.2, 0.0, 0.0, 0.0, 0.0, 1.0), + (0.0, -0.2, 0.0, 0.0, 0.0, 0.0, 1.0), + (0.0, -0.2, 0.0, 0.0, 0.0, 0.0, 1.0), + (0.1, -0.1, 0.0, 0.0, 0.0, 0.0, 1.0), + (0.4, 0.2, 0.0, 0.0, 0.0, 0.0, 1.0), + (0.8, 0.6, 0.0, 0.0, 0.0, 0.0, 1.0), + ) + ), + "finger_joint_position": torch.tensor( + ((0.028, 0.028), (0.028, 0.028), (0.04, 0.04), (0.03, 0.03), (0.02, 0.02), (0.0, 0.0)) + ), + } + + +class FakeResetDatasetEnv: + """Small manager-compatible environment stub.""" + + def __init__(self, *, freeze: bool = False, top_grasp_count: int | None = None): + self.num_envs = 4 + self.device = "cpu" + self._uses_reset_dataset = True + self._reset_dataset_states = _states() + self._reset_dataset_metadata = {"task_contract": _task_contract()} + self.cfg = SimpleNamespace( + reset_dataset_sampler=AdaptiveResetSamplerCfg( + target_success_rate=0.5, + temperature=0.1, + history_capacity=8, + prior_strength=2.0, + initial_frontier_size=2, + probe_size=2, + probe_fraction=0.1, + replay_fraction=0.1, + frontier_evidence=1.0, + ), + reset_dataset_top_grasp_count=top_grasp_count, + curriculum_freeze=freeze, + pour_target_frac=0.3, + ) + self.reset_dataset_row_id = torch.full((self.num_envs,), -1, dtype=torch.long) + self._forced_reset_dataset_row = torch.full_like(self.reset_dataset_row_id, -1) + self.episode_length_buf = torch.zeros(self.num_envs, dtype=torch.long) + self.episode_succeeded = torch.zeros(self.num_envs, dtype=torch.bool) + self.pour_target_frac = torch.zeros(self.num_envs) + + +def _term(env: FakeResetDatasetEnv) -> PourResetDatasetCurriculum: + return PourResetDatasetCurriculum(CurriculumTermCfg(func=PourResetDatasetCurriculum), env) + + +def test_reset_dataset_difficulty_orders_objectives_then_grades_non_grasps(): + difficulty = reset_dataset_difficulty(_states(), _task_contract()) + + assert difficulty.shape == (6,) + assert difficulty.tolist() == sorted(difficulty.tolist()) + assert difficulty[0] == 0.0 + assert difficulty[1] == pytest.approx(0.5) + assert bool(((difficulty >= 0.0) & (difficulty <= 1.0)).all()) + + +def test_reset_dataset_curriculum_honors_exact_raw_row_overrides(): + env = FakeResetDatasetEnv() + env._forced_reset_dataset_row[:] = torch.tensor((5, -1, 3, -1)) + term = _term(env) + + metrics = term(env, slice(None)) + + assert env.reset_dataset_row_id[[0, 2]].tolist() == [5, 3] + assert bool(torch.isin(env.reset_dataset_row_id, torch.arange(6)).all()) + assert env.pour_target_frac.tolist() == pytest.approx([0.3] * 4) + assert set(metrics) == { + "predicted_success_rate", + "observed_success_rate", + "dataset_success_rate", + "dataset_ever_solved_fraction", + "frontier_fraction", + "effective_pool_size", + } + + +def test_reset_dataset_curriculum_records_completed_rows_and_reports_compact_progress(): + env = FakeResetDatasetEnv() + term = _term(env) + env._forced_reset_dataset_row[:] = torch.arange(4) + term(env, slice(None)) + completed_rows = env.reset_dataset_row_id.clone() + env.episode_length_buf[:] = 1 + env.episode_succeeded[:] = torch.tensor((True, False, True, False)) + + metrics = term(env, slice(None)) + + assert 0.0 <= metrics["predicted_success_rate"] <= 1.0 + assert 0.0 <= metrics["observed_success_rate"] <= 1.0 + assert metrics["dataset_success_rate"] > 0.0 + assert metrics["dataset_ever_solved_fraction"] > 0.0 + assert metrics["dataset_success_rate"] <= metrics["dataset_ever_solved_fraction"] + assert completed_rows.min() >= 0 + + +def test_frozen_reset_dataset_samples_only_top_grasps_but_allows_exact_diagnostics(monkeypatch): + env = FakeResetDatasetEnv(freeze=True, top_grasp_count=1) + term = _term(env) + monkeypatch.setattr(torch, "randint", lambda *_args, **_kwargs: torch.zeros(4, dtype=torch.long)) + + metrics = term(env, slice(None)) + assert env.reset_dataset_row_id.tolist() == [0, 0, 0, 0] + assert metrics == {"frozen_pool_fraction": pytest.approx(1.0 / 6.0), "frozen_pool_size": 1.0} + + env._forced_reset_dataset_row[:] = torch.tensor((5, -1, -1, -1)) + env.episode_length_buf[:] = 1 + env.episode_succeeded[:] = True + state_before = term._sampler.state_dict() + term(env, slice(None)) + assert env.reset_dataset_row_id.tolist() == [5, 0, 0, 0] + assert all(torch.equal(value, term._sampler.state_dict()[name]) for name, value in state_before.items()) + + +def test_reset_dataset_curriculum_rejects_unknown_forced_row(): + env = FakeResetDatasetEnv() + env._forced_reset_dataset_row[0] = 99 + term = _term(env) + + with pytest.raises(ValueError, match="Unknown raw reset-row IDs"): + term(env, slice(None)) diff --git a/source/isaaclab_tasks/test/contrib/test_franka_pour_reset_dataset_generator.py b/source/isaaclab_tasks/test/contrib/test_franka_pour_reset_dataset_generator.py new file mode 100644 index 000000000000..ff3fb9d53860 --- /dev/null +++ b/source/isaaclab_tasks/test/contrib/test_franka_pour_reset_dataset_generator.py @@ -0,0 +1,534 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Focused CPU tests for the Franka Pour reset-dataset cache.""" + +import math +from copy import deepcopy +from types import SimpleNamespace + +import gymnasium as gym +import pytest +import torch + +from isaaclab_tasks.contrib.franka_pour.reset_dataset_generator import ( + FRANKA_POUR_RESET_DATASET_TASK_ID, + FrankaPourResetDatasetGenerator, + FrankaPourResetDatasetGeneratorCfg, + above_target_tilted_mask, + build_reset_dataset_payload, + grasp_objective_components, + normalize_grasp_objectives, + oriented_box_supported_by_bounds, + oriented_boxes_overlap, + reset_dataset_content_sha256, + select_production_reset_rows, + source_root_position_from_tcp_grasp, + validate_production_reset_dataset, + validate_reset_dataset, +) + + +def _identity_poses(positions: torch.Tensor) -> torch.Tensor: + poses = torch.zeros((positions.shape[0], 7), dtype=positions.dtype) + poses[:, :3] = positions + poses[:, 6] = 1.0 + return poses + + +def _tiny_states() -> dict[str, torch.Tensor]: + """Build two grasping and two non-grasping states with valid cache semantics.""" + count = 4 + source_pose = _identity_poses( + torch.tensor( + ( + (0.50, 0.00, 0.10), + (0.55, 0.05, 0.20), + (0.48, -0.12, 0.00), + (0.62, 0.10, 0.00), + ), + dtype=torch.float32, + ) + ) + target_pose = _identity_poses( + torch.tensor( + ( + (0.62, 0.00, 0.00), + (0.56, 0.04, 0.00), + (0.65, 0.10, 0.00), + (0.50, -0.18, 0.00), + ), + dtype=torch.float32, + ) + ) + return { + "arm_joint_position": torch.linspace(-0.5, 0.5, count * 7, dtype=torch.float32).reshape(count, 7), + "arm_joint_velocity": torch.zeros((count, 7), dtype=torch.float32), + "finger_joint_position": torch.tensor( + ((0.028, 0.028), (0.028, 0.028), (0.000, 0.000), (0.040, 0.040)), dtype=torch.float32 + ), + "finger_joint_velocity": torch.zeros((count, 2), dtype=torch.float32), + "finger_joint_target": torch.tensor( + ((0.021, 0.021), (0.021, 0.021), (0.000, 0.000), (0.040, 0.040)), dtype=torch.float32 + ), + "source_root_pose": source_pose, + "source_root_velocity": torch.zeros((count, 6), dtype=torch.float32), + "target_root_pose": target_pose, + "target_root_velocity": torch.zeros((count, 6), dtype=torch.float32), + "category": torch.tensor((1, 1, 0, 0), dtype=torch.int8), + "objective": torch.tensor((0.0, 1.0, -1.0, -1.0), dtype=torch.float32), + "objective_raw": torch.tensor((0.2, 0.8, -1.0, -1.0), dtype=torch.float32), + "objective_components": torch.tensor( + ((0.1, 0.2, 0.3), (0.8, 0.7, 0.9), (-1.0, -1.0, -1.0), (-1.0, -1.0, -1.0)), + dtype=torch.float32, + ), + "grasp_region": torch.tensor((0, 0, -1, -1), dtype=torch.int8), + "grasp_side": torch.tensor((0, 3, -1, -1), dtype=torch.int8), + "attempt_id": torch.arange(count, dtype=torch.int64), + "particle_layout_id": torch.zeros(count, dtype=torch.int32), + "ik_cost": torch.tensor((1.0e-5, 2.0e-5, 3.0e-5, 4.0e-5), dtype=torch.float32), + "ik_position_residual": torch.tensor((1.0e-4, 2.0e-4, 3.0e-4, 4.0e-4), dtype=torch.float32), + "ik_rotation_residual": torch.tensor((1.0e-3, 2.0e-3, 3.0e-3, 4.0e-3), dtype=torch.float32), + } + + +def _tiny_metadata() -> dict: + return { + "seed": 7, + "state_count": 4, + "category_names": ("non_grasping", "grasping"), + "category_counts": torch.tensor((2, 2), dtype=torch.int64), + "joint_names": tuple(f"panda_joint{index}" for index in range(1, 8)), + "frame": "environment", + "quaternion_order": "xyzw", + "particle_solver_state": "fresh_zero", + "source_region_center": torch.tensor((0.5, 0.0, 0.0595), dtype=torch.float32), + "objective_weights": torch.full((3,), 1.0 / 3.0, dtype=torch.float32), + "objective_component_names": ( + "source_distance", + "target_gated_inversion", + "target_alignment", + ), + "objective_raw_min_max": torch.tensor((0.2, 0.8), dtype=torch.float32), + "attempt_counts": torch.tensor((3, 4), dtype=torch.int64), + "rejection_counts": {"collision": torch.tensor((1, 2), dtype=torch.int64)}, + "sampling_and_validation_config": {"central_workspace_fraction": 0.9}, + "task_contract": { + "source_box_half": (0.028, 0.028, 0.0595), + "gripper_position_range": (0.0, 0.04), + "cup_grasp_height": 0.083, + "gripper_preload_pos": 0.024, + "gripper_grasp_reset_target": 0.021, + "gripper_contact_min_deflection": 0.001, + }, + } + + +def _tiny_payload( + states: dict[str, torch.Tensor] | None = None, + metadata: dict | None = None, +) -> dict: + particle_local_positions = torch.tensor( + (((-0.01, -0.01, 0.01), (0.01, -0.01, 0.01), (0.0, 0.01, 0.02)),), dtype=torch.float32 + ) + return build_reset_dataset_payload( + _tiny_states() if states is None else states, + particle_local_positions, + _tiny_metadata() if metadata is None else metadata, + FrankaPourResetDatasetGeneratorCfg( + grasping_count=2, + non_grasping_count=2, + near_pour_grasp_count=0, + batch_size=4, + ), + ) + + +def _tiny_production_payload() -> dict: + """Add the provenance emitted by dynamic validation to a tiny payload.""" + payload = _tiny_payload() + payload["metadata"]["dynamic_validation"] = { + "source_content_sha256": "0" * 64, + "steps": 60, + "settle_steps": 8, + "failure_dwell_steps": 2, + "failure_counts": {"nonfinite": 0, "grasp_lost": 1}, + "balance_trimmed": 0, + } + payload["content_sha256"] = reset_dataset_content_sha256(payload) + return payload + + +def _validate_tiny_production_payload(payload: dict) -> None: + """Validate provenance while retaining the compact test fixture quotas.""" + validate_production_reset_dataset( + payload, + expected_grasping_count=2, + expected_non_grasping_count=2, + ) + + +def test_grasp_objective_components_match_known_geometric_states(): + half_sqrt_two = math.sqrt(0.5) + source_pose = torch.tensor( + ( + (0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0), + (0.15, 0.0, 0.15, 1.0, 0.0, 0.0, 0.0), + (0.075, 0.0, 0.075, half_sqrt_two, 0.0, 0.0, half_sqrt_two), + ), + dtype=torch.float64, + ) + target_pose = _identity_poses( + torch.tensor(((1.0, 0.0, 0.0), (0.15, 0.0, 0.0), (0.0, 0.0, 0.0)), dtype=torch.float64) + ) + + components = grasp_objective_components( + source_pose, + target_pose, + source_region_center=(0.0, 0.0, 0.0), + cup_center_offset=(0.0, 0.0, 0.0), + target_rim_height=0.0, + ) + + expected = torch.tensor( + ((0.0, 0.0, 0.0), (1.0, 1.0, 1.0), (math.sqrt(0.5), 0.125, 0.25)), + dtype=torch.float64, + ) + torch.testing.assert_close(components, expected, rtol=1.0e-12, atol=1.0e-12) + + +def test_inversion_credit_is_zero_away_from_target(): + source_pose = torch.tensor(((0.10, 0.0, 0.2, 1.0, 0.0, 0.0, 0.0),), dtype=torch.float64) + target_pose = _identity_poses(torch.tensor(((0.0, 0.0, 0.0),), dtype=torch.float64)) + + components = grasp_objective_components( + source_pose, + target_pose, + source_region_center=(0.0, 0.0, 0.0), + cup_center_offset=(0.0, 0.0, 0.0), + target_rim_height=0.05, + inversion_gate_horizontal_threshold=0.07, + ) + + assert components[0, 0] > 0.0 + assert components[0, 1] == 0.0 + assert components[0, 2] > 0.0 + + +def test_grasp_objective_components_transform_cup_center_offset(): + source_pose = _identity_poses(torch.zeros((1, 3), dtype=torch.float64)) + target_pose = _identity_poses(torch.zeros((1, 3), dtype=torch.float64)) + + components = grasp_objective_components( + source_pose, + target_pose, + source_region_center=(0.0, 0.0, 0.1), + cup_center_offset=(0.0, 0.0, 0.1), + target_rim_height=0.025, + ) + + torch.testing.assert_close(components, torch.tensor(((0.0, 0.0, 0.5),), dtype=torch.float64)) + + +def test_source_root_position_seats_configured_grasp_point_at_tcp_not_cup_center(): + half_sqrt_two = math.sqrt(0.5) + tcp_position = torch.tensor(((0.4, -0.1, 0.3),), dtype=torch.float64) + tcp_quaternion = torch.tensor(((0.0, 0.0, 0.0, 1.0),), dtype=torch.float64) + source_quaternion = torch.tensor(((0.0, half_sqrt_two, 0.0, half_sqrt_two),), dtype=torch.float64) + grasp_offset = torch.tensor((0.0, 0.0, 0.083), dtype=torch.float64) + seating_offset = torch.tensor(((0.001, 0.0, -0.002),), dtype=torch.float64) + + source_position = source_root_position_from_tcp_grasp( + tcp_position, + tcp_quaternion, + source_quaternion, + grasp_offset, + seating_offset, + ) + restored_grasp_position = source_position + torch.nn.functional.normalize( + source_quaternion, + dim=-1, + ).new_tensor(((0.083, 0.0, 0.0),)) + + torch.testing.assert_close(restored_grasp_position, tcp_position + seating_offset) + + +def test_normalize_grasp_objectives_maps_global_extrema_to_unit_interval(): + raw = torch.tensor((2.0, 4.0, 8.0), dtype=torch.float64) + + normalized = normalize_grasp_objectives(raw) + + torch.testing.assert_close(normalized, torch.tensor((0.0, 1.0 / 3.0, 1.0), dtype=torch.float64)) + assert normalized.dtype == raw.dtype + assert normalized.device == raw.device + + +def test_normalize_grasp_objectives_rejects_degenerate_range(): + with pytest.raises(ValueError): + normalize_grasp_objectives(torch.full((3,), 0.4)) + + +def test_sampler_defaults_define_the_required_exact_twenty_thousand_states(): + cfg = FrankaPourResetDatasetGeneratorCfg() + + assert cfg.grasping_count == 10_000 + assert cfg.non_grasping_count == 10_000 + assert cfg.grasping_count + cfg.non_grasping_count == 20_000 + + +def test_sampler_runtime_interface_accepts_oversampled_candidate_quotas(monkeypatch): + cfg = FrankaPourResetDatasetGeneratorCfg(grasping_count=8, non_grasping_count=4, near_pour_grasp_count=4) + env = SimpleNamespace(cfg=object(), device="cpu") + monkeypatch.setattr(FrankaPourResetDatasetGenerator, "_build_ik_context", lambda _self: None) + + sampler = FrankaPourResetDatasetGenerator(env, cfg) + + assert sampler.cfg.grasping_count == 8 + + +@pytest.mark.parametrize("near_pour_count", (0, 8)) +def test_sampler_runtime_interface_requires_both_grasp_regions(near_pour_count): + cfg = FrankaPourResetDatasetGeneratorCfg( + grasping_count=8, + non_grasping_count=4, + near_pour_grasp_count=near_pour_count, + ) + + with pytest.raises(ValueError, match="requires both broad and near-pour"): + FrankaPourResetDatasetGenerator(None, cfg) + + +def test_sampler_runtime_interface_rejects_unbalanced_grasp_region_quota(): + cfg = FrankaPourResetDatasetGeneratorCfg(grasping_count=8, non_grasping_count=4, near_pour_grasp_count=2) + + with pytest.raises(ValueError, match="near_pour_grasp_count must be divisible by four"): + FrankaPourResetDatasetGenerator(None, cfg) + + +def test_production_row_selection_enforces_exact_balanced_quotas(): + non_grasping_count = 10_002 + broad_per_side = 2_252 + near_per_side = 252 + category = [torch.zeros(non_grasping_count, dtype=torch.int8)] + region = [torch.full((non_grasping_count,), -1, dtype=torch.int8)] + side = [torch.full((non_grasping_count,), -1, dtype=torch.int8)] + for region_id, per_side in ((0, broad_per_side), (1, near_per_side)): + for side_id in range(4): + category.append(torch.ones(per_side, dtype=torch.int8)) + region.append(torch.full((per_side,), region_id, dtype=torch.int8)) + side.append(torch.full((per_side,), side_id, dtype=torch.int8)) + states = { + "category": torch.cat(category), + "grasp_region": torch.cat(region), + "grasp_side": torch.cat(side), + } + states["objective"] = torch.linspace(-1.0, 1.0, states["category"].numel()) + valid = torch.ones(states["category"].numel(), dtype=torch.bool) + + keep, trimmed = select_production_reset_rows(states, valid) + + assert int(keep.sum()) == 20_000 + assert int((keep & (states["category"] == 0)).sum()) == 10_000 + assert int((keep & (states["category"] == 1)).sum()) == 10_000 + assert int((keep & (states["grasp_region"] == 1)).sum()) == 1_000 + assert torch.equal(trimmed, valid & ~keep) + + near_side_zero = (states["grasp_region"] == 1) & (states["grasp_side"] == 0) + valid[torch.nonzero(near_side_zero, as_tuple=False).flatten()[:3]] = False + with pytest.raises(RuntimeError, match="near-pour grasping states for side 0"): + select_production_reset_rows(states, valid) + + +def test_oriented_box_overlap_checks_all_axes(): + centers_a = torch.zeros((3, 3), dtype=torch.float64) + centers_b = torch.tensor(((0.5, 0.0, 0.0), (2.1, 0.0, 0.0), (1.2, 1.2, 0.0)), dtype=torch.float64) + identity = torch.zeros((3, 4), dtype=torch.float64) + identity[:, 3] = 1.0 + rotated = identity.clone() + rotated[2, 2] = math.sin(math.pi / 8.0) + rotated[2, 3] = math.cos(math.pi / 8.0) + + overlap = oriented_boxes_overlap( + centers_a, + identity, + (1.0, 1.0, 1.0), + centers_b, + rotated, + (1.0, 1.0, 1.0), + ) + + assert torch.equal(overlap, torch.tensor((True, False, True))) + + +def test_near_pour_mask_requires_alignment_clearance_and_tilt(): + source_pose = torch.tensor( + ( + (0.0, 0.0, 0.25, 1.0, 0.0, 0.0, 0.0), + (0.2, 0.0, 0.25, 1.0, 0.0, 0.0, 0.0), + (0.0, 0.0, 0.10, 1.0, 0.0, 0.0, 0.0), + (0.0, 0.0, 0.25, 0.0, 0.0, 0.0, 1.0), + ), + dtype=torch.float64, + ) + target_pose = _identity_poses(torch.zeros((4, 3), dtype=torch.float64)) + + valid = above_target_tilted_mask( + source_pose, + target_pose, + cup_center_offset=(0.0, 0.0, 0.0), + target_rim_height=0.05, + max_horizontal_distance=0.02, + min_vertical_clearance=0.15, + min_tilt_angle=math.radians(120.0), + ) + + assert valid.tolist() == [True, False, False, False] + + +def test_oriented_box_support_requires_the_complete_footprint(): + poses = _identity_poses(torch.tensor(((0.5, 0.5, 0.0), (0.95, 0.5, 0.0)), dtype=torch.float64)) + + supported = oriented_box_supported_by_bounds( + poses, + (0.1, 0.1, 0.1), + (0.0, 0.0), + (1.0, 1.0), + clearance=0.01, + ) + + assert supported.tolist() == [True, False] + + +def test_reset_dataset_cache_payload_has_exact_schema_categories_values_and_hashes(): + payload = _tiny_payload() + + validate_reset_dataset(payload, expected_grasping_count=2, expected_non_grasping_count=2) + assert payload["schema_version"] == 6 + assert payload["format"] == "franka_pour_reset_dataset" + assert len(payload["contract_sha256"]) == 64 + assert len(payload["content_sha256"]) == 64 + assert torch.equal(payload["states"]["category"], torch.tensor((1, 1, 0, 0), dtype=torch.int8)) + torch.testing.assert_close( + payload["states"]["objective"], torch.tensor((0.0, 1.0, -1.0, -1.0), dtype=torch.float32) + ) + assert tuple(payload["particle_layouts"]["local_position"].shape) == (1, 3, 3) + assert tuple(payload["particle_layouts"]["local_velocity"].shape) == (1, 3, 3) + + +def test_reset_dataset_cache_hashes_are_deterministic_and_detect_tensor_tampering(): + first = _tiny_payload() + second = _tiny_payload() + + assert first["contract_sha256"] == second["contract_sha256"] + assert first["content_sha256"] == second["content_sha256"] + + tampered = deepcopy(first) + tampered["states"]["arm_joint_position"][0, 0] += 0.125 + with pytest.raises(ValueError, match="(?i)content|hash"): + validate_reset_dataset(tampered, expected_grasping_count=2, expected_non_grasping_count=2) + + +def test_reset_dataset_cache_round_trips_with_weights_only_loader(tmp_path): + payload = _tiny_payload() + path = tmp_path / "reset_dataset.pt" + torch.save(payload, path) + + loaded = torch.load(path, map_location="cpu", weights_only=True) + + validate_reset_dataset(loaded, expected_grasping_count=2, expected_non_grasping_count=2) + assert loaded["content_sha256"] == payload["content_sha256"] + + +def test_production_validation_rejects_candidate_dataset_without_dynamic_provenance(): + with pytest.raises(ValueError, match="dynamic_validation"): + _validate_tiny_production_payload(_tiny_payload()) + + +def test_production_validation_accepts_complete_dynamic_provenance(): + _validate_tiny_production_payload(_tiny_production_payload()) + + +def test_production_validation_rejects_mismatched_current_task_contract(): + payload = _tiny_production_payload() + current_contract = deepcopy(payload["metadata"]["task_contract"]) + current_contract["source_box_half"] = (0.030, 0.028, 0.0595) + + with pytest.raises(ValueError, match="task contract does not match"): + validate_production_reset_dataset( + payload, + expected_grasping_count=2, + expected_non_grasping_count=2, + expected_task_contract=current_contract, + ) + + +@pytest.mark.parametrize( + ("mutation", "message"), + ( + (lambda marker: marker.update(source_content_sha256="A" * 64), "source_content_sha256"), + (lambda marker: marker.update(steps=0), "steps"), + (lambda marker: marker.update(settle_steps=60), "settle_steps"), + (lambda marker: marker.update(failure_counts={}), "failure_counts"), + (lambda marker: marker.update(balance_trimmed=-1), "balance_trimmed"), + ), +) +def test_production_validation_rejects_malformed_dynamic_provenance(mutation, message): + payload = _tiny_production_payload() + mutation(payload["metadata"]["dynamic_validation"]) + payload["content_sha256"] = reset_dataset_content_sha256(payload) + + with pytest.raises(ValueError, match=message): + _validate_tiny_production_payload(payload) + + +@pytest.mark.parametrize( + ("field", "index", "value"), + ( + ("category", 0, 2), + ("objective", 0, 1.1), + ("objective", 2, 0.0), + ), +) +def test_reset_dataset_cache_rejects_invalid_category_or_objective_semantics(field, index, value): + states = _tiny_states() + states[field][index] = value + + with pytest.raises(ValueError): + payload = _tiny_payload(states=states) + validate_reset_dataset(payload, expected_grasping_count=2, expected_non_grasping_count=2) + + +def test_reset_dataset_cache_rejects_missing_required_state_and_wrong_expected_counts(): + states = _tiny_states() + del states["source_root_pose"] + with pytest.raises(ValueError): + payload = _tiny_payload(states=states) + validate_reset_dataset(payload, expected_grasping_count=2, expected_non_grasping_count=2) + + payload = _tiny_payload() + with pytest.raises(ValueError): + validate_reset_dataset(payload, expected_grasping_count=3, expected_non_grasping_count=2) + + +def test_reset_dataset_cache_rejects_grasp_target_without_required_close_command(): + states = _tiny_states() + states["finger_joint_target"][0] = 0.028 + + with pytest.raises(ValueError, match="reset target"): + _tiny_payload(states=states) + + +def test_reset_dataset_cache_rejects_non_grasp_target_that_changes_sampled_opening(): + states = _tiny_states() + states["finger_joint_target"][2, 0] -= 0.001 + + with pytest.raises(ValueError, match="opening"): + _tiny_payload(states=states) + + +def test_reset_dataset_task_id_selects_the_production_registration(): + spec = gym.spec(FRANKA_POUR_RESET_DATASET_TASK_ID) + + assert spec.kwargs["env_cfg_entry_point"].endswith(":FrankaPourEnvCfg_RESET_DATASET") diff --git a/source/isaaclab_tasks/test/contrib/test_franka_pour_reset_utils.py b/source/isaaclab_tasks/test/contrib/test_franka_pour_reset_utils.py new file mode 100644 index 000000000000..f978eb93aac3 --- /dev/null +++ b/source/isaaclab_tasks/test/contrib/test_franka_pour_reset_utils.py @@ -0,0 +1,409 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Focused tests for Franka Pour reset-bank tensor utilities.""" + +import math + +import pytest +import torch + +from isaaclab_tasks.contrib.franka_pour.reset_utils import ( + asymmetric_reset_offset_samples, + polar_workspace_cells, + reset_rotation_vector_samples, + sample_index_pools, + scale_randomization_rows_by_extent, +) + + +def test_scale_randomization_rows_preserves_balanced_rows_and_full_extent(): + rows = torch.tensor( + [ + [0.0, 0.0, 0.0], + [0.10, -0.04, 0.02], + [-0.10, 0.04, -0.02], + ] + ) + extents = (0.0, 0.01, 0.025, 0.05, 0.10, 0.20, 0.35, 0.55, 0.75, 1.0) + + scaled = scale_randomization_rows_by_extent(rows, extents) + + assert scaled.shape == (len(extents), *rows.shape) + for level, extent in enumerate(extents): + torch.testing.assert_close(scaled[level], rows * extent) + if extent == 0.0: + torch.testing.assert_close(scaled[level], torch.zeros_like(rows), rtol=0.0, atol=0.0) + else: + torch.testing.assert_close(torch.sign(scaled[level]), torch.sign(rows)) + torch.testing.assert_close(scaled[-1], rows, rtol=0.0, atol=0.0) + + +def test_reset_rotation_vectors_cover_angle_range_without_aligned_full_extent_sample(): + angle_range = (math.radians(15.0), math.radians(35.0)) + + vectors = reset_rotation_vector_samples(angle_range, 11, dtype=torch.float64) + magnitudes = torch.linalg.vector_norm(vectors, dim=-1) + + assert vectors.shape == (11, 3) + assert vectors.dtype == torch.float64 + assert float(magnitudes.amin()) == pytest.approx(angle_range[0]) + assert float(magnitudes.amax()) == pytest.approx(angle_range[1]) + assert bool(torch.all(magnitudes > 0.0)) + assert int(torch.linalg.matrix_rank(vectors)) == 3 + torch.testing.assert_close( + vectors, + reset_rotation_vector_samples(angle_range, 11, dtype=torch.float64), + rtol=0.0, + atol=0.0, + ) + + +def test_reset_rotation_vectors_allow_disabled_rotation_randomization(): + vectors = reset_rotation_vector_samples((0.0, 0.0), 5) + + torch.testing.assert_close(vectors, torch.zeros_like(vectors), rtol=0.0, atol=0.0) + + +@pytest.mark.parametrize( + "angle_range, sample_count, message", + [ + ((0.0,), 5, "two values"), + ((-0.1, 0.2), 5, "within"), + ((0.3, 0.2), 5, "within"), + ((0.0, math.pi + 0.1), 5, "within"), + ((0.0, float("nan")), 5, "finite"), + ((0.0, 0.2), 0, "positive integer"), + ], +) +def test_reset_rotation_vectors_reject_invalid_configuration(angle_range, sample_count, message): + with pytest.raises(ValueError, match=message): + reset_rotation_vector_samples(angle_range, sample_count) + + +def test_polar_workspace_cells_cover_sector_and_preserve_exact_nominal_row(): + nominal = torch.tensor([0.50, 0.0, 0.02], dtype=torch.float64) + radius_range = (0.48, 0.74) + azimuth_range = math.radians(25.0) + grid_size = 7 + + cells = polar_workspace_cells( + nominal, + radius_range=radius_range, + azimuth_half_range=azimuth_range, + grid_size=grid_size, + ) + + assert cells.shape == (grid_size**2, 3) + assert cells.dtype == nominal.dtype + assert cells.device == nominal.device + torch.testing.assert_close(cells[:, 2], nominal[2].expand(cells.shape[0]), rtol=0.0, atol=0.0) + nominal_ring = grid_size // 2 + nominal_row = nominal_ring * grid_size + grid_size // 2 + torch.testing.assert_close(cells[nominal_row], nominal, rtol=0.0, atol=1.0e-12) + + radius = torch.linalg.vector_norm(cells[:, :2], dim=-1).reshape(grid_size, grid_size) + azimuth = torch.atan2(cells[:, 1], cells[:, 0]).reshape(grid_size, grid_size) + expected_radii = torch.cat( + ( + torch.linspace(radius_range[0], float(nominal[0]), nominal_ring + 1, dtype=nominal.dtype)[:-1], + torch.linspace(float(nominal[0]), radius_range[1], nominal_ring + 1, dtype=nominal.dtype), + ) + ) + torch.testing.assert_close(radius, expected_radii[:, None].expand_as(radius)) + torch.testing.assert_close(radius[:, 0], radius[:, -1], rtol=0.0, atol=1.0e-12) + torch.testing.assert_close(azimuth[0], azimuth[-1], rtol=0.0, atol=1.0e-12) + assert bool(torch.all(radius[1:, 0] > radius[:-1, 0])) + assert bool(torch.all(azimuth[0, 1:] > azimuth[0, :-1])) + assert float(radius.amin()) == pytest.approx(radius_range[0]) + assert float(radius.amax()) == pytest.approx(radius_range[1]) + assert float(azimuth.amin()) == pytest.approx(-azimuth_range) + assert float(azimuth.amax()) == pytest.approx(azimuth_range) + + # The task's conservative Cartesian half-extents must contain the complete polar sector. + xy_offset = torch.abs(cells[:, :2] - nominal[:2]) + assert bool(torch.all(xy_offset <= nominal.new_tensor((0.25, 0.34)) + 1.0e-12)) + + +def test_polar_workspace_cells_centers_azimuth_on_nonzero_nominal_bearing(): + nominal_azimuth = 0.4 + nominal = torch.tensor( + (0.5 * math.cos(nominal_azimuth), 0.5 * math.sin(nominal_azimuth), 0.02), + dtype=torch.float64, + ) + + cells = polar_workspace_cells( + nominal, + radius_range=(0.4, 0.7), + azimuth_half_range=0.3, + grid_size=3, + ) + + azimuth = torch.atan2(cells[:, 1], cells[:, 0]).reshape(3, 3) + offsets = torch.atan2( + torch.sin(azimuth - nominal_azimuth), + torch.cos(azimuth - nominal_azimuth), + ) + torch.testing.assert_close(offsets, nominal.new_tensor((-0.3, 0.0, 0.3)).expand_as(offsets)) + torch.testing.assert_close( + cells, + polar_workspace_cells( + nominal, + radius_range=(0.4, 0.7), + azimuth_half_range=0.3, + grid_size=3, + ), + rtol=0.0, + atol=0.0, + ) + + +@pytest.mark.parametrize( + "nominal_radius, nominal_ring", + [ + (0.4, 0), + (0.8, 4), + ], +) +def test_polar_workspace_cells_supports_nominal_radius_at_range_endpoint(nominal_radius, nominal_ring): + grid_size = 5 + nominal = torch.tensor((nominal_radius, 0.0, 0.03), dtype=torch.float64) + + cells = polar_workspace_cells( + nominal, + radius_range=(0.4, 0.8), + azimuth_half_range=0.2, + grid_size=grid_size, + ) + + radii = torch.linalg.vector_norm(cells[:, :2], dim=-1).reshape(grid_size, grid_size) + expected_radii = torch.linspace(0.4, 0.8, grid_size, dtype=nominal.dtype) + torch.testing.assert_close(radii, expected_radii[:, None].expand_as(radii)) + nominal_row = nominal_ring * grid_size + grid_size // 2 + torch.testing.assert_close(cells[nominal_row], nominal, rtol=0.0, atol=0.0) + assert torch.unique(cells, dim=0).shape[0] == cells.shape[0] + + +@pytest.mark.parametrize("grid_size", [2, 4, 3.0, True]) +def test_polar_workspace_cells_rejects_invalid_grid_size(grid_size): + with pytest.raises(ValueError, match="odd integer"): + polar_workspace_cells( + (0.5, 0.0, 0.0), + radius_range=(0.4, 0.6), + azimuth_half_range=0.2, + grid_size=grid_size, + ) + + +@pytest.mark.parametrize( + "nominal, radius_range, azimuth_half_range, message", + [ + ((0.5, 0.0), (0.4, 0.6), 0.2, "three coordinates"), + ((float("nan"), 0.0, 0.0), (0.4, 0.6), 0.2, "finite"), + ((0.0, 0.0, 0.0), (0.4, 0.6), 0.2, "closed interval"), + ((0.5, 0.0, 0.0), (0.4,), 0.2, "two values"), + ((0.3, 0.0, 0.0), (0.4, 0.6), 0.2, "closed interval"), + ((0.7, 0.0, 0.0), (0.4, 0.6), 0.2, "closed interval"), + ((0.5, 0.0, 0.0), (-0.1, 0.6), 0.2, "strictly increasing"), + ((0.5, 0.0, 0.0), (0.6, 0.4), 0.2, "strictly increasing"), + ((0.5, 0.0, 0.0), (0.4, 0.6), 0.0, "azimuth_half_range"), + ((0.5, 0.0, 0.0), (0.4, 0.6), math.pi, "azimuth_half_range"), + ], +) +def test_polar_workspace_cells_rejects_invalid_geometry(nominal, radius_range, azimuth_half_range, message): + with pytest.raises(ValueError, match=message): + polar_workspace_cells( + nominal, + radius_range=radius_range, + azimuth_half_range=azimuth_half_range, + grid_size=3, + ) + + +def test_polar_workspace_cells_rejects_nonfloating_tensor_nominal(): + with pytest.raises(ValueError, match="floating-point dtype"): + polar_workspace_cells( + torch.tensor((1, 0, 0)), + radius_range=(0.4, 1.2), + azimuth_half_range=0.2, + grid_size=3, + ) + + +def test_asymmetric_reset_offsets_include_bounds_zero_and_antithetic_interior(): + lower = torch.tensor((0.0, -0.10, 0.0), dtype=torch.float64) + upper = torch.tensor((0.0, 0.10, 0.15), dtype=torch.float64) + + sample_count = 11 + offsets = asymmetric_reset_offset_samples(lower, upper, sample_count=sample_count) + + assert offsets.shape == (sample_count, 3) + assert offsets.dtype == lower.dtype + assert offsets.device == lower.device + torch.testing.assert_close(offsets[0], torch.zeros_like(lower), rtol=0.0, atol=0.0) + torch.testing.assert_close(offsets[1], lower, rtol=0.0, atol=0.0) + torch.testing.assert_close(offsets[2], upper, rtol=0.0, atol=0.0) + interior_pair_count = (sample_count - 3) // 2 + torch.testing.assert_close( + offsets[3::2] + offsets[4::2], + (lower + upper).expand(interior_pair_count, -1), + rtol=0.0, + atol=1.0e-12, + ) + assert bool(torch.all(offsets >= lower)) + assert bool(torch.all(offsets <= upper)) + torch.testing.assert_close(offsets.amin(dim=0), lower, rtol=0.0, atol=0.0) + torch.testing.assert_close(offsets.amax(dim=0), upper, rtol=0.0, atol=0.0) + + active_axes = upper > lower + normalized_interior = (offsets[3:, active_axes] - lower[active_axes]) / (upper[active_axes] - lower[active_axes]) + assert bool(torch.all((normalized_interior > 0.0) & (normalized_interior < 1.0))) + torch.testing.assert_close( + normalized_interior.reshape(-1, 2, int(active_axes.sum())).sum(dim=1), + torch.ones((interior_pair_count, int(active_axes.sum())), dtype=lower.dtype), + ) + expected_first = torch.frac(lower.new_tensor((0.754877666, 0.569840296, 0.438447187)) * 0.5)[active_axes] + torch.testing.assert_close(normalized_interior[0], expected_first) + torch.testing.assert_close( + offsets, + asymmetric_reset_offset_samples(lower, upper, sample_count), + rtol=0.0, + atol=0.0, + ) + + +def test_asymmetric_reset_offsets_support_minimum_count_and_degenerate_axes(): + offsets = asymmetric_reset_offset_samples( + (-0.2, 0.0, 0.0), + (0.0, 0.0, 0.15), + 3, + dtype=torch.float64, + ) + + expected = torch.tensor( + ((0.0, 0.0, 0.0), (-0.2, 0.0, 0.0), (0.0, 0.0, 0.15)), + dtype=torch.float64, + ) + torch.testing.assert_close(offsets, expected, rtol=0.0, atol=0.0) + + +@pytest.mark.parametrize("sample_count", [1, 4, 5.0, True]) +def test_asymmetric_reset_offsets_reject_invalid_sample_count(sample_count): + with pytest.raises(ValueError, match="odd integer"): + asymmetric_reset_offset_samples((-0.2, -0.1, 0.0), (0.0, 0.1, 0.15), sample_count) + + +@pytest.mark.parametrize( + "lower, upper, message", + [ + ((-0.2, -0.1), (0.0, 0.1, 0.15), "three coordinates"), + ((-0.2, -0.1, 0.0), (0.0, 0.1), "three coordinates"), + ((float("nan"), -0.1, 0.0), (0.0, 0.1, 0.15), "finite"), + ((0.01, -0.1, 0.0), (0.1, 0.1, 0.15), "contain zero"), + ((-0.2, -0.1, 0.0), (0.0, -0.01, 0.15), "contain zero"), + ((0.0, 0.0, 0.0), (0.0, 0.0, 0.0), "positive width"), + ], +) +def test_asymmetric_reset_offsets_reject_invalid_bounds(lower, upper, message): + with pytest.raises(ValueError, match=message): + asymmetric_reset_offset_samples(lower, upper, 5) + + +def test_asymmetric_reset_offsets_reject_nonfloating_or_mismatched_tensor_bounds(): + with pytest.raises(ValueError, match="floating-point dtype"): + asymmetric_reset_offset_samples(torch.tensor((-1, 0, 0)), torch.tensor((0, 1, 1)), 5) + with pytest.raises(ValueError, match="same dtype"): + asymmetric_reset_offset_samples( + torch.tensor((-0.2, -0.1, 0.0), dtype=torch.float32), + torch.tensor((0.0, 0.1, 0.15), dtype=torch.float64), + 5, + ) + + +def test_weighted_index_pool_sampling_is_repeatable_from_torch_seed(): + pools = (torch.tensor([4, 8, 12]), torch.tensor([1, 5, 9, 13])) + weights = (torch.tensor([1.0, 2.0, 1.0]), torch.tensor([1.0, 1.0, 3.0, 1.0])) + pool_ids = torch.tensor([0, 1, 1, 0, 1, 0, 0, 1]) + + with torch.random.fork_rng(): + torch.manual_seed(1234) + first = sample_index_pools(pools, pool_ids, weights=weights) + torch.manual_seed(1234) + second = sample_index_pools(pools, pool_ids, weights=weights) + + torch.testing.assert_close(first, second, rtol=0.0, atol=0.0) + + +@pytest.mark.parametrize( + "extents, message", + [ + ((), "must not be empty"), + ((-0.1, 1.0), "values in"), + ((0.8, 0.7, 1.0), "strictly increasing"), + ((0.5, 0.8), "must end at 1.0"), + ], +) +def test_scale_randomization_rows_rejects_invalid_extent_levels(extents, message): + with pytest.raises(ValueError, match=message): + scale_randomization_rows_by_extent(torch.ones((3, 2)), extents) + + +def test_scale_randomization_rows_rejects_nonfinite_offsets(): + with pytest.raises(ValueError, match="finite floating-point"): + scale_randomization_rows_by_extent(torch.tensor([[0.0, float("nan")]]), (1.0,)) + + +def test_scale_randomization_rows_normalizes_tolerant_final_extent_to_full_amplitude(): + rows = torch.tensor([[0.10, -0.04]]) + + scaled = scale_randomization_rows_by_extent(rows, (0.5, 1.0 - 5.0e-10)) + + torch.testing.assert_close(scaled[-1], rows, rtol=0.0, atol=0.0) + + +def test_weighted_index_pool_sampling_uses_multinomial_and_maps_global_rows(monkeypatch): + pools = (torch.tensor([4, 8]), torch.tensor([1, 5, 9])) + weights = (torch.tensor([0.75, 0.25]), torch.tensor([0.0, 1.0, 3.0])) + pool_ids = torch.tensor([0, 1, 1, 0]) + sampled_slots = iter((torch.tensor([1, 0]), torch.tensor([2, 1]))) + calls = [] + + def sample_weighted(pool_weights, count, replacement): + calls.append((pool_weights.clone(), count, replacement)) + return next(sampled_slots) + + monkeypatch.setattr(torch, "multinomial", sample_weighted) + sampled = sample_index_pools(pools, pool_ids, weights=weights) + + assert sampled.tolist() == [8, 9, 5, 4] + assert len(calls) == 2 + torch.testing.assert_close(calls[0][0], weights[0]) + torch.testing.assert_close(calls[1][0], weights[1]) + assert [(count, replacement) for _, count, replacement in calls] == [(2, True), (2, True)] + + +@pytest.mark.parametrize( + "weights, message", + [ + ((), "one tensor per index pool"), + ((torch.ones(3),), "shape and device"), + ((torch.ones(2, dtype=torch.int64),), "finite, nonnegative floating-point"), + ((torch.tensor([1.0, float("nan")]),), "finite, nonnegative floating-point"), + ((torch.tensor([1.0, -0.1]),), "finite, nonnegative floating-point"), + ((torch.zeros(2),), "positive sum"), + ], +) +def test_weighted_index_pool_sampling_rejects_invalid_weights(weights, message): + with pytest.raises(ValueError, match=message): + sample_index_pools((torch.tensor([4, 8]),), torch.tensor([0]), weights=weights) + + +def test_weighted_index_pool_sampling_rejects_weight_device_mismatch(): + with pytest.raises(ValueError, match="shape and device"): + sample_index_pools( + (torch.tensor([4, 8]),), + torch.tensor([0]), + weights=(torch.ones(2, device="meta"),), + ) diff --git a/source/isaaclab_tasks/test/contrib/test_franka_pour_runtime.py b/source/isaaclab_tasks/test/contrib/test_franka_pour_runtime.py new file mode 100644 index 000000000000..f7965b9e22a6 --- /dev/null +++ b/source/isaaclab_tasks/test/contrib/test_franka_pour_runtime.py @@ -0,0 +1,1283 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Headless runtime integration for batched Franka Pour MPM.""" + +from __future__ import annotations + +import math +import os +from unittest import mock + +import pytest + +_RUNTIME_UNAVAILABLE_REASON = "Isaac Sim runtime is unavailable because EXP_PATH is not set." +_RUNTIME_AVAILABLE = bool(os.environ.get("EXP_PATH")) + +if _RUNTIME_AVAILABLE: + from isaaclab.app import AppLauncher + + # Launch Kit before importing simulation-dependent modules. + app_launcher = AppLauncher(headless=True) + simulation_app = app_launcher.app + + import gymnasium as gym + import newton + import torch + import warp as wp + from isaaclab_newton.physics import NewtonManager + + import isaaclab.sim as sim_utils + import isaaclab.utils.math as math_utils + + import isaaclab_tasks # noqa: F401 + from isaaclab_tasks.contrib.franka_pour.pour_env_cfg import MPM_ENTRY, PANDA_ARM_JOINT_LIMITS + from isaaclab_tasks.contrib.franka_pour.reset_utils import ( + asymmetric_reset_offset_samples, + polar_workspace_cells, + ) + from isaaclab_tasks.utils.parse_cfg import parse_env_cfg + +pytestmark = [pytest.mark.isaacsim_ci, pytest.mark.newton_ci] + +_TASK_ID = "Isaac-Pour-Franka-v0" + + +def _require_cuda() -> None: + """Require a CUDA device for the Franka Pour runtime integration.""" + if not wp.is_cuda_available(): + pytest.skip("Franka Pour runtime integration requires a CUDA device.") + + +def _assert_scene_solver_roles(model) -> None: + """Check exact per-world task bodies and solver-only collision roles.""" + body_world = model.body_world.numpy() + body_mass = model.body_mass.numpy() + body_inv_mass = model.body_inv_mass.numpy() + body_flags = model.body_flags.numpy() + shape_body = model.shape_body.numpy() + shape_flags = model.shape_flags.numpy() + collide_shapes = int(newton.ShapeFlags.COLLIDE_SHAPES) + collide_particles = int(newton.ShapeFlags.COLLIDE_PARTICLES) + visible = int(newton.ShapeFlags.VISIBLE) + + for world in range(2): + bodies_by_name: dict[str, list[int]] = {} + for body_id, label in enumerate(model.body_label): + if int(body_world[body_id]) == world: + bodies_by_name.setdefault(str(label).rsplit("/", 1)[-1], []).append(body_id) + for name in ("SourceCup", "TargetCup", "SpillFloor"): + assert len(bodies_by_name.get(name, [])) == 1, (world, name, bodies_by_name.get(name)) + assert "TargetCupRigid" not in bodies_by_name + + target_body = bodies_by_name["TargetCup"][0] + assert int(body_flags[target_body]) & int(newton.BodyFlags.KINEMATIC) + assert float(body_mass[target_body]) == 0.0 + assert float(body_inv_mass[target_body]) == 0.0 + + expected_shapes = ( + ("SourceCup", "/SourceCup/ParticleCollider", False, True, False), + ("TargetCup", "/TargetCup/ParticleCollider", False, True, False), + ("TargetCup", "/TargetCup/Collision", True, False, False), + ("SpillFloor", "/SpillFloor/Collision", False, True, False), + ) + for body_name, suffix, rigid, particles, is_visible in expected_shapes: + body_id = bodies_by_name[body_name][0] + matches = [ + shape_id + for shape_id, label in enumerate(model.shape_label) + if int(shape_body[shape_id]) == body_id and str(label).endswith(suffix) + ] + assert len(matches) == 1, (world, body_name, matches) + flags = int(shape_flags[matches[0]]) + assert bool(flags & collide_shapes) is rigid + assert bool(flags & collide_particles) is particles + assert bool(flags & visible) is is_visible + + +def _make_runtime_cfg( + *, + use_cuda_graph: bool = False, + env_spacing: float = 2.5, + mpm_iterations: int = 2, +): + cfg = parse_env_cfg(_TASK_ID, device="cuda:0", num_envs=2) + cfg.seed = 37 + # Keep the existing solver/reset regressions on the full-task open-hand approach contract. + # The dedicated curriculum test below exercises mixed easy/full stages explicitly. + cfg.curriculum_start_stage = cfg.curriculum_stage_names.index("full") + cfg.curriculum_freeze = True + cfg.scene.env_spacing = env_spacing + cfg.decimation = 1 + cfg.physics_substeps = 1 + cfg.mpm_iterations = mpm_iterations + cfg.use_cuda_graph = use_cuda_graph + cfg.sim.render_interval = 1 + + entries = {entry.name: entry for entry in cfg.sim.physics.solver_cfg.entries} + assert entries[MPM_ENTRY].in_place + return cfg + + +def _assert_full_task_reset_state(task, *, arm_atol: float = 1.0e-5, gripper_atol: float = 1.0e-5) -> None: + """Check that a full-task reset is finite, open, and begins from a real approach.""" + arm_action = task.action_manager.get_term("arm_action") + arm_q = task._robot.data.joint_pos.torch[:, task._arm_joint_ids] + assert bool(torch.isfinite(arm_q).all()) + joint_limits = torch.as_tensor(PANDA_ARM_JOINT_LIMITS, device=task.device) + assert bool(torch.all(arm_q >= joint_limits[:, 0] - arm_atol)) + assert bool(torch.all(arm_q <= joint_limits[:, 1] + arm_atol)) + torch.testing.assert_close( + arm_action.processed_actions, + torch.zeros_like(arm_action.processed_actions), + rtol=0.0, + atol=0.0, + ) + torch.testing.assert_close( + arm_action.raw_actions, + torch.zeros_like(arm_action.raw_actions), + rtol=0.0, + atol=0.0, + ) + torch.testing.assert_close( + task.gripper_width(), + torch.full((task.num_envs,), 0.08, device=task.device), + rtol=0.0, + atol=gripper_atol, + ) + tcp_distance = torch.linalg.vector_norm(task.tcp_pos_e() - task.cup_grasp_point_e(), dim=-1) + minimum_distance = task.cfg.curriculum_randomized_reset_tcp_min_grasp_distance - 0.005 + assert bool(torch.all(tcp_distance >= minimum_distance)), ( + f"Full-task TCP reset is not a real approach: distance={tcp_distance.tolist()}" + ) + + +@pytest.mark.skipif(not _RUNTIME_AVAILABLE, reason=_RUNTIME_UNAVAILABLE_REASON) +def test_franka_pour_reset_uses_direct_actions_and_zero_arm_action_holds(): + """The task exposes direct 7+1 actions and zero arm increments hold the reset pose.""" + _require_cuda() + sim_utils.create_new_stage() + env = None + try: + env = gym.make(_TASK_ID, cfg=_make_runtime_cfg()) + task = env.unwrapped + task.sim._app_control_on_stop_handle = None + env.reset() + + assert task.action_manager.active_terms == ["arm_action", "gripper_action"] + assert task.action_manager.action_term_dim == [7, 1] + assert task.action_manager.total_action_dim == 8 + _assert_full_task_reset_state(task) + arm_action = task.action_manager.get_term("arm_action") + arm_before = task._robot.data.joint_pos.torch[:, task._arm_joint_ids].clone() + actions = torch.zeros((task.num_envs, task.action_manager.total_action_dim), device=task.device) + actions[:, -1] = 1.0 + env.step(actions) + torch.testing.assert_close( + arm_action.processed_actions, + torch.zeros_like(arm_action.processed_actions), + rtol=0.0, + atol=0.0, + ) + arm_after = task._robot.data.joint_pos.torch[:, task._arm_joint_ids] + torch.testing.assert_close(arm_after, arm_before, rtol=0.0, atol=2.0e-3) + assert bool(torch.all(task.state_finite())) + finally: + if env is not None: + env.close() + + +@pytest.mark.skipif(not _RUNTIME_AVAILABLE, reason=_RUNTIME_UNAVAILABLE_REASON) +def test_franka_pour_direct_arm_action_moves_without_hidden_reference(): + """A bounded arm increment moves its joint while all other direct coordinates remain zero.""" + _require_cuda() + sim_utils.create_new_stage() + env = None + try: + cfg = _make_runtime_cfg(mpm_iterations=24) + full_stage = cfg.curriculum_stage_names.index("full") + cfg.curriculum_start_stage = full_stage + cfg.curriculum_freeze = True + cfg.decimation = 2 + cfg.physics_substeps = 2 + env = gym.make(_TASK_ID, cfg=cfg) + task = env.unwrapped + task.sim._app_control_on_stop_handle = None + + env.reset() + assert bool(torch.all(task.curriculum_stage == full_stage)) + assert task.cfg.curriculum_freeze is True + _assert_full_task_reset_state(task) + + arm_action = task.action_manager.get_term("arm_action") + arm_before = task._robot.data.joint_pos.torch[:, task._arm_joint_ids].clone() + actions = torch.zeros((task.num_envs, task.action_manager.total_action_dim), device=task.device) + actions[:, 0] = 0.25 + actions[:, -1] = 1.0 + for _ in range(3): + env.step(actions) + wp.synchronize_device(NewtonManager.get_model().device) + torch.testing.assert_close( + arm_action.processed_actions[:, 0], + torch.full((task.num_envs,), 0.25 * task.cfg.actions.arm_action.scale, device=task.device), + rtol=0.0, + atol=1.0e-7, + ) + assert bool(torch.all(task._robot.data.joint_pos.torch[:, task._arm_joint_ids[0]] > arm_before[:, 0])) + assert bool(torch.all(task.state_finite())) + assert bool(torch.all(task.rigid_state_in_bounds())) + assert bool(torch.all(task.particles_in_workspace())) + finally: + if env is not None: + env.close() + + +@pytest.mark.skipif(not _RUNTIME_AVAILABLE, reason=_RUNTIME_UNAVAILABLE_REASON) +def test_franka_pour_randomized_stage_builds_safe_ik_resets_and_resets_selected_world(): + """Randomized full-task resets stay reachable, separated, bounded, and world-selective.""" + _require_cuda() + sim_utils.create_new_stage() + env = None + try: + cfg = _make_runtime_cfg() + randomized_stage = cfg.curriculum_stage_names.index("randomized") + cfg.curriculum_start_stage = randomized_stage + cfg.curriculum_randomization_start_level = len(cfg.curriculum_randomization_extent_levels) - 1 + cfg.curriculum_freeze = True + env = gym.make(_TASK_ID, cfg=cfg) + task = env.unwrapped + task.sim._app_control_on_stop_handle = None + env.reset() + + source_bank = task._randomized_source_pos_bank_t + source_yaw_bank = task._randomized_source_yaw_bank_t + source_quat_bank = task._randomized_source_quat_bank_t + target_bank = task._randomized_target_pos_bank_t + tcp_start_bank = task._randomized_tcp_pos_bank_t + tcp_start_quat_bank = task._randomized_tcp_quat_bank_t + tcp_rotation_vector_bank = task._randomized_tcp_rotation_vector_bank_t + arm_bank = task._randomized_arm_q_bank_t + assert source_bank.shape == target_bank.shape + assert source_bank.shape == tcp_start_bank.shape + assert source_yaw_bank.shape == (source_bank.shape[0],) + assert source_quat_bank.shape == tcp_start_quat_bank.shape == (source_bank.shape[0], 4) + assert tcp_rotation_vector_bank.shape == (source_bank.shape[0], 3) + assert source_bank.shape[0] == arm_bank.shape[0] > 1 + rows_per_extent = ( + task.cfg.curriculum_randomized_reset_ik_grid_size**2 + * task.cfg.curriculum_randomized_reset_ik_samples_per_source + ) + extent_levels = task.cfg.curriculum_randomization_extent_levels + assert source_bank.shape[0] == len(extent_levels) * rows_per_extent + source_center = torch.as_tensor(task.cfg.cup_reset_pos[:2], device=task.device) + source_range = torch.as_tensor(task.cfg.curriculum_randomized_source_position_range, device=task.device) + extent_pools = task._randomized_extent_index_pools + extent_weights = task._randomized_extent_index_weights + source_cell_counts = task._randomized_extent_source_cell_counts + minimum_variant_counts = task._randomized_extent_minimum_variant_counts + assert len(extent_pools) == len(extent_levels) + assert len(extent_weights) == len(extent_pools) + assert len(source_cell_counts) == len(extent_pools) + assert len(minimum_variant_counts) == len(extent_pools) + collision_free_counts = task._randomized_collision_free_candidate_count_t + assert collision_free_counts.shape == (source_bank.shape[0],) + joint6_within_open_branch = arm_bank[:, 5] <= task.cfg.curriculum_randomized_reset_joint6_max + assert all(0 < pool.numel() <= rows_per_extent for pool in extent_pools) + assert torch.unique(torch.cat(extent_pools)).numel() == sum(pool.numel() for pool in extent_pools) + target_center = torch.as_tensor(task.cfg.curriculum_randomized_target_center_xy, device=task.device) + target_range = torch.as_tensor(task.cfg.curriculum_randomized_target_position_range, device=task.device) + tcp_standoff = torch.as_tensor(task.cfg.curriculum_randomized_reset_tcp_standoff, device=task.device) + reset_offset_lower = torch.as_tensor( + task.cfg.curriculum_randomized_reset_tcp_offset_lower, + device=task.device, + ) + reset_offset_upper = torch.as_tensor( + task.cfg.curriculum_randomized_reset_tcp_offset_upper, + device=task.device, + ) + grasp_points = source_bank.clone() + grasp_points[:, 2] += task.cfg.cup_grasp_height + tcp_displacements = tcp_start_bank - grasp_points + cos_yaw = torch.cos(source_yaw_bank) + sin_yaw = torch.sin(source_yaw_bank) + tcp_displacements_c = torch.stack( + ( + cos_yaw * tcp_displacements[:, 0] + sin_yaw * tcp_displacements[:, 1], + -sin_yaw * tcp_displacements[:, 0] + cos_yaw * tcp_displacements[:, 1], + tcp_displacements[:, 2], + ), + dim=-1, + ) + tcp_jitter = tcp_displacements_c - tcp_standoff + grasp_tcp_quat_c = torch.as_tensor(task.cfg.cup_grasp_tcp_quat_c, device=task.device) + aligned_tcp_quat_bank = math_utils.quat_mul( + source_quat_bank, + grasp_tcp_quat_c.expand(source_quat_bank.shape[0], -1), + ) + tcp_rotation_angles = torch.linalg.vector_norm(tcp_rotation_vector_bank, dim=-1) + tcp_quat_error = math_utils.quat_mul( + math_utils.quat_conjugate(aligned_tcp_quat_bank), + tcp_start_quat_bank, + ) + tcp_quat_error_angles = torch.linalg.vector_norm(math_utils.axis_angle_from_quat(tcp_quat_error), dim=-1) + torch.testing.assert_close(tcp_quat_error_angles, tcp_rotation_angles, rtol=0.0, atol=1.0e-5) + samples_per_source = task.cfg.curriculum_randomized_reset_ik_samples_per_source + grid_size = task.cfg.curriculum_randomized_reset_ik_grid_size + nominal_source = torch.as_tensor(task.cfg.cup_reset_pos, device=task.device) + full_source_cells = polar_workspace_cells( + nominal_source, + radius_range=task.cfg.curriculum_randomized_source_radius_range, + azimuth_half_range=task.cfg.curriculum_randomized_source_azimuth_range, + grid_size=grid_size, + ) + full_reset_offsets = asymmetric_reset_offset_samples( + reset_offset_lower, + reset_offset_upper, + samples_per_source, + ) + nominal_source_cell = int(torch.argmin(torch.linalg.vector_norm(full_source_cells - nominal_source, dim=-1))) + zero_jitter_slot = int(torch.argmin(torch.linalg.vector_norm(full_reset_offsets, dim=-1))) + zero_bank_row = nominal_source_cell * samples_per_source + zero_jitter_slot + minimum_source_cells = math.ceil(grid_size**2 * task.cfg.curriculum_randomized_min_source_cell_fraction) + source_outer_half_x = task.cfg.source_cup_inner_width / 2.0 + task.cfg.source_cup_wall_thickness + source_outer_half_y = task.cfg.source_cup_inner_depth / 2.0 + task.cfg.source_cup_wall_thickness + target_outer_half_y = task.cfg.target_cup_inner_depth / 2.0 + task.cfg.target_cup_wall_thickness + for level, (extent, pool, pool_weights) in enumerate( + zip(extent_levels, extent_pools, extent_weights, strict=True) + ): + level_indices = torch.arange( + level * rows_per_extent, + (level + 1) * rows_per_extent, + device=task.device, + ) + all_source_cell_ids = torch.arange(rows_per_extent, device=task.device) // samples_per_source + branch_continuous = ( + task._randomized_reference_branch_delta_t[level_indices] + <= task._randomized_reference_branch_limit_t[level_indices] + ) + raw_valid = ( + (collision_free_counts[level_indices] > 0) + & joint6_within_open_branch[level_indices] + & branch_continuous + ) + raw_variants_per_cell = torch.bincount( + all_source_cell_ids[raw_valid], + minlength=grid_size**2, + ) + retained_cell = raw_variants_per_cell >= task.cfg.curriculum_randomized_min_reset_variants_per_source + if extent == 0.0: + expected_pool = torch.as_tensor((zero_bank_row,), device=task.device) + else: + eligible = raw_valid & retained_cell[all_source_cell_ids] + variant_limit = max(1, math.ceil(extent * samples_per_source)) + score_by_cell = torch.where( + eligible, + task._randomized_reference_branch_delta_t[level_indices], + torch.full_like(task._randomized_reference_branch_delta_t[level_indices], torch.inf), + ).reshape(grid_size**2, samples_per_source) + ordered_slots = torch.argsort(score_by_cell, dim=1, stable=True) + selected_by_cell = torch.zeros_like(score_by_cell, dtype=torch.bool) + selected_by_cell.scatter_(1, ordered_slots[:, :variant_limit], True) + expected_pool = level_indices[eligible & selected_by_cell.reshape(-1)] + torch.testing.assert_close(pool, expected_pool) + assert bool(torch.all(collision_free_counts[pool] > 0)) + assert bool(torch.all(joint6_within_open_branch[pool])) + assert pool_weights.shape == pool.shape + assert pool_weights.device == pool.device + assert bool(torch.all(torch.isfinite(pool_weights) & (pool_weights > 0.0))) + source_cell_ids = (pool - level * rows_per_extent) // samples_per_source + unique_cells, inverse_cells = torch.unique(source_cell_ids, sorted=True, return_inverse=True) + total_weight_per_cell = torch.zeros(unique_cells.numel(), device=task.device).scatter_add_( + 0, + inverse_cells, + pool_weights, + ) + torch.testing.assert_close( + total_weight_per_cell, + torch.ones_like(total_weight_per_cell), + rtol=1.0e-6, + atol=1.0e-6, + ) + assert unique_cells.numel() == source_cell_counts[level] + assert source_cell_counts[level] >= (1 if extent == 0.0 else minimum_source_cells) + if level == len(extent_levels) - 1: + radial_rings = torch.div(unique_cells, grid_size, rounding_mode="floor") + azimuth_slots = unique_cells.remainder(grid_size) + torch.testing.assert_close( + torch.unique(radial_rings, sorted=True), + torch.arange(grid_size, device=task.device), + ) + assert int(azimuth_slots.amin()) < grid_size // 2 + assert int(azimuth_slots.amax()) > grid_size // 2 + assert int(azimuth_slots.amax() - azimuth_slots.amin()) >= grid_size // 2 + rows_per_cell = torch.bincount(inverse_cells, minlength=unique_cells.numel()) + assert int(rows_per_cell.min()) == minimum_variant_counts[level] + expected_minimum_variants = min( + task.cfg.curriculum_randomized_min_reset_variants_per_source, + max(1, math.ceil(extent * samples_per_source)), + ) + assert minimum_variant_counts[level] >= (1 if extent == 0.0 else expected_minimum_variants) + + eligible_arm_q = arm_bank[pool] + if extent == 0.0: + torch.testing.assert_close( + eligible_arm_q, + eligible_arm_q[:1].expand_as(eligible_arm_q), + rtol=0.0, + atol=0.0, + ) + else: + arm_span = eligible_arm_q.amax(dim=0) - eligible_arm_q.amin(dim=0) + assert int((arm_span > 1.0e-3).sum()) >= 4 + arm_rank = int(torch.linalg.matrix_rank(eligible_arm_q - eligible_arm_q.mean(dim=0), tol=1.0e-4)) + assert arm_rank >= 3 + if level == len(extent_levels) - 1: + # The final task must contain genuinely different whole-arm postures, not only + # Cartesian cup motion with numerically distinct IK rows. The validated default + # bank spans 0.7--1.7 rad over all seven joints; keep a conservative margin for + # solver/platform variation while still enforcing substantial diversity. + assert int((arm_span > 0.5).sum()) >= 6 + assert arm_rank == eligible_arm_q.shape[1] + quantized_arm_q = torch.round(eligible_arm_q * 1.0e4).to(dtype=torch.int64) + assert torch.unique(quantized_arm_q, dim=0).shape[0] >= source_cell_counts[level] + for cell_slot in range(unique_cells.numel()): + cell_arm_q = quantized_arm_q[inverse_cells == cell_slot] + assert torch.unique(cell_arm_q, dim=0).shape[0] >= expected_minimum_variants + + source_offsets = source_bank[level_indices, :2] - source_center + expected_source_range = source_range * extent + source_cells = source_offsets.reshape(-1, samples_per_source, 2) + source_positions_by_cell = source_bank[level_indices, :2].reshape(-1, samples_per_source, 2) + torch.testing.assert_close( + source_cells, + source_cells[:, :1].expand_as(source_cells), + rtol=0.0, + atol=0.0, + ) + expected_cells = nominal_source + extent * (full_source_cells - nominal_source) + torch.testing.assert_close(source_cells[:, 0], expected_cells[:, :2] - source_center, rtol=0.0, atol=1.0e-6) + assert bool(torch.all(torch.abs(source_offsets) <= expected_source_range + 1.0e-6)) + + nominal_source_rows = torch.isclose( + source_bank[level_indices, :2], source_center, atol=1.0e-7, rtol=0.0 + ).all(dim=-1) + nominal_source_yaws = source_yaw_bank[level_indices][nominal_source_rows] + expected_nominal_source_rows = rows_per_extent if extent == 0.0 else samples_per_source + assert nominal_source_yaws.shape == (expected_nominal_source_rows,) + expected_yaw_range = task.cfg.curriculum_randomized_source_yaw_range * extent + assert float(nominal_source_yaws.amin()) == pytest.approx(-expected_yaw_range) + assert float(nominal_source_yaws.amax()) == pytest.approx(expected_yaw_range) + + expected_target_range = target_range * extent + assert bool( + torch.all(torch.abs(target_bank[level_indices, :2] - target_center) <= expected_target_range + 1.0e-6) + ) + target_positions_by_cell = target_bank[level_indices, :2].reshape( + grid_size**2, + samples_per_source, + 2, + ) + # Receiver X is cell-local. Receiver Y may vary across source-yaw variants because the + # exact rotated source support changes the minimum collision-free cup separation. + torch.testing.assert_close( + target_positions_by_cell[:, :, 0], + target_positions_by_cell[:, :1, 0].expand_as(target_positions_by_cell[:, :, 0]), + rtol=0.0, + atol=0.0, + ) + expected_unique_targets = 1 if extent == 0.0 else grid_size**2 + assert torch.unique(target_positions_by_cell[:, 0, 0]).numel() == expected_unique_targets + cell_index = torch.arange(grid_size**2, device=task.device, dtype=target_bank.dtype) + center_cell = torch.argmin(torch.linalg.vector_norm(full_source_cells - nominal_source, dim=-1)).to( + dtype=target_bank.dtype + ) + target_unit_x = torch.remainder((cell_index - center_cell) * 0.754877666 + 0.5, 1.0) + target_unit_y = torch.remainder((cell_index - center_cell) * 0.569840296 + 0.5, 1.0) + expected_target_x = ( + target_center[0] - expected_target_range[0] + 2.0 * expected_target_range[0] * target_unit_x + ) + torch.testing.assert_close( + target_positions_by_cell[:, 0, 0], + expected_target_x, + rtol=0.0, + atol=1.0e-6, + ) + level_minimum_y_separation = ( + source_outer_half_x * torch.abs(torch.sin(source_yaw_bank[level_indices])) + + source_outer_half_y * torch.abs(torch.cos(source_yaw_bank[level_indices])) + + target_outer_half_y + + task.cfg.curriculum_randomized_cup_clearance + ) + allowed_target_y_upper = torch.minimum( + target_center[1] + expected_target_range[1], + source_bank[level_indices, 1] - level_minimum_y_separation, + ) + target_y_lower = target_center[1] - expected_target_range[1] + expected_target_y = target_y_lower + target_unit_y.repeat_interleave(samples_per_source) * ( + allowed_target_y_upper - target_y_lower + ) + torch.testing.assert_close( + target_bank[level_indices, 1], + expected_target_y, + rtol=0.0, + atol=1.0e-6, + ) + expected_offsets = (full_reset_offsets * extent).repeat(grid_size**2, 1) + torch.testing.assert_close( + tcp_jitter[level_indices], + expected_offsets, + rtol=0.0, + atol=1.0e-6, + ) + assert bool(torch.all(tcp_jitter[level_indices] >= reset_offset_lower * extent - 1.0e-6)) + assert bool(torch.all(tcp_jitter[level_indices] <= reset_offset_upper * extent + 1.0e-6)) + + level_rotation_angles = tcp_rotation_angles[level_indices] + rotation_lower, rotation_upper = task.cfg.curriculum_randomized_reset_tcp_rotation_angle_range + if extent == 0.0: + torch.testing.assert_close( + level_rotation_angles, + torch.zeros_like(level_rotation_angles), + rtol=0.0, + atol=0.0, + ) + else: + assert float(level_rotation_angles.amin()) == pytest.approx(rotation_lower * extent, abs=1.0e-6) + assert float(level_rotation_angles.amax()) == pytest.approx(rotation_upper * extent, abs=1.0e-6) + + yaw_by_source = source_yaw_bank[level_indices].reshape( + grid_size**2, + samples_per_source, + ) + source_bearing = torch.atan2(source_positions_by_cell[:, 0, 1], source_positions_by_cell[:, 0, 0]) + local_yaw_by_source = yaw_by_source - source_bearing.unsqueeze(-1) + expected_yaw_marginal = torch.sort(local_yaw_by_source[0]).values.expand_as(local_yaw_by_source) + torch.testing.assert_close(torch.sort(local_yaw_by_source, dim=-1).values, expected_yaw_marginal) + assert float(local_yaw_by_source.amin()) == pytest.approx(-expected_yaw_range) + assert float(local_yaw_by_source.amax()) == pytest.approx(expected_yaw_range) + + assert bool(torch.all(task.curriculum_stage == randomized_stage)) + torch.testing.assert_close( + torch.linalg.vector_norm(source_quat_bank, dim=-1), + torch.ones(source_bank.shape[0], device=task.device), + rtol=0.0, + atol=1.0e-6, + ) + expected_source_quat = torch.zeros_like(source_quat_bank) + expected_source_quat[:, 2] = torch.sin(0.5 * source_yaw_bank) + expected_source_quat[:, 3] = torch.cos(0.5 * source_yaw_bank) + torch.testing.assert_close(source_quat_bank, expected_source_quat, rtol=0.0, atol=1.0e-6) + + final_pool = extent_pools[-1] + final_rotation_angles = tcp_rotation_angles[final_pool] + assert float(final_rotation_angles.amin()) >= ( + task.cfg.curriculum_randomized_reset_tcp_rotation_angle_range[0] - 1.0e-6 + ) + assert int(torch.linalg.matrix_rank(tcp_rotation_vector_bank[final_pool], tol=1.0e-5)) == 3 + final_tcp_jitter = tcp_jitter[final_pool] + final_tcp_jitter_span = final_tcp_jitter.amax(dim=0) - final_tcp_jitter.amin(dim=0) + assert bool(torch.all(final_tcp_jitter_span > torch.tensor((0.04, 0.08, 0.06), device=task.device))) + + assert bool( + torch.all( + torch.linalg.vector_norm(tcp_displacements, dim=-1) + >= task.cfg.curriculum_randomized_reset_tcp_min_grasp_distance - 1.0e-6 + ) + ) + assert bool(torch.all(task._reach_source_yaw_bank_t == 0.0)) + reach_count = task._reach_arm_q_bank_t.shape[0] + assert task._reach_tcp_pos_bank_t.shape == (reach_count, 3) + assert task._reach_source_yaw_bank_t.shape == (reach_count,) + assert task._reach_pregrasp_arm_q_bank_t.shape == task._reach_arm_q_bank_t.shape + assert task._reach_midgrasp_arm_q_bank_t.shape == task._reach_arm_q_bank_t.shape + assert task._reach_grasp_arm_q_bank_t.shape == task._reach_arm_q_bank_t.shape + assert task._reach_reset_ik_cost_t.shape == (reach_count,) + assert task._reach_reset_ik_margin_t.shape == (reach_count,) + assert task._reach_collision_free_candidate_count_t.shape == (reach_count,) + assert bool(torch.all(task._reach_collision_free_candidate_count_t > 0)) + assert task.cfg.curriculum_randomized_min_reset_variants_per_source <= reach_count <= samples_per_source + assert bool(torch.all(task._reach_reset_ik_cost_t <= task.cfg.curriculum_randomized_reset_ik_max_cost)) + assert bool(torch.all(task._reach_reset_ik_margin_t >= task.cfg.curriculum_randomized_reset_ik_joint_margin)) + quantized_reach_arm_q = torch.round(task._reach_arm_q_bank_t * 1.0e4).to(dtype=torch.int64) + assert torch.unique(quantized_reach_arm_q, dim=0).shape[0] >= ( + task.cfg.curriculum_randomized_min_reset_variants_per_source + ) + + minimum_y_separation = ( + source_outer_half_x * torch.abs(torch.sin(source_yaw_bank)) + + source_outer_half_y * torch.abs(torch.cos(source_yaw_bank)) + + target_outer_half_y + + task.cfg.curriculum_randomized_cup_clearance + ) + assert bool(torch.all(source_bank[:, 1] - target_bank[:, 1] >= minimum_y_separation - 1.0e-6)) + waypoint_costs = ( + task._randomized_reset_ik_cost_t, + task._randomized_pregrasp_ik_cost_t, + task._randomized_midgrasp_ik_cost_t, + task._randomized_grasp_ik_cost_t, + task._randomized_carry_ik_cost_t, + task._randomized_pour_ik_cost_t, + task._randomized_tilt_ik_cost_t, + ) + waypoint_margins = ( + task._randomized_reset_ik_margin_t, + task._randomized_pregrasp_ik_margin_t, + task._randomized_midgrasp_ik_margin_t, + task._randomized_grasp_ik_margin_t, + task._randomized_carry_ik_margin_t, + task._randomized_pour_ik_margin_t, + task._randomized_tilt_ik_margin_t, + ) + assert all(cost.shape == (source_bank.shape[0],) for cost in waypoint_costs) + assert all(margin.shape == (source_bank.shape[0],) for margin in waypoint_margins) + eligible_indices = torch.cat(extent_pools) + assert all( + bool(torch.all(cost[eligible_indices] <= task.cfg.curriculum_randomized_reset_ik_max_cost)) + for cost in waypoint_costs + ) + assert all( + bool(torch.all(margin[eligible_indices] >= task.cfg.curriculum_randomized_reset_ik_joint_margin)) + for margin in waypoint_margins + ) + waypoint_bank = torch.stack( + ( + task._randomized_arm_q_bank_t, + task._randomized_pregrasp_arm_q_bank_t, + task._randomized_midgrasp_arm_q_bank_t, + task._randomized_grasp_arm_q_bank_t, + task._randomized_carry_arm_q_bank_t, + task._randomized_pour_arm_q_bank_t, + task._randomized_tilt_arm_q_bank_t, + ), + dim=1, + ) + joint_limits = torch.as_tensor(PANDA_ARM_JOINT_LIMITS, device=task.device) + joint_lower = joint_limits[:, 0] + joint_upper = joint_limits[:, 1] + eligible_waypoint_bank = waypoint_bank[eligible_indices] + assert bool(torch.isfinite(eligible_waypoint_bank).all()) + assert bool(torch.all(eligible_waypoint_bank >= joint_lower)) + assert bool(torch.all(eligible_waypoint_bank <= joint_upper)) + + # The zero-amplitude frontier is the exact behavioral continuation of the mastered full + # task: source/receiver geometry and zero-jitter insertion path are nominal, and the + # transport/pour tail stays on the authored safe branch. + assert extent_levels[0] == 0.0 + zero_pool = extent_pools[0] + torch.testing.assert_close( + source_bank[zero_pool], + nominal_source.expand(zero_pool.numel(), -1), + rtol=0.0, + atol=0.0, + ) + nominal_target = torch.as_tensor(task.cfg.target_cup_reset_pos, device=task.device) + torch.testing.assert_close( + target_bank[zero_pool], + nominal_target.expand(zero_pool.numel(), -1), + rtol=0.0, + atol=0.0, + ) + reach_waypoint_prefix = torch.stack( + ( + task._reach_arm_q_bank_t, + task._reach_pregrasp_arm_q_bank_t, + task._reach_midgrasp_arm_q_bank_t, + task._reach_grasp_arm_q_bank_t, + ), + dim=1, + ) + torch.testing.assert_close( + waypoint_bank[zero_pool, :4], + reach_waypoint_prefix[0].expand(zero_pool.numel(), -1, -1), + rtol=0.0, + atol=0.0, + ) + nominal_waypoint_tail = torch.as_tensor( + ( + task.cfg.curriculum_carry_arm_q, + task.cfg.curriculum_pour_arm_q, + task.cfg.curriculum_pour_target_arm_q, + ), + device=task.device, + dtype=waypoint_bank.dtype, + ) + torch.testing.assert_close( + waypoint_bank[zero_pool, 4:], + nominal_waypoint_tail.expand(zero_pool.numel(), -1, -1), + rtol=0.0, + atol=0.0, + ) + + # Every row in the first nonzero frontier must stay on the nominal local IK branch. The + # production filter permits 0.04 rad plus three times the physical extent, or 0.07 rad at + # the one-percent frontier. Assert the actual configured contract rather than a loose + # order-one bound that could admit the elbow/wrist jump this curriculum is designed to avoid. + assert len(extent_levels) > 1 and extent_levels[1] > 0.0 + level_one_pool = extent_pools[1] + maximum_level_one_joint_delta = torch.abs( + waypoint_bank[level_one_pool] - waypoint_bank[zero_pool[0]].unsqueeze(0) + ).amax() + maximum_level_one_limit = task._randomized_reference_branch_limit_t[level_one_pool].amax() + assert float(maximum_level_one_joint_delta) <= float(maximum_level_one_limit) + 1.0e-6 + + selected = torch.tensor([0], device=task.device, dtype=torch.long) + actions = torch.zeros((task.num_envs, task.action_manager.total_action_dim), device=task.device) + actions[:, -1] = 1.0 + fixed_robot_root = task._robot.data.root_link_pose_w.torch[0].clone() + # The full-amplitude bank remains the physical safety census. Probe both boundaries of the + # preceding levels as well, which verifies that level-local random slots map to their own + # disjoint global bank rows rather than accidentally sampling the full-amplitude block. + for level, pool in enumerate(extent_pools): + task.set_curriculum_randomization_level(selected, level) + local_slots = range(pool.numel()) if level == len(extent_pools) - 1 else (0, pool.numel() - 1) + for local_slot in local_slots: + bank_index = int(pool[local_slot]) + sampled_slot = torch.tensor([local_slot], device=task.device, dtype=torch.long) + with mock.patch.object(torch, "multinomial", return_value=sampled_slot): + task.reset_pour_scene(selected) + torch.testing.assert_close( + task._robot.data.root_link_pose_w.torch[0], + fixed_robot_root, + rtol=0.0, + atol=0.0, + ) + achieved_tcp_pose = task.tcp_pose_e()[0] + arm_index = int(task._last_arm_bank_index[0]) + assert int(task._last_source_bank_index[0]) == bank_index + expected_tcp_position = tcp_start_bank[arm_index] + tcp_position_error = torch.linalg.vector_norm(achieved_tcp_pose[:3] - expected_tcp_position) + tcp_quat_alignment = torch.abs(torch.dot(achieved_tcp_pose[3:7], tcp_start_quat_bank[arm_index])) + assert float(tcp_position_error) < 0.005, ( + f"Reset bank row {bank_index} TCP position error is {float(tcp_position_error):.6g} m." + ) + assert float(tcp_quat_alignment) > 1.0 - 1.0e-5, ( + f"Reset bank row {bank_index} TCP quaternion alignment is {float(tcp_quat_alignment):.6g}." + ) + _, _, terminated, truncated, _ = env.step(actions) + assert not bool(terminated[0]), f"Reset bank row {bank_index} terminated after one step." + assert not bool(truncated[0]), f"Reset bank row {bank_index} timed out after one step." + assert bool(torch.all(task.state_finite())), f"Reset bank row {bank_index} produced non-finite state." + assert bool(torch.all(task.rigid_state_in_bounds())), ( + f"Reset bank row {bank_index} exceeded rigid bounds." + ) + assert bool(torch.all(task.particles_in_workspace())), f"Reset bank row {bank_index} lost particles." + torch.testing.assert_close( + task._robot.data.root_link_pose_w.torch[0], + fixed_robot_root, + rtol=0.0, + atol=0.0, + ) + env.reset() + + arm_action = task.action_manager.get_term("arm_action") + world_0_root = task._robot.data.root_link_pose_w.torch[0].clone() + world_1_q = task._robot.data.joint_pos.torch[1].clone() + world_1_root = task._robot.data.root_link_pose_w.torch[1].clone() + world_1_source = task._source_cup.data.root_link_pose_w.torch[1].clone() + world_1_target = task._target_cup.data.root_link_pose_w.torch[1].clone() + world_1_media = task._media.data.particle_pos_w.torch[1].clone() + world_1_raw_action = arm_action.raw_actions[1].clone() + world_1_command = arm_action.processed_actions[1].clone() + + task.set_curriculum_randomization_level(selected, len(extent_pools) - 1) + bank_index = int(extent_pools[-1][-1]) + sampled_slot = torch.tensor([extent_pools[-1].numel() - 1], device=task.device, dtype=torch.long) + with mock.patch.object(torch, "multinomial", return_value=sampled_slot): + task.reset_pour_scene(selected) + wp.synchronize_device(NewtonManager.get_model().device) + + source_index = int(task._last_source_bank_index[0]) + arm_index = int(task._last_arm_bank_index[0]) + target_index = int(task._last_target_bank_index[0]) + assert source_index == arm_index == target_index == bank_index + torch.testing.assert_close(task.cup_pose_e()[0, :3], source_bank[source_index], rtol=0.0, atol=1.0e-6) + torch.testing.assert_close(task.cup_pose_e()[0, 3:7], source_quat_bank[source_index], rtol=0.0, atol=1.0e-6) + torch.testing.assert_close(task.target_pose_e()[0, :3], target_bank[target_index], rtol=0.0, atol=1.0e-6) + torch.testing.assert_close( + task.target_pose_e()[0, 3:7], + torch.tensor((0.0, 0.0, 0.0, 1.0), device=task.device), + rtol=0.0, + atol=1.0e-6, + ) + torch.testing.assert_close( + task._robot.data.joint_pos.torch[0, task._arm_joint_ids], + arm_bank[arm_index], + rtol=0.0, + atol=1.0e-6, + ) + torch.testing.assert_close( + task.gripper_width()[0], + torch.tensor(task.gripper_open_width, device=task.device), + rtol=0.0, + atol=1.0e-6, + ) + tcp_start_error = torch.linalg.vector_norm(task.tcp_pos_e()[0] - tcp_start_bank[arm_index]) + assert float(tcp_start_error) < 0.005, f"Randomized reset TCP-start error is {float(tcp_start_error):.6g} m." + tcp_quat_alignment = torch.abs(torch.dot(task.tcp_pose_e()[0, 3:7], tcp_start_quat_bank[arm_index])) + assert float(tcp_quat_alignment) > 1.0 - 1.0e-5 + tcp_grasp_distance = torch.linalg.vector_norm(task.tcp_pos_e()[0] - task.cup_grasp_point_e()[0]) + assert float(tcp_grasp_distance) >= task.cfg.curriculum_independent_arm_min_tcp_distance - 0.005 + + torch.testing.assert_close(task._robot.data.root_link_pose_w.torch[0], world_0_root, rtol=0.0, atol=0.0) + torch.testing.assert_close(task._robot.data.joint_pos.torch[1], world_1_q, rtol=0.0, atol=0.0) + torch.testing.assert_close(task._robot.data.root_link_pose_w.torch[1], world_1_root, rtol=0.0, atol=0.0) + torch.testing.assert_close(task._source_cup.data.root_link_pose_w.torch[1], world_1_source, rtol=0.0, atol=0.0) + torch.testing.assert_close(task._target_cup.data.root_link_pose_w.torch[1], world_1_target, rtol=0.0, atol=0.0) + torch.testing.assert_close(task._media.data.particle_pos_w.torch[1], world_1_media, rtol=0.0, atol=0.0) + torch.testing.assert_close(arm_action.raw_actions[1], world_1_raw_action, rtol=0.0, atol=0.0) + torch.testing.assert_close(arm_action.processed_actions[1], world_1_command, rtol=0.0, atol=0.0) + + randomized_target = task._target_cup.data.root_link_pose_w.torch[0].clone() + env.step(actions) + wp.synchronize_device(NewtonManager.get_model().device) + torch.testing.assert_close( + task._target_cup.data.root_link_pose_w.torch[0], + randomized_target, + rtol=0.0, + atol=0.0, + ) + assert bool(torch.all(task.state_finite())) + assert bool(torch.all(task.particles_in_workspace())) + + assert "reach" not in task.cfg.curriculum_stage_names + approach_distances = [] + for approach_index, stage_name in enumerate( + ("approach_1", "approach_2", "approach_3", "approach_4", "approach_5", "approach_6") + ): + approach_stage = task.cfg.curriculum_stage_names.index(stage_name) + task.set_curriculum_stage(selected, approach_stage) + task.reset_pour_scene(selected) + wp.synchronize_device(NewtonManager.get_model().device) + torch.testing.assert_close( + task._robot.data.joint_pos.torch[0, task._arm_joint_ids], + task._curriculum_approach_arm_q_t[approach_index], + rtol=0.0, + atol=1.0e-6, + ) + approach_distances.append( + float(torch.linalg.vector_norm(task.tcp_pos_e()[0] - task.cup_grasp_point_e()[0])) + ) + assert all( + near_distance < far_distance + for near_distance, far_distance in zip(approach_distances, approach_distances[1:]) + ) + assert 0.005 < approach_distances[0] < task.cfg.rewards.grasp_lift_progress.params["grasp_reach_std"] + expected_pre_grasp_distance = 0.5 * torch.linalg.vector_norm( + torch.tensor(task.cfg.curriculum_randomized_reset_tcp_standoff, device=task.device) + ) + torch.testing.assert_close( + torch.tensor(approach_distances[-1], device=task.device), + expected_pre_grasp_distance, + rtol=0.0, + atol=0.005, + ) + + full_stage = task.cfg.curriculum_stage_names.index("full") + task.set_curriculum_stage(selected, full_stage) + task.reset_pour_scene(selected) + wp.synchronize_device(NewtonManager.get_model().device) + torch.testing.assert_close( + task.cup_pose_e()[0, :3], + torch.tensor(task.cfg.cup_reset_pos, device=task.device), + rtol=0.0, + atol=0.0, + ) + torch.testing.assert_close( + task.target_pose_e()[0, :3], + torch.tensor(task.cfg.target_cup_reset_pos, device=task.device), + rtol=0.0, + atol=0.0, + ) + torch.testing.assert_close( + task._robot.data.joint_pos.torch[0, task._arm_joint_ids], + task._reach_pregrasp_arm_q_bank_t[0], + rtol=0.0, + atol=1.0e-6, + ) + pregrasp_distance = torch.linalg.vector_norm(task.tcp_pos_e()[0] - task.cup_grasp_point_e()[0]) + expected_pregrasp_distance = torch.linalg.vector_norm( + torch.tensor(task.cfg.curriculum_randomized_reset_tcp_standoff, device=task.device) + ) + torch.testing.assert_close(pregrasp_distance, expected_pregrasp_distance, rtol=0.0, atol=0.005) + finally: + if env is not None: + env.close() + + +@pytest.mark.skipif(not _RUNTIME_AVAILABLE, reason=_RUNTIME_UNAVAILABLE_REASON) +def test_extreme_finite_rigid_state_resets_before_returned_observation(): + """A finite dynamics outlier must be reset before RSL-RL can normalize the next observation.""" + _require_cuda() + sim_utils.create_new_stage() + env = None + try: + env = gym.make(_TASK_ID, cfg=_make_runtime_cfg()) + task = env.unwrapped + task.sim._app_control_on_stop_handle = None + env.reset() + + real_scene_update = task.scene.update + arm_joint_id = task._arm_joint_ids[0] + + def update_then_inject_outlier(dt: float) -> None: + real_scene_update(dt) + task._robot.data.joint_vel.torch[0, arm_joint_id] = 1.0e6 + + actions = torch.zeros((task.num_envs, task.action_manager.total_action_dim), device=task.device) + with mock.patch.object(task.scene, "update", side_effect=update_then_inject_outlier): + observations, _, terminated, truncated, _ = env.step(actions) + + assert terminated.tolist() == [True, False] + assert truncated.tolist() == [False, False] + assert task.termination_manager.get_term("extreme_rigid_state").tolist() == [True, False] + assert bool(torch.isfinite(observations["policy"]).all()) + assert bool(torch.isfinite(observations["privileged"]).all()) + assert bool(torch.all(torch.abs(observations["policy"][0, 7:14]) <= task.cfg.state_bound_max_joint_velocity)) + assert float(task._success_dwell_count[0]) == 0.0 + finally: + if env is not None: + env.close() + + +@pytest.mark.skipif(not _RUNTIME_AVAILABLE, reason=_RUNTIME_UNAVAILABLE_REASON) +def test_franka_pour_curriculum_mixed_stage_reset_uses_direct_actions_and_stays_isolated(): + """Reset two worlds at different backward stages, command both directly, and step eagerly.""" + _require_cuda() + sim_utils.create_new_stage() + env = None + try: + env = gym.make(_TASK_ID, cfg=_make_runtime_cfg()) + task = env.unwrapped + task.sim._app_control_on_stop_handle = None + observations, info = env.reset() + + assert task.curriculum_manager.active_terms == ["stage"] + assert set(observations) == {"policy", "privileged"} + assert observations["policy"].shape == (2, 62) + assert observations["privileged"].shape == (2, 20) + assert bool(torch.isfinite(observations["policy"]).all()) + assert bool(torch.isfinite(observations["privileged"]).all()) + full_stage = task.cfg.curriculum_stage_names.index("full") + grasp_stage = task.cfg.curriculum_stage_names.index("grasp") + assert info["log"]["Curriculum/stage/stage"] == float(full_stage) + assert "Curriculum/stage/success_rate" in info["log"] + assert "Curriculum/stage/completed_episodes" in info["log"] + assert "Curriculum/stage/required_completed_episodes" in info["log"] + assert "Curriculum/stage/mastered" in info["log"] + + selected = torch.tensor([0], device=task.device, dtype=torch.long) + transport_distance = [] + for stage_name in ("pour", "near_carry", "mid_carry", "carry"): + stage_index = task.cfg.curriculum_stage_names.index(stage_name) + task.set_curriculum_stage(selected, stage_index) + task.reset_pour_scene(selected) + distance = torch.linalg.vector_norm(task.cup_pose_e()[0, :2] - task.target_pose_e()[0, :2]) + transport_distance.append(float(distance)) + torch.testing.assert_close( + task._robot.data.joint_pos.torch[0, task._arm_joint_ids], + task._curriculum_arm_q_t[stage_index], + rtol=0.0, + atol=1.0e-6, + ) + assert all(left < right for left, right in zip(transport_distance, transport_distance[1:], strict=True)) + + env_ids = torch.tensor([0, 1], device=task.device, dtype=torch.long) + task.set_curriculum_stage(torch.tensor([0], device=task.device), 0) + task.set_curriculum_stage(torch.tensor([1], device=task.device), full_stage) + task.reset_pour_scene(env_ids) + + arm_action = task.action_manager.get_term("arm_action") + arm_q = task._robot.data.joint_pos.torch[:, task._arm_joint_ids] + torch.testing.assert_close(arm_q[0], task._curriculum_arm_q_t[0], rtol=0.0, atol=1.0e-6) + assert bool(torch.isfinite(arm_q[1]).all()) + torch.testing.assert_close( + arm_action.processed_actions, + torch.zeros_like(arm_action.processed_actions), + rtol=0.0, + atol=0.0, + ) + torch.testing.assert_close( + task.gripper_width(), + torch.tensor([task.gripper_grasp_width, task.gripper_open_width], device=task.device), + rtol=0.0, + atol=1.0e-6, + ) + gripper_action = task.action_manager.get_term("gripper_action") + torch.testing.assert_close( + gripper_action.action_offset[:, 0], + torch.full((2,), task.cfg.gripper_preload_pos, device=task.device), + rtol=0.0, + atol=0.0, + ) + torch.testing.assert_close( + task.pour_target_frac, + torch.tensor([0.05, 0.30], device=task.device), + rtol=0.0, + atol=0.0, + ) + grasp_error = torch.linalg.vector_norm(task.tcp_pos_e() - task.cup_grasp_point_e(), dim=-1) + assert float(grasp_error[0]) < 2.0e-5 + assert float(grasp_error[1]) >= task.cfg.curriculum_randomized_reset_tcp_min_grasp_distance - 0.005 + assert float(task.cup_pose_e()[0, 2]) > 0.16 + source_target_distance = torch.linalg.vector_norm(task.cup_pose_e()[0, :2] - task.target_pose_e()[0, :2]) + assert float(source_target_distance) < task.cfg.rewards.task_progress.params["alignment_radius"] + torch.testing.assert_close( + task.cup_pose_e()[1, :3], + torch.tensor(task.cfg.cup_reset_pos, device=task.device), + rtol=0.0, + atol=0.0, + ) + + source_pose = task._source_cup.data.root_link_pose_w.torch + expected_media = task._sample_cup_media(source_pose[:, :3], source_pose[:, 3:7]) + torch.testing.assert_close(task._media.data.particle_pos_w.torch, expected_media, rtol=0.0, atol=0.0) + assert bool(torch.all(task.particles_in_workspace())) + + actions = torch.zeros((task.num_envs, task.action_manager.total_action_dim), device=task.device) + gripper_before = gripper_action.processed_actions.clone() + task.action_manager.process_action(actions) + torch.testing.assert_close( + arm_action.processed_actions, + torch.zeros_like(arm_action.processed_actions), + rtol=0.0, + atol=0.0, + ) + gripper_target = torch.full_like(gripper_before, task.cfg.gripper_preload_pos) + expected_gripper = torch.lerp(gripper_before, gripper_target, task.cfg.actions.gripper_action.alpha) + torch.testing.assert_close( + gripper_action.processed_actions, + expected_gripper, + rtol=0.0, + atol=0.0, + ) + + actions[1, -1] = 1.0 + env.step(actions) + for _ in range(9): + env.step(actions) + wp.synchronize_device(NewtonManager.get_model().device) + assert bool(torch.all(task.state_finite())) + assert bool(torch.all(task.particles_in_workspace())) + assert float(task.cup_pose_e()[0, 2]) > 0.12 + assert float(torch.linalg.vector_norm(task.tcp_pos_e()[0] - task.cup_grasp_point_e()[0])) < 0.03 + + world_1_q = task._robot.data.joint_pos.torch[1].clone() + world_1_cup = task._source_cup.data.root_link_pose_w.torch[1].clone() + world_1_media = task._media.data.particle_pos_w.torch[1].clone() + world_1_raw_action = arm_action.raw_actions[1].clone() + world_1_command = arm_action.processed_actions[1].clone() + world_1_threshold = task.pour_target_frac[1].clone() + + selected = torch.tensor([0], device=task.device, dtype=torch.long) + task.set_curriculum_stage(selected, 1) + task.reset_pour_scene(selected) + wp.synchronize_device(NewtonManager.get_model().device) + + torch.testing.assert_close(task._robot.data.joint_pos.torch[1], world_1_q, rtol=0.0, atol=0.0) + torch.testing.assert_close(task._source_cup.data.root_link_pose_w.torch[1], world_1_cup, rtol=0.0, atol=0.0) + torch.testing.assert_close(task._media.data.particle_pos_w.torch[1], world_1_media, rtol=0.0, atol=0.0) + torch.testing.assert_close(arm_action.raw_actions[1], world_1_raw_action, rtol=0.0, atol=0.0) + torch.testing.assert_close(arm_action.processed_actions[1], world_1_command, rtol=0.0, atol=0.0) + torch.testing.assert_close(task.pour_target_frac[1], world_1_threshold, rtol=0.0, atol=0.0) + + torch.testing.assert_close( + task._robot.data.joint_pos.torch[0, task._arm_joint_ids], + task._curriculum_arm_q_t[1], + rtol=0.0, + atol=1.0e-6, + ) + assert ( + float(torch.linalg.vector_norm(task.cup_pose_e()[0, :2] - task.target_pose_e()[0, :2])) + < (task.cfg.rewards.task_progress.params["alignment_radius"]) + ) + assert float(task.cup_pose_e()[0, 2]) > 0.16 + for _ in range(3): + env.step(actions) + assert bool(torch.all(task.state_finite())) + assert float(task.cup_pose_e()[0, 2]) > 0.12 + + task.set_curriculum_stage(selected, grasp_stage) + task.reset_pour_scene(selected) + torch.testing.assert_close( + task._robot.data.joint_pos.torch[0, task._arm_joint_ids], + task._reach_grasp_arm_q_bank_t[0], + rtol=0.0, + atol=1.0e-6, + ) + torch.testing.assert_close( + task.cup_pose_e()[0, :3], + torch.tensor(task.cfg.cup_reset_pos, device=task.device), + rtol=0.0, + atol=0.0, + ) + torch.testing.assert_close(task.gripper_width()[0], torch.tensor(task.gripper_open_width, device=task.device)) + torch.testing.assert_close(task.pour_target_frac[0], torch.tensor(0.3, device=task.device)) + for _ in range(3): + env.step(actions) + assert bool(torch.all(task.state_finite())) + assert float(task.cup_pose_e()[0, 2]) > -0.02 + + # Exercise the real manager lifecycle: consume grasp-stage success, advance globally, then + # let the reset event apply the near-grasp approach and clear only the selected world's latches. + curriculum_term = task.curriculum_manager._term_cfgs[0].func + curriculum_term.stage = grasp_stage + curriculum_term.success_rate = 0.0 + curriculum_term.resets_in_stage = 0 + task.cfg.curriculum_freeze = False + task.cfg.curriculum_min_resets_per_stage = 1 + task.cfg.curriculum_min_reset_cohorts_per_stage = 0.0 + task.cfg.curriculum_previous_stage_replay_fraction = 0.0 + task.cfg.curriculum_success_threshold = 0.5 + task.episode_length_buf[0] = 10 + task.episode_succeeded[0] = True + task.ep_max_target_frac[0] = 0.4 + task.episode_succeeded[1] = True + task.ep_max_target_frac[1] = 0.25 + delivered_term = task.reward_manager.get_term_cfg("delivered").func + task._target_entry_seen[:, 0] = True + task._held_delivered[:, 0] = True + delivered_term._previous_credit[:] = 0.1 + world_1_target_entry_seen = task._target_entry_seen[1].clone() + world_1_held_delivered = task._held_delivered[1].clone() + world_1_delivery_credit = delivered_term._previous_credit[1].clone() + + task._reset_idx(selected) + + first_approach_stage = task.cfg.curriculum_stage_names.index("approach_1") + assert curriculum_term.stage == first_approach_stage + assert int(task.curriculum_stage[0]) == first_approach_stage + assert not bool(task.episode_succeeded[0]) + assert float(task.ep_max_target_frac[0]) == 0.0 + assert not bool(torch.any(task._target_entry_seen[0])) + assert not bool(torch.any(task._held_delivered[0])) + assert float(delivered_term._previous_credit[0]) == 0.0 + assert bool(task.episode_succeeded[1]) + assert float(task.ep_max_target_frac[1]) == pytest.approx(0.25) + torch.testing.assert_close(task._target_entry_seen[1], world_1_target_entry_seen, rtol=0.0, atol=0.0) + torch.testing.assert_close(task._held_delivered[1], world_1_held_delivered, rtol=0.0, atol=0.0) + torch.testing.assert_close(delivered_term._previous_credit[1], world_1_delivery_credit, rtol=0.0, atol=0.0) + assert bool(torch.isfinite(task._robot.data.joint_pos.torch[0, task._arm_joint_ids]).all()) + torch.testing.assert_close( + task._robot.data.joint_pos.torch[0, task._arm_joint_ids], + task._curriculum_approach_arm_q_t[0], + rtol=0.0, + atol=1.0e-6, + ) + torch.testing.assert_close(task.gripper_width()[0], torch.tensor(0.08, device=task.device)) + torch.testing.assert_close(task.pour_target_frac[0], torch.tensor(0.30, device=task.device)) + assert task.extras["log"]["Curriculum/stage/stage"] == float(first_approach_stage) + env.step(actions) + assert bool(torch.all(task.state_finite())) + finally: + if env is not None: + env.close() + + +@pytest.mark.skipif(not _RUNTIME_AVAILABLE, reason=_RUNTIME_UNAVAILABLE_REASON) +def test_franka_pour_scene_owned_cups_use_public_state_and_leave_caller_cfg_unmodified(): + """The task resolves cloned cup assets without mutating its caller-owned config.""" + _require_cuda() + sim_utils.create_new_stage() + env = None + caller_cfg = _make_runtime_cfg(use_cuda_graph=False, env_spacing=2.5) + # Keep public views on the authoritative state after the eager step below. + caller_cfg.physics_substeps = 2 + assert caller_cfg.scene.source_cup is None + assert caller_cfg.scene.target_cup is None + assert caller_cfg.scene.media is None + try: + env = gym.make(_TASK_ID, cfg=caller_cfg) + task = env.unwrapped + task.sim._app_control_on_stop_handle = None + env.reset() + + assert caller_cfg.scene.source_cup is None + assert caller_cfg.scene.target_cup is None + assert caller_cfg.scene.media is None + assert task.cfg is not caller_cfg + + source_cup = task.scene["source_cup"] + target_cup = task.scene["target_cup"] + assert source_cup.num_instances == 2 + assert target_cup.num_instances == 2 + + source_pose_e = source_cup.data.root_link_pose_w.torch.clone() + source_pose_e[:, :3] -= task.scene.env_origins + target_pose_e = target_cup.data.root_link_pose_w.torch.clone() + target_pose_e[:, :3] -= task.scene.env_origins + torch.testing.assert_close(task.cup_pose_e(), source_pose_e) + torch.testing.assert_close(task.target_pose_e(), target_pose_e) + _assert_scene_solver_roles(NewtonManager.get_model()) + assert NewtonManager.get_model().particle_max_velocity == pytest.approx(task.cfg.particle_max_velocity) + + for legacy_attribute in ( + "_cup_body_ids", + "_target_body_ids", + "_cup_joint_q", + "_cup_joint_qd", + "_cup_articulation_ids", + ): + assert not hasattr(task, legacy_attribute) + + target_pose_before = target_cup.data.root_link_pose_w.torch.clone() + actions = torch.zeros((task.num_envs, task.action_manager.total_action_dim), device=task.device) + actions[:, -1] = 1.0 + env.step(actions) + wp.synchronize_device(NewtonManager.get_model().device) + torch.testing.assert_close( + target_cup.data.root_link_pose_w.torch, + target_pose_before, + rtol=0.0, + atol=0.0, + ) + finally: + if env is not None: + env.close() + + +@pytest.mark.skipif(not _RUNTIME_AVAILABLE, reason=_RUNTIME_UNAVAILABLE_REASON) +def test_franka_pour_offset_world_tables_support_both_cups_without_reset(): + """Supported cups retain their media beyond the former delayed-leak horizon.""" + _require_cuda() + sim_utils.create_new_stage() + env = None + try: + runtime_cfg = _make_runtime_cfg(use_cuda_graph=False, env_spacing=2.5, mpm_iterations=24) + runtime_cfg.physics_substeps = 2 + env = gym.make(_TASK_ID, cfg=runtime_cfg) + task = env.unwrapped + task.sim._app_control_on_stop_handle = None + env.reset() + + entries = {entry.name: entry for entry in task.cfg.sim.physics.solver_cfg.entries} + assert entries[MPM_ENTRY].solver_cfg.project_outside_colliders is False + assert task.num_envs == 2 + assert not torch.equal(task.scene.env_origins[0], task.scene.env_origins[1]) + torch.testing.assert_close(task.cup_pose_e()[:, 2], torch.zeros(2, device=task.device), rtol=0.0, atol=0.0) + + actions = torch.zeros((task.num_envs, task.action_manager.total_action_dim), device=task.device) + actions[:, -1] = 1.0 # keep the gripper open so this tests table support only + for step in range(180): + _, _, terminated, truncated, _ = env.step(actions) + done = terminated | truncated + assert not bool(torch.any(done)), ( + f"A resting environment reset at step {step}: " + f"terminated={terminated.tolist()}, truncated={truncated.tolist()}" + ) + + wp.synchronize_device(NewtonManager.get_model().device) + cup_z = task.cup_pose_e()[:, 2] + assert bool(torch.all(cup_z > -0.02)), f"A source cup fell through its local table: z={cup_z.tolist()}" + expected_count = torch.full_like(task.count_in_source(), task.particle_pos_e().shape[1]) + torch.testing.assert_close(task.count_in_source(), expected_count, rtol=0.0, atol=0.0) + torch.testing.assert_close(task.count_spilled(), torch.zeros_like(expected_count), rtol=0.0, atol=0.0) + torch.testing.assert_close(cup_z[0], cup_z[1], rtol=0.0, atol=2.0e-3) + source_fraction = task.count_in_source() / float(task.num_particles) + assert bool(torch.all(source_fraction >= 0.99)), ( + f"PIC27 failed to retain resting media in the source cup: fraction={source_fraction.tolist()}" + ) + finally: + if env is not None: + env.close() diff --git a/source/isaaclab_tasks/test/contrib/test_franka_pour_spacemouse_teleop.py b/source/isaaclab_tasks/test/contrib/test_franka_pour_spacemouse_teleop.py new file mode 100644 index 000000000000..41b9c6ef0621 --- /dev/null +++ b/source/isaaclab_tasks/test/contrib/test_franka_pour_spacemouse_teleop.py @@ -0,0 +1,128 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Tests for the Franka Pour joint-position SpaceMouse adapter.""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path + +import pytest +import torch + +from isaaclab_tasks.contrib.franka_pour.pour_env_cfg import FrankaPourEnvCfg_TELEOP + + +def _load_teleop_script(): + repo_root = Path(__file__).parents[4] + script_path = repo_root / "scripts/environments/teleoperation/teleop_franka_pour_spacemouse.py" + spec = importlib.util.spec_from_file_location("teleop_franka_pour_spacemouse", script_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_joint_targets_are_encoded_as_joint_position_actions(): + teleop = _load_teleop_script() + default_joint_pos = torch.tensor([[0.0, 1.0, -1.0]]) + joint_targets = torch.tensor([[0.25, 0.5, 2.0]]) + lower_limits = torch.tensor([[-1.0, -1.0, -1.5]]) + upper_limits = torch.tensor([[1.0, 2.0, 1.5]]) + + actions = teleop.joint_targets_to_actions( + joint_targets=joint_targets, + action_offset=default_joint_pos, + action_scale=0.5, + lower_limits=lower_limits, + upper_limits=upper_limits, + ) + + torch.testing.assert_close(actions, torch.tensor([[0.5, -1.0, 5.0]])) + + +def test_joint_targets_support_per_joint_action_scales(): + teleop = _load_teleop_script() + + actions = teleop.joint_targets_to_actions( + joint_targets=torch.tensor([[0.5, 1.5, -0.5]]), + action_offset=torch.zeros((1, 3)), + action_scale=torch.tensor([[0.5, 1.5, 0.5]]), + lower_limits=torch.full((1, 3), -2.0), + upper_limits=torch.full((1, 3), 2.0), + ) + + torch.testing.assert_close(actions, torch.tensor([[1.0, 1.0, -1.0]])) + + +def test_environment_action_keeps_gripper_separate_from_seven_arm_joints(): + teleop = _load_teleop_script() + arm_action = torch.arange(14, dtype=torch.float32).reshape(2, 7) + gripper_command = torch.tensor([1.0, -1.0]) + + action = teleop.compose_env_action(arm_action, gripper_command) + + assert action.shape == (2, 8) + torch.testing.assert_close(action[:, :7], arm_action) + torch.testing.assert_close(action[:, 7], gripper_command) + + +def test_tcp_offset_updates_the_translational_jacobian(): + teleop = _load_teleop_script() + jacobian = torch.zeros((1, 6, 3)) + jacobian[:, 3:, :] = torch.eye(3) + body_quat = torch.tensor([[0.0, 0.0, 0.0, 1.0]]) + offset_pos = torch.tensor([[0.0, 0.0, 0.1]]) + + shifted = teleop.apply_tcp_offset_to_jacobian(jacobian, body_quat, offset_pos) + + expected_linear = torch.tensor([[[0.0, 0.1, 0.0], [-0.1, 0.0, 0.0], [0.0, 0.0, 0.0]]]) + torch.testing.assert_close(shifted[:, :3], expected_linear) + torch.testing.assert_close(shifted[:, 3:], torch.eye(3).unsqueeze(0)) + + +def test_tcp_offset_is_rotated_from_the_hand_frame_before_shifting_the_jacobian(): + teleop = _load_teleop_script() + jacobian = torch.zeros((1, 6, 3)) + jacobian[:, 3:, :] = torch.eye(3) + sin_cos_45 = 2.0**-0.5 + body_quat = torch.tensor([[0.0, 0.0, sin_cos_45, sin_cos_45]]) + offset_pos = torch.tensor([[0.1, 0.0, 0.0]]) + + shifted = teleop.apply_tcp_offset_to_jacobian(jacobian, body_quat, offset_pos) + + expected_linear = torch.tensor([[[0.0, 0.0, -0.1], [0.0, 0.0, 0.0], [0.1, 0.0, 0.0]]]) + torch.testing.assert_close(shifted[:, :3], expected_linear) + torch.testing.assert_close(shifted[:, 3:], torch.eye(3).unsqueeze(0)) + + +def test_teleop_uses_the_registered_joint_position_task_and_kitless_launcher(): + repo_root = Path(__file__).parents[4] + script = repo_root / "scripts/environments/teleoperation/teleop_franka_pour_spacemouse.py" + source = script.read_text(encoding="utf-8") + + assert 'DEFAULT_TASK = "Isaac-Pour-Franka-Teleop-v0"' in source + assert "launch_simulation" in source + assert "AppLauncher" not in source + + +def test_teleop_config_finalizes_without_trajectory_controller_or_rl_distribution(): + import gymnasium as gym + + import isaaclab_tasks # noqa: F401 + + cfg = FrankaPourEnvCfg_TELEOP().finalize() + task_spec = gym.spec("Isaac-Pour-Franka-Teleop-v0") + + assert cfg.actions.arm_action.class_type.__name__ == "CurriculumJointPositionAction" + assert cfg.actions.gripper_action.force_open_before_phase_stage == -1 + assert cfg.actions.gripper_action.limit_to_preload is False + assert cfg.actions.gripper_action.default_position == pytest.approx(cfg.gripper_open_pos) + assert cfg.actions.gripper_action.neutral_position == pytest.approx(cfg.gripper_open_pos) + assert cfg.actions.gripper_action.scale == pytest.approx( + cfg.gripper_open_pos - cfg.actions.gripper_action.close_position + ) + assert cfg.terminations.time_out is None + assert "rsl_rl_cfg_entry_point" not in task_spec.kwargs diff --git a/source/isaaclab_tasks/test/contrib/test_franka_pour_visualization.py b/source/isaaclab_tasks/test/contrib/test_franka_pour_visualization.py new file mode 100644 index 000000000000..f47c8dab9d90 --- /dev/null +++ b/source/isaaclab_tasks/test/contrib/test_franka_pour_visualization.py @@ -0,0 +1,417 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Visualization regressions for the Franka Pour task.""" + +from __future__ import annotations + +import os + +import pytest + +_RUNTIME_AVAILABLE = bool(os.environ.get("EXP_PATH")) +_RUNTIME_UNAVAILABLE_REASON = "Isaac Sim runtime is unavailable because EXP_PATH is not set." +_TEST_DEVICE = os.environ.get("ISAACLAB_TEST_DEVICE", "cuda:0") + +if _RUNTIME_AVAILABLE: + from isaaclab.app import AppLauncher + + # Launch Kit before importing simulation-dependent modules. + app_launcher = AppLauncher(headless=True, enable_cameras=True, device=_TEST_DEVICE) + simulation_app = app_launcher.app + + import gymnasium as gym + import newton + import numpy as np + import torch + import warp as wp + from isaaclab_newton.physics import NewtonManager + from isaaclab_visualizers.kit import KitVisualizer, KitVisualizerCfg + from isaaclab_visualizers.newton import NewtonVisualizer, NewtonVisualizerCfg + + import usdrt + from pxr import Usd, UsdGeom + + import isaaclab.sim as sim_utils + + import isaaclab_tasks # noqa: F401 + from isaaclab_tasks.utils.parse_cfg import parse_env_cfg + +from isaaclab_tasks.contrib.franka_pour.pour_env_cfg import FrankaPourEnvCfg + +pytestmark = [pytest.mark.isaacsim_ci, pytest.mark.newton_ci] + +_TASK_ID = "Isaac-Pour-Franka-v0" +_RESET_DATASET_PLAY_TASK_ID = "Isaac-Pour-Franka-Reset-Dataset-Play-v0" +_SCENE_PARTITION_ENV_VAR = "ISAAC_LAB_ENABLE_ISAAC_RTX_PER_ENV_SCENE_PARTITION" + + +def test_franka_pour_viewer_frames_first_environment(): + """The task viewer should interpret its camera pose relative to environment zero.""" + cfg = FrankaPourEnvCfg() + + assert cfg.viewer.origin_type == "env" + assert cfg.viewer.env_index == 0 + + +def _make_visualization_cfg(): + cfg = parse_env_cfg(_TASK_ID, device=_TEST_DEVICE, num_envs=2) + cfg.seed = 37 + cfg.curriculum_start_stage = cfg.curriculum_stage_names.index("full") + cfg.curriculum_freeze = True + cfg.scene.env_spacing = 2.5 + cfg.decimation = 1 + cfg.physics_substeps = 2 + cfg.mpm_iterations = 2 + cfg.use_cuda_graph = False + cfg.sim.render_interval = 1 + # Keep the eager double-buffered manager on an even substep count, matching + # the production configuration's stable public state bindings. + cfg.sim.visualizer_cfgs = [ + KitVisualizerCfg(headless=True, randomly_sample_visible_envs=False), + NewtonVisualizerCfg( + headless=True, + show_particles=True, + enable_shadows=False, + enable_sky=False, + randomly_sample_visible_envs=False, + ), + ] + return cfg + + +def _assert_unpartitioned(prim, attribute_name: str) -> None: + attribute = prim.GetAttribute(attribute_name) + assert not attribute.IsValid() or not attribute.HasAuthoredValueOpinion(), ( + f"Unexpected authored {attribute_name!r} on {prim.GetPath()}." + ) + + +def _shape_matches(model, label_fragment: str) -> list[tuple[int, int, bool]]: + visible_flag = int(newton.ShapeFlags.VISIBLE) + flags = model.shape_flags.numpy() + worlds = model.shape_world.numpy() + matches = [ + (shape_id, int(worlds[shape_id]), bool(int(flags[shape_id]) & visible_flag)) + for shape_id, label in enumerate(model.shape_label) + if label_fragment in str(label) + ] + assert matches, f"No Newton shapes matched {label_fragment!r}." + return matches + + +def _assert_shape_distribution( + model, label_fragment: str, *, visible: bool, expected_per_world: dict[int, int] +) -> None: + matches = _shape_matches(model, label_fragment) + actual_per_world = { + world: sum(match_world == world for _, match_world, _ in matches) for world in expected_per_world + } + assert actual_per_world == expected_per_world, (label_fragment, matches) + assert all(match_visible is visible for _, _, match_visible in matches), (label_fragment, matches) + + +def _pose_xyzw_to_fabric_matrix(pose: np.ndarray) -> np.ndarray: + """Convert an XYZ + XYZW pose to Fabric's row-vector matrix convention.""" + position = pose[:3] + x, y, z, w = pose[3:7] / np.linalg.norm(pose[3:7]) + rotation = np.array( + [ + [1.0 - 2.0 * (y * y + z * z), 2.0 * (x * y - z * w), 2.0 * (x * z + y * w)], + [2.0 * (x * y + z * w), 1.0 - 2.0 * (x * x + z * z), 2.0 * (y * z - x * w)], + [2.0 * (x * z - y * w), 2.0 * (y * z + x * w), 1.0 - 2.0 * (x * x + y * y)], + ], + dtype=np.float64, + ) + matrix = np.eye(4, dtype=np.float64) + matrix[:3, :3] = rotation.T + matrix[3, :3] = position + return matrix + + +def _fabric_world_matrix(path: str) -> np.ndarray: + """Read a fully materialized Fabric hierarchy matrix without creating a view.""" + stage = NewtonManager._usdrt_stage + assert stage is not None + prim = stage.GetPrimAtPath(usdrt.Sdf.Path(path)) + assert prim.IsValid(), path + xformable = usdrt.Rt.Xformable(prim) + assert xformable.GetFabricHierarchyLocalMatrixAttr().IsValid(), path + attribute = xformable.GetFabricHierarchyWorldMatrixAttr() + assert attribute.IsValid(), path + return np.asarray(attribute.Get(), dtype=np.float64) + + +def _assert_descendant_follows_body( + usd_stage, xform_cache: UsdGeom.XformCache, body_path: str, descendant_path: str, expected_body: np.ndarray +) -> None: + """Compare a descendant against its authored body-relative transform.""" + body_prim = usd_stage.GetPrimAtPath(body_path) + descendant_prim = usd_stage.GetPrimAtPath(descendant_path) + assert body_prim.IsValid(), body_path + assert descendant_prim.IsValid(), descendant_path + authored_body = np.asarray(xform_cache.GetLocalToWorldTransform(body_prim), dtype=np.float64) + authored_descendant = np.asarray(xform_cache.GetLocalToWorldTransform(descendant_prim), dtype=np.float64) + body_relative = authored_descendant @ np.linalg.inv(authored_body) + expected_descendant = body_relative @ expected_body + np.testing.assert_allclose( + _fabric_world_matrix(descendant_path), + expected_descendant, + rtol=0.0, + atol=1.0e-5, + err_msg=descendant_path, + ) + + +def _assert_visual_descendants_follow_fabric_bodies(task) -> None: + """Check body roots, visual instance roots, and proxy meshes in raw Fabric.""" + usd_stage = sim_utils.get_current_stage() + xform_cache = UsdGeom.XformCache(Usd.TimeCode.Default()) + robot_poses = task.scene["robot"].data.body_link_pose_w.torch.detach().cpu().numpy() + + for env_id in range(task.num_envs): + robot_path = f"/World/envs/env_{env_id}/Robot" + for body_id, body_name in enumerate(task.scene["robot"].body_names): + body_path = f"{robot_path}/{body_name}" + expected_body = _pose_xyzw_to_fabric_matrix(robot_poses[env_id, body_id]) + np.testing.assert_allclose( + _fabric_world_matrix(body_path), + expected_body, + rtol=0.0, + atol=1.0e-5, + err_msg=body_path, + ) + _assert_descendant_follows_body(usd_stage, xform_cache, body_path, f"{body_path}/visuals", expected_body) + _assert_descendant_follows_body( + usd_stage, xform_cache, body_path, f"{body_path}/visuals/{body_name}", expected_body + ) + + table_path = f"/World/envs/env_{env_id}/Table" + table_world = _fabric_world_matrix(table_path) + _assert_descendant_follows_body(usd_stage, xform_cache, table_path, f"{table_path}/Visuals", table_world) + _assert_descendant_follows_body( + usd_stage, xform_cache, table_path, f"{table_path}/Visuals/TableGeom", table_world + ) + + for scene_name, prim_name in (("source_cup", "SourceCup"), ("target_cup", "TargetCup")): + poses = task.scene[scene_name].data.root_link_pose_w.torch.detach().cpu().numpy() + for env_id in range(task.num_envs): + body_path = f"/World/envs/env_{env_id}/{prim_name}" + expected_body = _pose_xyzw_to_fabric_matrix(poses[env_id]) + np.testing.assert_allclose( + _fabric_world_matrix(body_path), + expected_body, + rtol=0.0, + atol=1.0e-5, + err_msg=body_path, + ) + _assert_descendant_follows_body(usd_stage, xform_cache, body_path, f"{body_path}/geometry", expected_body) + _assert_descendant_follows_body( + usd_stage, xform_cache, body_path, f"{body_path}/geometry/mesh", expected_body + ) + + +@pytest.mark.skipif(not _RUNTIME_AVAILABLE, reason=_RUNTIME_UNAVAILABLE_REASON) +def test_franka_pour_reset_dataset_play_captures_sparse_outer_graph(): + """Sparse reset-dataset playback should capture and replay the complete coupled solve graph.""" + sim_utils.create_new_stage() + env = None + try: + cfg = parse_env_cfg(_RESET_DATASET_PLAY_TASK_ID, device=_TEST_DEVICE, num_envs=1) + cfg.seed = 37 + cfg.decimation = 1 + cfg.physics_substeps = 1 + cfg.mpm_iterations = 2 + cfg.sim.render_interval = 1 + cfg.sim.visualizer_cfgs = [KitVisualizerCfg(headless=True, randomly_sample_visible_envs=False)] + env = gym.make(_RESET_DATASET_PLAY_TASK_ID, cfg=cfg) + task = env.unwrapped + task.sim._app_control_on_stop_handle = None + env.reset() + + actions = torch.zeros((1, task.action_manager.total_action_dim), device=task.device) + env.step(actions) + env.step(actions) + wp.synchronize_device(task.device) + + mpm_solver = NewtonManager._solver.solver("media") + assert NewtonManager._graph is not None + assert mpm_solver._use_cuda_graph is True + assert bool(torch.all(task.state_finite())) + finally: + if env is not None: + env.close() + + +@pytest.mark.skipif(not _RUNTIME_AVAILABLE, reason=_RUNTIME_UNAVAILABLE_REASON) +def test_franka_pour_kit_and_newton_visualize_both_worlds(monkeypatch: pytest.MonkeyPatch): + """Both renderers should consume the two spaced worlds without adding another world offset.""" + monkeypatch.delenv(_SCENE_PARTITION_ENV_VAR, raising=False) + sim_utils.create_new_stage() + env = None + try: + env = gym.make(_TASK_ID, cfg=_make_visualization_cfg()) + task = env.unwrapped + task.sim._app_control_on_stop_handle = None + + # Validate the natural Newton-to-Fabric setup before a view can initialize + # or otherwise mask missing descendant hierarchy attributes. + _assert_visual_descendants_follow_fabric_bodies(task) + + env.reset() + + stage = sim_utils.get_current_stage() + model = NewtonManager.get_model() + wp.synchronize_device(model.device) + assert str(model.device) == _TEST_DEVICE + + _assert_visual_descendants_follow_fabric_bodies(task) + + expected_origins = np.array([[1.25, 0.0, 0.0], [-1.25, 0.0, 0.0]], dtype=np.float32) + origins = task.scene.env_origins.detach().cpu().numpy() + np.testing.assert_allclose(origins, expected_origins, rtol=0.0, atol=1.0e-6) + + local_positions = { + "Robot": tuple(task.cfg.scene.robot.init_state.pos), + "Table": tuple(task.cfg.scene.table.init_state.pos), + "SourceCup": tuple(task.cfg.cup_reset_pos), + "TargetCup": tuple(task.cfg.target_cup_reset_pos), + } + for env_id, origin in enumerate(origins): + env_root = stage.GetPrimAtPath(f"/World/envs/env_{env_id}") + assert env_root.IsValid() + assert UsdGeom.Imageable(env_root).ComputeVisibility() == UsdGeom.Tokens.inherited + _assert_unpartitioned(env_root, "primvars:omni:scenePartition") + + for asset_name, expected_local_position in local_positions.items(): + asset_path = f"/World/envs/env_{env_id}/{asset_name}" + asset_prim = stage.GetPrimAtPath(asset_path) + assert asset_prim.IsValid(), asset_path + local_position, _ = sim_utils.resolve_prim_pose(asset_prim, ref_prim=env_root) + world_position, _ = sim_utils.resolve_prim_pose(asset_prim) + np.testing.assert_allclose(local_position, expected_local_position, rtol=0.0, atol=1.0e-6) + np.testing.assert_allclose( + world_position, + origin + np.asarray(expected_local_position), + rtol=0.0, + atol=1.0e-6, + ) + + for cup_name in ("SourceCup", "TargetCup"): + mesh_path = f"/World/envs/env_{env_id}/{cup_name}/geometry/mesh" + mesh = UsdGeom.Mesh.Get(stage, mesh_path) + assert mesh.GetPrim().IsValid(), mesh_path + assert UsdGeom.Imageable(mesh).ComputeVisibility() == UsdGeom.Tokens.inherited + assert not stage.GetPrimAtPath(f"/World/envs/env_{env_id}/Cup").IsValid() + assert not stage.GetPrimAtPath(f"/World/envs/env_{env_id}/SpillFloor").IsValid() + + source_positions = task.scene["source_cup"].data.root_link_pose_w.torch[:, :3].detach().cpu().numpy() + target_positions = task.scene["target_cup"].data.root_link_pose_w.torch[:, :3].detach().cpu().numpy() + np.testing.assert_allclose( + source_positions, + origins + np.asarray(task.cfg.cup_reset_pos), + rtol=0.0, + atol=1.0e-6, + ) + np.testing.assert_allclose( + target_positions, + origins + np.asarray(task.cfg.target_cup_reset_pos), + rtol=0.0, + atol=1.0e-6, + ) + body_labels = [str(label) for label in model.body_label] + for env_id in (0, 1): + for body_name in ("SourceCup", "TargetCup", "SpillFloor"): + expected_label = f"/World/envs/env_{env_id}/{body_name}" + assert body_labels.count(expected_label) == 1, expected_label + assert f"/World/envs/env_{env_id}/TargetCupRigid" not in body_labels + assert f"/World/envs/env_{env_id}/Cup" not in body_labels + + assert {world for _, world, _ in _shape_matches(model, "/Robot/")} == {0, 1} + _assert_shape_distribution(model, "/SourceCup/geometry/mesh", visible=True, expected_per_world={0: 1, 1: 1}) + _assert_shape_distribution(model, "/TargetCup/geometry/mesh", visible=True, expected_per_world={0: 1, 1: 1}) + _assert_shape_distribution( + model, "/SourceCup/geometry/grasp_proxy", visible=False, expected_per_world={0: 1, 1: 1} + ) + table_shapes = _shape_matches(model, "/Table/Collisions/Cube") + table_shape_distribution = { + (world, visible): sum( + match_world == world and match_visible is visible for _, match_world, match_visible in table_shapes + ) + for world in (0, 1) + for visible in (False, True) + } + assert table_shape_distribution == { + (0, False): 1, + (0, True): 1, + (1, False): 1, + (1, True): 1, + }, table_shapes + grasp_proxy_shape_ids = [ + shape_id for shape_id, _, _ in _shape_matches(model, "/SourceCup/geometry/grasp_proxy") + ] + for shape_id in grasp_proxy_shape_ids: + np.testing.assert_allclose( + model.shape_scale.numpy()[shape_id], + task.cfg.cup_grasp_box_half, + rtol=0.0, + atol=1.0e-6, + ) + _assert_shape_distribution(model, "/ParticleCollider", visible=False, expected_per_world={0: 2, 1: 2}) + _assert_shape_distribution(model, "/TargetCup/Collision", visible=False, expected_per_world={0: 1, 1: 1}) + _assert_shape_distribution(model, "/SpillFloor/Collision", visible=False, expected_per_world={0: 1, 1: 1}) + assert set(model.particle_world.numpy().tolist()) == {0, 1} + + point_paths = { + prim.GetPath().pathString + for prim in stage.Traverse() + if prim.IsA(UsdGeom.Points) and prim.GetPath().pathString.startswith("/World/Visuals/MPMParticles/") + } + assert len(point_paths) == 2 + assert {path.rsplit("/", 1)[-1] for path in point_paths} == {"env_0", "env_1"} + point_path_by_env = {path.rsplit("/", 1)[-1]: path for path in point_paths} + + cameras = [prim for prim in stage.Traverse() if prim.IsA(UsdGeom.Camera)] + assert cameras + for camera in cameras: + _assert_unpartitioned(camera, "omni:scenePartition") + + kit_visualizers = [visualizer for visualizer in task.sim.visualizers if isinstance(visualizer, KitVisualizer)] + newton_visualizers = [ + visualizer for visualizer in task.sim.visualizers if isinstance(visualizer, NewtonVisualizer) + ] + assert len(kit_visualizers) == 1 + assert len(newton_visualizers) == 1 + assert kit_visualizers[0].cfg.max_visible_envs is None + assert kit_visualizers[0].get_visualized_env_ids() is None + + newton_visualizer = newton_visualizers[0] + assert newton_visualizer.cfg.max_visible_envs is None + assert newton_visualizer.get_visualized_env_ids() is None + # NewtonVisualizer has no public accessor for the native viewer. + viewer = newton_visualizer._viewer + assert viewer is not None + assert model.world_count == 2 + assert viewer._visible_worlds is None + assert viewer._visible_worlds_mask is None + np.testing.assert_array_equal(viewer.world_offsets.numpy(), np.zeros((2, 3), dtype=np.float32)) + assert viewer.show_particles is True + + actions = torch.zeros((task.num_envs, task.action_manager.total_action_dim), device=task.device) + env.step(actions) + wp.synchronize_device(model.device) + + _assert_visual_descendants_follow_fabric_bodies(task) + + particle_world = model.particle_world.numpy() + particle_q = NewtonManager.get_state_0().particle_q.numpy() + for env_id in range(task.num_envs): + points = UsdGeom.Points.Get(stage, point_path_by_env[f"env_{env_id}"]).GetPointsAttr().Get() + actual = np.asarray(points, dtype=np.float32) + expected = particle_q[particle_world == env_id] + np.testing.assert_allclose(actual, expected, rtol=0.0, atol=1.0e-6) + finally: + if env is not None: + env.close() diff --git a/source/isaaclab_tasks/test/core/test_adaptive_reset_sampler.py b/source/isaaclab_tasks/test/core/test_adaptive_reset_sampler.py new file mode 100644 index 000000000000..7e2bfdff91b0 --- /dev/null +++ b/source/isaaclab_tasks/test/core/test_adaptive_reset_sampler.py @@ -0,0 +1,158 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Tests for the task-agnostic adaptive reset sampler.""" + +import pytest +import torch + +from isaaclab_tasks.utils.adaptive_reset_sampler import AdaptiveResetSampler, AdaptiveResetSamplerCfg + + +def test_config_rejects_invalid_values(): + """Configuration validation rejects values that break the probability model.""" + with pytest.raises(ValueError, match="target_success_rate"): + AdaptiveResetSamplerCfg(target_success_rate=1.0) + with pytest.raises(ValueError, match="history_capacity"): + AdaptiveResetSamplerCfg(history_capacity=0) + with pytest.raises(ValueError, match="probe_fraction"): + AdaptiveResetSamplerCfg(probe_fraction=1.0) + + +def test_sample_preserves_raw_ids_and_exact_overrides(): + """Sampling returns opaque raw IDs and honors overrides outside the frontier.""" + order = torch.tensor([50, 3, 99, 7, 42], dtype=torch.long) + cfg = AdaptiveResetSamplerCfg(initial_frontier_size=2, probe_size=1, probe_fraction=0.2) + sampler = AdaptiveResetSampler(order, cfg) + + forced = torch.tensor([-1, 42, 7, -1], dtype=torch.long) + samples = sampler.sample(4, forced) + + assert samples[1].item() == 42 + assert samples[2].item() == 7 + assert bool(torch.isin(samples[[0, 3]], order[:3]).all()) + with pytest.raises(ValueError, match="Unknown raw reset-row IDs"): + sampler.sample(1, torch.tensor([1234], dtype=torch.long)) + with pytest.raises(ValueError, match="-1"): + sampler.sample(1, torch.tensor([-2], dtype=torch.long)) + assert sampler.sample(0).shape == (0,) + + +def test_record_bounds_effective_outcome_history(): + """Recent evidence remains bounded while retaining the newest batch statistics.""" + sampler = AdaptiveResetSampler( + torch.tensor([10, 20], dtype=torch.long), + AdaptiveResetSamplerCfg(history_capacity=4, initial_frontier_size=2, probe_size=0), + ) + + sampler.record(torch.full((4,), 10), torch.ones(4, dtype=torch.bool)) + sampler.record(torch.full((4,), 10), torch.zeros(4, dtype=torch.bool)) + state = sampler.state_dict() + assert state["effective_attempts"][0].item() == pytest.approx(4.0) + assert state["effective_successes"][0].item() == pytest.approx(0.0) + + sampler.record(torch.full((2,), 10), torch.ones(2, dtype=torch.bool)) + state = sampler.state_dict() + assert state["effective_attempts"][0].item() == pytest.approx(4.0) + assert state["effective_successes"][0].item() == pytest.approx(2.0) + assert state["total_attempts"][0].item() == 10 + + +def test_probabilities_target_aggregate_success_with_replay(): + """The calibrated softmax approaches the requested aggregate success rate.""" + order = torch.arange(10, dtype=torch.long) + cfg = AdaptiveResetSamplerCfg( + target_success_rate=0.3, + temperature=0.02, + history_capacity=16, + prior_strength=0.01, + initial_frontier_size=10, + probe_size=0, + replay_fraction=0.1, + ) + sampler = AdaptiveResetSampler(order, cfg) + rows = order.repeat_interleave(16) + successes = (rows < 5).to(dtype=torch.bool) + sampler.record(rows, successes) + + predicted = torch.dot(sampler.sampling_probabilities, sampler.success_estimates) + assert predicted.item() == pytest.approx(0.3, abs=2.0e-3) + assert sampler.sampling_probabilities.min().item() >= 0.01 - 1.0e-6 + + +def test_probe_and_replay_probability_mass_are_retained(): + """Probe rows and active rows retain their configured probability floors.""" + cfg = AdaptiveResetSamplerCfg( + initial_frontier_size=4, + probe_size=2, + probe_fraction=0.2, + replay_fraction=0.1, + ) + sampler = AdaptiveResetSampler(torch.arange(8, dtype=torch.long), cfg) + probabilities = sampler.sampling_probabilities + + assert probabilities[4:6].sum().item() == pytest.approx(0.2) + assert probabilities[:4].min().item() >= 0.8 * 0.1 / 4 - 1.0e-6 + assert probabilities[6:].sum().item() == 0.0 + assert probabilities.sum().item() == pytest.approx(1.0) + + +def test_frontier_advances_monotonically_from_local_evidence(): + """Success near the frontier exposes harder rows and later failures never retract it.""" + cfg = AdaptiveResetSamplerCfg( + target_success_rate=0.5, + initial_frontier_size=2, + probe_size=2, + frontier_evidence=1.0, + ) + sampler = AdaptiveResetSampler(torch.arange(8, dtype=torch.long), cfg) + + sampler.record(torch.tensor([1, 2]), torch.tensor([True, True])) + promoted_size = sampler.frontier_size + assert promoted_size == 3 + + sampler.record(torch.tensor([2, 3, 4]), torch.tensor([False, False, False])) + assert sampler.frontier_size == promoted_size + + +def test_state_dict_round_trip_restores_sampling_state(): + """A checkpoint round trip preserves estimates, frontier, and sampling probabilities.""" + order = torch.tensor([7, 19, 2, 31, 4], dtype=torch.long) + cfg = AdaptiveResetSamplerCfg(initial_frontier_size=2, probe_size=2, frontier_evidence=1.0) + sampler = AdaptiveResetSampler(order, cfg) + sampler.record(torch.tensor([7, 19, 2]), torch.tensor([True, True, False])) + state = sampler.state_dict() + + restored = AdaptiveResetSampler(order, cfg) + restored.load_state_dict(state) + + assert restored.frontier_size == sampler.frontier_size + assert torch.equal(restored.success_estimates, sampler.success_estimates) + assert torch.equal(restored.sampling_probabilities, sampler.sampling_probabilities) + assert restored.metrics() == pytest.approx(sampler.metrics()) + + state["effective_attempts"].zero_() + assert bool(torch.any(sampler.state_dict()["effective_attempts"] > 0)) + + +def test_state_dict_rejects_a_different_reset_cache(): + """Checkpoint restore cannot silently attach outcome history to different rows.""" + cfg = AdaptiveResetSamplerCfg(initial_frontier_size=2) + state = AdaptiveResetSampler(torch.tensor([1, 2, 3]), cfg).state_dict() + sampler = AdaptiveResetSampler(torch.tensor([1, 3, 2]), cfg) + + with pytest.raises(ValueError, match="difficulty_order"): + sampler.load_state_dict(state) + + +@pytest.mark.parametrize("field", ("effective_successes", "effective_attempts", "frontier_credit")) +def test_state_dict_rejects_nonfinite_sampling_state(field): + """Non-finite checkpoint counters cannot poison restored sampling probabilities.""" + sampler = AdaptiveResetSampler(torch.tensor([1, 2, 3]), AdaptiveResetSamplerCfg(initial_frontier_size=2)) + state = sampler.state_dict() + state[field].reshape(-1)[0] = torch.nan + + with pytest.raises(ValueError, match="invalid|outside"): + sampler.load_state_dict(state) diff --git a/source/isaaclab_tasks/test/utils/test_reset_dataset.py b/source/isaaclab_tasks/test/utils/test_reset_dataset.py new file mode 100644 index 000000000000..1c3f6f9157a7 --- /dev/null +++ b/source/isaaclab_tasks/test/utils/test_reset_dataset.py @@ -0,0 +1,214 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Tests for task-agnostic reset-dataset utilities.""" + +from copy import deepcopy + +import pytest +import torch + +from isaaclab_tasks.utils.reset_dataset import ( + reset_dataset_collect_batches, + reset_dataset_content_digest, + reset_dataset_digest, + reset_dataset_save_atomic, + reset_dataset_validate_header, +) + +_FORMAT = "example_reset_states" +_SCHEMA_VERSION = 3 +_CONTRACT = {"generator": {"seed": 7, "workspace": (-1.0, 1.0)}} + + +def _payload() -> dict: + payload = { + "format": _FORMAT, + "schema_version": _SCHEMA_VERSION, + "contract_sha256": reset_dataset_digest(_CONTRACT), + "metadata": {"state_count": 2}, + "states": { + "position": torch.tensor(((0.0, 1.0), (2.0, 3.0)), dtype=torch.float32), + "category": torch.tensor((0, 1), dtype=torch.int64), + }, + } + payload["content_sha256"] = reset_dataset_content_digest(payload) + return payload + + +def _validate(payload: dict) -> None: + reset_dataset_validate_header( + payload, + expected_format=_FORMAT, + expected_schema_version=_SCHEMA_VERSION, + expected_contract=_CONTRACT, + ) + + +def test_reset_dataset_digest_is_stable_over_mapping_order_and_tensor_device_metadata(): + first = { + "nested": [{"enabled": True, "gain": 0.25}, None], + "tensor": torch.arange(6, dtype=torch.float32).reshape(2, 3), + } + second = { + "tensor": first["tensor"].clone(), + "nested": [{"gain": 0.25, "enabled": True}, None], + } + + assert reset_dataset_digest(first) == reset_dataset_digest(second) + + +@pytest.mark.parametrize( + "changed", + [ + {"value": True}, + {"value": [1, 2]}, + {"value": torch.tensor((1, 2), dtype=torch.int32)}, + {"value": torch.tensor(((1, 2),), dtype=torch.int64)}, + ], +) +def test_reset_dataset_digest_preserves_type_shape_and_dtype_boundaries(changed): + baseline = {"value": (1, 2)} + + assert reset_dataset_digest(changed) != reset_dataset_digest(baseline) + + +def test_reset_dataset_digest_supports_scalar_and_bfloat16_tensors(): + payload = { + "scalar": torch.tensor(2.0, dtype=torch.float64), + "bfloat16": torch.tensor((1.0, 2.0), dtype=torch.bfloat16), + } + + assert len(reset_dataset_digest(payload)) == 64 + + +def test_reset_dataset_content_digest_excludes_only_its_own_top_level_field(): + payload = _payload() + original = reset_dataset_content_digest(payload) + payload["content_sha256"] = "not-the-content-digest" + + assert reset_dataset_content_digest(payload) == original + + payload["metadata"]["content_sha256"] = "nested-data-remains-content" + assert reset_dataset_content_digest(payload) != original + + +def test_reset_dataset_validate_header_returns_common_mappings(): + payload = _payload() + + metadata, states = reset_dataset_validate_header( + payload, + expected_format=_FORMAT, + expected_schema_version=_SCHEMA_VERSION, + expected_contract=_CONTRACT, + ) + + assert metadata is payload["metadata"] + assert states is payload["states"] + + +@pytest.mark.parametrize( + ("mutation", "message"), + [ + (lambda payload: payload.update(format="other"), "format"), + (lambda payload: payload.update(schema_version=4), "schema version"), + (lambda payload: payload.update(metadata=[]), "metadata"), + (lambda payload: payload.update(states=[]), "states"), + (lambda payload: payload.update(contract_sha256="wrong"), "contract digest"), + (lambda payload: payload["states"]["position"].add_(1.0), "content digest"), + ], +) +def test_reset_dataset_validate_header_rejects_invalid_envelopes(mutation, message): + payload = _payload() + mutation(payload) + + with pytest.raises((TypeError, ValueError), match=message): + _validate(payload) + + +def test_reset_dataset_save_atomic_validates_and_round_trips(tmp_path): + payload = _payload() + calls = [] + + def validator(candidate): + calls.append(candidate) + _validate(candidate) + + output_path = tmp_path / "nested" / "states.pt" + saved_path = reset_dataset_save_atomic(payload, output_path, validator=validator) + restored = torch.load(saved_path, map_location="cpu", weights_only=False) + + assert calls == [payload] + assert saved_path == output_path.resolve() + assert reset_dataset_content_digest(restored) == payload["content_sha256"] + assert not list(output_path.parent.glob(".*.tmp")) + + +def test_reset_dataset_save_atomic_does_not_replace_destination_when_validation_fails(tmp_path): + destination = tmp_path / "states.pt" + destination.write_bytes(b"existing") + + with pytest.raises(ValueError, match="invalid dataset"): + reset_dataset_save_atomic( + _payload(), + destination, + validator=lambda _payload: (_ for _ in ()).throw(ValueError("invalid dataset")), + ) + + assert destination.read_bytes() == b"existing" + + +def test_reset_dataset_collect_batches_trims_the_final_accepted_batch(): + evaluated_ranges = [] + + def evaluate(candidate_ids: range) -> list[int]: + evaluated_ranges.append(candidate_ids) + return [candidate_id for candidate_id in candidate_ids if candidate_id % 2 == 0] + + batches, evaluated_count = reset_dataset_collect_batches( + 5, + batch_size=4, + max_candidate_count=12, + evaluate_batch=evaluate, + batch_count=len, + batch_slice=lambda batch, count: batch[:count], + ) + + assert batches == [[0, 2], [4, 6], [8]] + assert evaluated_ranges == [range(0, 4), range(4, 8), range(8, 12)] + assert evaluated_count == 12 + + +def test_reset_dataset_collect_batches_reports_candidate_budget_exhaustion(): + with pytest.raises(RuntimeError, match="accepted 2/3.*6 candidates"): + reset_dataset_collect_batches( + 3, + batch_size=2, + max_candidate_count=6, + evaluate_batch=lambda candidate_ids: [value for value in candidate_ids if value % 5 == 0], + batch_count=len, + batch_slice=lambda batch, count: batch[:count], + ) + + +def test_reset_dataset_collect_batches_rejects_impossible_callback_counts(): + with pytest.raises(ValueError, match="outside the candidate range"): + reset_dataset_collect_batches( + 1, + batch_size=2, + max_candidate_count=2, + evaluate_batch=lambda _candidate_ids: "invalid", + batch_count=lambda _batch: 3, + batch_slice=lambda batch, _count: batch, + ) + + +def test_reset_dataset_content_tampering_does_not_mutate_fixture(): + payload = _payload() + modified = deepcopy(payload) + modified["states"]["position"][0, 0] = 9.0 + + assert payload["states"]["position"][0, 0] == 0.0 + assert reset_dataset_content_digest(modified) != payload["content_sha256"] diff --git a/source/isaaclab_visualizers/changelog.d/pin-newton-c7ae7c7.rst b/source/isaaclab_visualizers/changelog.d/pin-newton-c7ae7c7.rst new file mode 100644 index 000000000000..2f3e20289546 --- /dev/null +++ b/source/isaaclab_visualizers/changelog.d/pin-newton-c7ae7c7.rst @@ -0,0 +1,6 @@ +Changed +^^^^^^^ + +* Changed the ``newton[sim]`` dependency pin of the visualizer extras to Newton + commit ``b691d94db1a03c514de03bfdaf27cb9136fc766f`` and required + ``newton-usd-schemas>=0.4.0`` for Newton's USD parsing. From f43d04da7b2545fb8f3dae78dcd8993f7ab23528 Mon Sep 17 00:00:00 2001 From: HuiDong Chen Date: Wed, 15 Jul 2026 16:24:43 +0800 Subject: [PATCH 02/14] draft of rendering correctness validation for franka pour --- .../test/core/test_rendering_franka_pour.py | 37 +++++ .../test_rendering_franka_pour_kitless.py | 33 ++++ .../test/rendering_test_utils.py | 142 ++++++++++++++++++ 3 files changed, 212 insertions(+) create mode 100644 source/isaaclab_tasks/test/core/test_rendering_franka_pour.py create mode 100644 source/isaaclab_tasks/test/core/test_rendering_franka_pour_kitless.py diff --git a/source/isaaclab_tasks/test/core/test_rendering_franka_pour.py b/source/isaaclab_tasks/test/core/test_rendering_franka_pour.py new file mode 100644 index 000000000000..ee2351e68897 --- /dev/null +++ b/source/isaaclab_tasks/test/core/test_rendering_franka_pour.py @@ -0,0 +1,37 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Rendering correctness tests for test-local Franka pour camera setup.""" + +# Launch Isaac Sim Simulator first for kit-based combinations. +from isaaclab.app import AppLauncher + +app_launcher = AppLauncher(headless=True, enable_cameras=True) +simulation_app = app_launcher.app + +from pathlib import Path # noqa: E402 + +import pytest # noqa: E402 +from rendering_test_utils import ( # noqa: E402 + PHYSICS_RENDERER_AOV_COMBINATIONS, + make_attach_comparison_properties_fixture, + make_determinism_fixture, + make_generate_html_report_fixture, + rendering_test_franka_pour, +) + +pytestmark = [pytest.mark.isaacsim_ci, pytest.mark.newton_ci] + +_COMPARISON_SCORES: list[dict] = [] + +_determinism_fixture = make_determinism_fixture() +_generate_html_report_fixture = make_generate_html_report_fixture(_COMPARISON_SCORES, Path(__file__).stem + ".html") +_attach_comparison_properties_fixture = make_attach_comparison_properties_fixture(_COMPARISON_SCORES) + + +@pytest.mark.parametrize("physics_backend,renderer,data_type", PHYSICS_RENDERER_AOV_COMBINATIONS) +def test_rendering_franka_pour(physics_backend, renderer, data_type): + """Test Franka pour rendering correctness.""" + rendering_test_franka_pour(physics_backend, renderer, data_type, _COMPARISON_SCORES) diff --git a/source/isaaclab_tasks/test/core/test_rendering_franka_pour_kitless.py b/source/isaaclab_tasks/test/core/test_rendering_franka_pour_kitless.py new file mode 100644 index 000000000000..e5b39548597a --- /dev/null +++ b/source/isaaclab_tasks/test/core/test_rendering_franka_pour_kitless.py @@ -0,0 +1,33 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Kit-less rendering correctness tests for test-local Franka pour camera setup.""" + +from pathlib import Path + +import pytest +from rendering_test_utils import ( + KITLESS_PHYSICS_RENDERER_AOV_COMBINATIONS, + make_attach_comparison_properties_fixture, + make_determinism_fixture, + make_generate_html_report_fixture, + make_require_ovlibs_install_fixture, + rendering_test_franka_pour, +) + +pytestmark = [pytest.mark.isaacsim_ci, pytest.mark.newton_ci] + +_COMPARISON_SCORES: list[dict] = [] + +_determinism_fixture = make_determinism_fixture() +_generate_html_report_fixture = make_generate_html_report_fixture(_COMPARISON_SCORES, Path(__file__).stem + ".html") +_attach_comparison_properties_fixture = make_attach_comparison_properties_fixture(_COMPARISON_SCORES) +_require_ovlibs_install_fixture = make_require_ovlibs_install_fixture() + + +@pytest.mark.parametrize("physics_backend,renderer,data_type", KITLESS_PHYSICS_RENDERER_AOV_COMBINATIONS) +def test_rendering_franka_pour_kitless(physics_backend, renderer, data_type): + """Camera output must match golden images for the Franka pour test setup.""" + rendering_test_franka_pour(physics_backend, renderer, data_type, _COMPARISON_SCORES) diff --git a/source/isaaclab_tasks/test/rendering_test_utils.py b/source/isaaclab_tasks/test/rendering_test_utils.py index cccaa7eb359f..e50bf9d724af 100644 --- a/source/isaaclab_tasks/test/rendering_test_utils.py +++ b/source/isaaclab_tasks/test/rendering_test_utils.py @@ -49,6 +49,8 @@ # Aliasing artifacts of shadow on the table. "franka_cloth": 8.0, "franka_soft": 8.0, + # Aliasing artifacts on the table and MPM particle noise. + "franka_pour": 8.0, # Shadow-hand renderings (incl. ``Isaac-Reorient-Cube-Shadow-Camera-Direct``) show up to # ~3.28 % per-pixel diff from anti-aliasing noise along the many finger/cube edges. 5.0 gives # headroom above that without masking real regressions, which the SSIM gate still catches. @@ -1908,3 +1910,143 @@ def rendering_test_franka_soft( # This invokes camera sensor and renderer cleanup explicitly before pytest teardown, otherwise OV # native code could probably complain about leaks and trigger segmentation fault. env = None + + +def _make_franka_pour_camera_env_cfg(data_type: str): + """Create a test-local Franka pour camera env cfg without exposing a production task.""" + import isaaclab.sim as sim_utils + from isaaclab.envs import mdp as env_mdp + from isaaclab.managers import ObservationGroupCfg as ObsGroup + from isaaclab.managers import ObservationTermCfg as ObsTerm + from isaaclab.managers import SceneEntityCfg + from isaaclab.sensors import CameraCfg + from isaaclab.utils.configclass import configclass + + from isaaclab_tasks.contrib.franka_pour.pour_env_cfg import FrankaPourEnvCfg, PourSceneCfg + from isaaclab_tasks.utils.presets import MultiBackendRendererCfg + + @configclass + class TestFrankaPourCameraSceneCfg(PourSceneCfg): + """Franka pour scene with a test-only camera sensor.""" + + tiled_camera: CameraCfg = CameraCfg( + prim_path="/World/envs/env_.*/Camera", + offset=CameraCfg.OffsetCfg( + pos=(0.85, -0.55, 0.42), + rot=(0.5080, 0.2114, 0.318, 0.7720), + convention="opengl", + ), + data_types=[data_type], + spawn=sim_utils.PinholeCameraCfg(clipping_range=(0.01, 3.0)), + width=128, + height=128, + renderer_cfg=MultiBackendRendererCfg(), + ) + + @configclass + class TestFrankaPourCameraObservationsCfg: + """Image-only observations for the local rendering test env.""" + + @configclass + class PolicyCfg(ObsGroup): + image = ObsTerm( + func=env_mdp.image, + params={"sensor_cfg": SceneEntityCfg("tiled_camera"), "data_type": data_type, "permute": True}, + ) + + def __post_init__(self) -> None: + self.enable_corruption = False + self.concatenate_terms = True + + policy: ObsGroup = PolicyCfg() + + @configclass + class TestFrankaPourCameraEnvCfg(FrankaPourEnvCfg): + """Test-only camera variant of ``Isaac-Pour-Franka-v0``.""" + + scene: TestFrankaPourCameraSceneCfg = TestFrankaPourCameraSceneCfg( + num_envs=4, env_spacing=2.5, replicate_physics=True + ) + observations: TestFrankaPourCameraObservationsCfg = TestFrankaPourCameraObservationsCfg() + + def __post_init__(self) -> None: + super().__post_init__() + self.seed = 42 + full_stage = self.curriculum_stage_names.index("full") + self.curriculum_start_stage = full_stage + self.curriculum_freeze = True + self.decimation = 1 + self.physics_substeps = 1 + self.mpm_iterations = 2 + self.use_cuda_graph = False + self.sim.render_interval = 1 + self.sim.device = "cuda:0" + + return TestFrankaPourCameraEnvCfg() + + +def rendering_test_franka_pour( + physics_backend: str, + renderer: str, + data_type: str, + comparison_scores: list[dict], +) -> None: + if physics_backend in ("physx", "ovphysx"): + pytest.skip("FrankaPour env cfg is Newton-only.") + + _skip_if_newton_motion_vectors(physics_backend, data_type) + + import warp as wp + + if not wp.is_cuda_available(): + pytest.skip("FrankaPour rendering tests require a CUDA device.") + + from isaaclab_tasks.contrib.franka_pour.pour_env import FrankaPourEnv + + import isaaclab.sim as sim_utils + + env_cfg = _make_franka_pour_camera_env_cfg(data_type) + env_cfg = _apply_overrides_to_env_cfg( + env_cfg, [f"presets={_physics_preset_name(physics_backend)},{renderer}"] + ) + + env_cfg.scene.num_envs = 4 + + if renderer == "ovrtx_renderer": + _redirect_ovrtx_renderer_log_to_stdout(env_cfg) + + test_name = "franka_pour" + env = None + + try: + # Kit-based RTX tests need a fresh stage between parametrized cases. Kitless OVRTX uses a + # different USD bootstrap path and can segfault if we force-create a stage here. + if renderer == "isaacsim_rtx_renderer": + sim_utils.create_new_stage() + env = FrankaPourEnv(env_cfg) + env.sim._app_control_on_stop_handle = None + + _maybe_disable_instancing_for_current_stage(physics_backend, renderer, data_type) + + maybe_save_stage(test_name, physics_backend, renderer, data_type) + + zero_actions = torch.zeros(env.num_envs, env.action_manager.total_action_dim, device=env.device) + maybe_step_env_for_motion(env, data_type, num_steps=2, action_value=0.0) + if data_type != "motion_vectors": + env.step(zero_actions) + + validate_camera_outputs( + test_name, + physics_backend, + renderer, + env.scene.sensors["tiled_camera"].data.output, + max_different_pixels_percentage=MAX_DIFFERENT_PIXELS_PERCENTAGE_BY_ENV_NAME[test_name], + comparison_scores=comparison_scores, + ) + finally: + if env is not None: + env.close() + + # This invokes camera sensor and renderer cleanup explicitly before pytest teardown, otherwise OV + # native code could probably complain about leaks and trigger segmentation fault. + env = None From cfadf7583402c280d0091c0ae6fd00a99dafd973 Mon Sep 17 00:00:00 2001 From: HuiDong Chen Date: Wed, 15 Jul 2026 21:46:29 +0800 Subject: [PATCH 03/14] Changes in MPM particle spawner - to be reviewed by Maximilian --- .../mpm-visualization-backend-agnostic.rst | 13 ++++++ .../assets/mpm_object/mpm_object.py | 42 ++++++++++++------- .../sim/spawners/mpm/mpm_cfg.py | 4 +- .../sim/spawners/mpm/visualization.py | 25 ++++++++--- .../test/assets/test_mpm_object.py | 6 +-- 5 files changed, 63 insertions(+), 27 deletions(-) create mode 100644 source/isaaclab_newton/changelog.d/mpm-visualization-backend-agnostic.rst diff --git a/source/isaaclab_newton/changelog.d/mpm-visualization-backend-agnostic.rst b/source/isaaclab_newton/changelog.d/mpm-visualization-backend-agnostic.rst new file mode 100644 index 000000000000..b5bca527c013 --- /dev/null +++ b/source/isaaclab_newton/changelog.d/mpm-visualization-backend-agnostic.rst @@ -0,0 +1,13 @@ +Fixed +^^^^^ + +* Fixed :class:`~isaaclab_newton.assets.MPMObject` not creating particle visualization + prims outside the Kit viewport visualizer, which left MPM particles invisible to Kit RTX + cameras. The spawner now creates ``UsdGeom.Points`` prims whenever the object is configured + visible, without inspecting the active render backend. +* Fixed MPM particle visualization prims being authored outside the environment hierarchy, + which left them without an ``omni:scenePartition`` and therefore invisible to tiled + renderers (Kit RTX cameras, OVRTX). The ``UsdGeom.Points`` prims are now authored as a + ``Particles`` child of the asset prim (``/World/envs/env_{idx}//Particles``) with the + reset-xform-stack flag set, so they inherit the environment's scene partition while keeping + their world-frame positions. diff --git a/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/mpm_object.py b/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/mpm_object.py index 485b631de9f4..1d896f0337c7 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/mpm_object.py +++ b/source/isaaclab_newton/isaaclab_newton/assets/mpm_object/mpm_object.py @@ -6,6 +6,7 @@ from __future__ import annotations import logging +import re from collections.abc import Sequence from dataclasses import dataclass, field from typing import TYPE_CHECKING @@ -341,14 +342,19 @@ def _create_buffers(self): ) self._data.default_nodal_state_w = ProxyArray(default_state) self._data.default_particle_state_w = self._data.default_nodal_state_w - self._create_kit_points() - - def _create_kit_points(self) -> None: - """Create Kit-visible ``UsdGeom.Points`` prims for the particles when the Kit visualizer is active.""" - from isaaclab.sim import SimulationContext # noqa: PLC0415 - - sim = SimulationContext.instance() - if sim is None or "kit" not in sim.resolve_visualizer_types() or not self.cfg.spawn.visible: + self._create_particle_visualization() + + def _create_particle_visualization(self) -> None: + """Create ``UsdGeom.Points`` prims mirroring the particle state for USD-based renderers. + + The prims are authored inside the environment hierarchy, as a ``Particles`` child of the + asset prim (e.g. ``/World/envs/env_{idx}/Media/Particles``), so each inherits its + environment's ``omni:scenePartition`` and is drawn in the correct tile by partition-aware + renderers. The prims are created whenever the object is configured visible; the spawner does + not inspect the active render backend. Per-frame position updates are handled by + :meth:`~isaaclab_newton.physics.NewtonManager.sync_particles_to_usd`. + """ + if not self.cfg.spawn.visible: return first_offset = self._recorded_particle_offsets[0] @@ -357,10 +363,11 @@ def _create_kit_points(self) -> None: .particle_radius[first_offset : first_offset + self._particles_per_object] .numpy() ) - base_path = _create_kit_visualization_path(self.cfg.prim_path) + positions = self.data.particle_pos_w.warp.numpy() + prim_paths = _particle_visualization_paths(self.cfg.prim_path, positions.shape[0]) prim_paths = create_mpm_particle_visualization( - prim_path=base_path, - positions=self.data.particle_pos_w.warp.numpy(), + prim_paths=prim_paths, + positions=positions, widths=2.0 * radii, color=self.cfg.spawn.visual_color, ) @@ -371,7 +378,7 @@ def _create_kit_points(self) -> None: particle_count=self._particles_per_object, sync_frequency=self.cfg.spawn.visual_update_frequency, ) - logger.info("Kit MPM particle visualization initialized at: %s", base_path) + logger.info("MPM particle visualization initialized under: %s", self.cfg.prim_path) def _resolve_env_ids(self, env_ids): if env_ids is None or (isinstance(env_ids, slice) and env_ids == slice(None)): @@ -439,6 +446,11 @@ def _compose_env_asset_pose( return (float(pos[0]), float(pos[1]), float(pos[2])), (float(rot[0]), float(rot[1]), float(rot[2]), float(rot[3])) -def _create_kit_visualization_path(prim_path: str) -> str: - sanitized = "".join(char if char.isalnum() else "_" for char in prim_path.strip("/")) - return f"/World/Visuals/MPMParticles/{sanitized or 'Object'}" +def _particle_visualization_paths(prim_path: str, num_envs: int) -> list[str]: + """Resolve one in-hierarchy ``UsdGeom.Points`` path per environment. + + The particle cloud is authored as a ``Particles`` child of the asset prim, resolving the + cloned environment wildcard (``env_.*``) to each concrete index, so the prim lives under + ``/World/envs/env_{idx}//Particles`` and inherits the environment's scene partition. + """ + return [re.sub(r"(?<=[Ee]nv_)\.\*", str(env_idx), prim_path) + "/Particles" for env_idx in range(num_envs)] diff --git a/source/isaaclab_newton/isaaclab_newton/sim/spawners/mpm/mpm_cfg.py b/source/isaaclab_newton/isaaclab_newton/sim/spawners/mpm/mpm_cfg.py index 4bbeb8fd8c39..ade529177cde 100644 --- a/source/isaaclab_newton/isaaclab_newton/sim/spawners/mpm/mpm_cfg.py +++ b/source/isaaclab_newton/isaaclab_newton/sim/spawners/mpm/mpm_cfg.py @@ -67,10 +67,10 @@ class MPMParticleSpawnerCfg(SpawnerCfg): """Material values applied to generated particles.""" visual_color: Sequence[float] = (0.7, 0.6, 0.4) - """Display color for Kit particle visualization.""" + """Display color for particle visualization.""" visual_update_frequency: int = 1 - """Kit particle visualization update frequency in render frames.""" + """Particle visualization update frequency in render frames.""" @configclass diff --git a/source/isaaclab_newton/isaaclab_newton/sim/spawners/mpm/visualization.py b/source/isaaclab_newton/isaaclab_newton/sim/spawners/mpm/visualization.py index b63bda8e31d2..6a4d78ae3fe7 100644 --- a/source/isaaclab_newton/isaaclab_newton/sim/spawners/mpm/visualization.py +++ b/source/isaaclab_newton/isaaclab_newton/sim/spawners/mpm/visualization.py @@ -11,12 +11,18 @@ def create_mpm_particle_visualization( - prim_path: str, + prim_paths: Sequence[str], positions: np.ndarray, widths: np.ndarray, color: Sequence[float], ) -> list[str]: - """Create one ``UsdGeom.Points`` prim per environment for Kit MPM particle rendering. + """Create one ``UsdGeom.Points`` prim per environment for USD-based MPM particle rendering. + + A ``Points`` prim is authored at each path in ``prim_paths``. The paths live inside the + environment hierarchy (a child of the asset prim, e.g. ``/World/envs/env_{idx}/Media/Particles``) + so the prim inherits its environment's ``omni:scenePartition`` and is drawn in the correct tile + by partition-aware renderers (Kit RTX cameras, OVRTX). The reset-xform-stack flag is set so the + world-frame positions are not offset a second time by the inherited environment/asset transform. The created prims are static USD containers: per-frame position updates are handled by :meth:`isaaclab_newton.physics.NewtonManager.sync_particles_to_usd` @@ -24,8 +30,7 @@ def create_mpm_particle_visualization( :meth:`isaaclab_newton.physics.NewtonManager.register_particle_visual_prim`. Args: - prim_path: Base prim path; one ``Points`` prim is created per environment - at ``{prim_path}/env_{idx}``. + prim_paths: Absolute ``Points`` prim paths, one per environment in environment order. positions: Initial world-frame particle positions [m], shape ``(num_envs, particles_per_env, 3)``. widths: Particle display widths (diameters) [m], one per particle. @@ -38,8 +43,13 @@ def create_mpm_particle_visualization( import isaaclab.sim as sim_utils + if len(prim_paths) != positions.shape[0]: + raise ValueError( + f"Expected one prim path per environment: got {len(prim_paths)} paths for" + f" {positions.shape[0]} environments." + ) + stage = sim_utils.get_current_stage() - prim_paths = [f"{prim_path}/env_{env_idx}" for env_idx in range(positions.shape[0])] points_prims = [UsdGeom.Points.Define(stage, path) for path in prim_paths] widths_vt = Vt.FloatArray.FromNumpy(np.ascontiguousarray(widths, dtype=np.float32)) @@ -49,5 +59,8 @@ def create_mpm_particle_visualization( points.GetPointsAttr().Set(Vt.Vec3fArray.FromNumpy(positions[env_idx])) points.CreateWidthsAttr(widths_vt) points.CreateDisplayColorAttr(color_vt) + # Positions are world-frame; drop the inherited env/asset transform so they are + # not offset a second time when the prim lives under ``/World/envs/env_{idx}``. + points.SetResetXformStack(True) - return prim_paths + return list(prim_paths) diff --git a/source/isaaclab_newton/test/assets/test_mpm_object.py b/source/isaaclab_newton/test/assets/test_mpm_object.py index 646cc6e8dfc5..5bbc0bf5f4fa 100644 --- a/source/isaaclab_newton/test/assets/test_mpm_object.py +++ b/source/isaaclab_newton/test/assets/test_mpm_object.py @@ -181,7 +181,7 @@ class MPMSceneCfg(InteractiveSceneCfg): np.testing.assert_allclose(body_q, root_pose.detach().cpu().numpy()[0], rtol=1.0e-5, atol=1.0e-6) -def test_mpm_object_creates_kit_points_when_kit_visualizer_requested(monkeypatch): +def test_mpm_object_creates_particle_visualization_prims(): @configclass class MPMSceneCfg(InteractiveSceneCfg): media = MPMObjectCfg( @@ -202,7 +202,6 @@ class MPMSceneCfg(InteractiveSceneCfg): ) with build_simulation_context(sim_cfg=sim_cfg) as sim: - monkeypatch.setattr(sim, "resolve_visualizer_types", lambda: ["kit"]) scene = InteractiveScene(MPMSceneCfg(num_envs=2, env_spacing=1.0)) sim.reset() @@ -225,7 +224,7 @@ class MPMSceneCfg(InteractiveSceneCfg): assert tuple(points.GetDisplayColorAttr().Get()[0]) == pytest.approx((0.1, 0.2, 0.3)) -def test_mpm_kit_points_follow_particle_state(monkeypatch): +def test_mpm_particle_visualization_follows_particle_state(): @configclass class MPMSceneCfg(InteractiveSceneCfg): media = MPMObjectCfg( @@ -246,7 +245,6 @@ class MPMSceneCfg(InteractiveSceneCfg): ) with build_simulation_context(sim_cfg=sim_cfg) as sim: - monkeypatch.setattr(sim, "resolve_visualizer_types", lambda: ["kit"]) scene = InteractiveScene(MPMSceneCfg(num_envs=1, env_spacing=0.0)) sim.reset() From abf4136b74abcfe81c9067e4ac77fc1f06dc9f6d Mon Sep 17 00:00:00 2001 From: HuiDong Chen Date: Wed, 15 Jul 2026 21:55:51 +0800 Subject: [PATCH 04/14] ovrtx changes --- .../changelog.d/ovrtx-mpm-particle-points.rst | 9 ++ .../test/rendering_test_utils.py | 145 +++++++++++++++--- 2 files changed, 136 insertions(+), 18 deletions(-) create mode 100644 source/isaaclab_ov/changelog.d/ovrtx-mpm-particle-points.rst diff --git a/source/isaaclab_ov/changelog.d/ovrtx-mpm-particle-points.rst b/source/isaaclab_ov/changelog.d/ovrtx-mpm-particle-points.rst new file mode 100644 index 000000000000..64e704fe268b --- /dev/null +++ b/source/isaaclab_ov/changelog.d/ovrtx-mpm-particle-points.rst @@ -0,0 +1,9 @@ +Fixed +^^^^^ + +* Fixed OVRTX rendering to stream Newton MPM particle positions into registered + ``UsdGeom.Points`` prims, so kitless cameras can visualize MPM particle clouds. +* Fixed MPM particle clouds silently disappearing in OVRTX because the per-frame + ``points`` update was written from a GPU buffer, which OVRTX does not use to refresh + ``UsdGeom.Points`` sphere geometry. Particle positions are now written from host memory + (deformable meshes keep the zero-copy GPU path). diff --git a/source/isaaclab_tasks/test/rendering_test_utils.py b/source/isaaclab_tasks/test/rendering_test_utils.py index e50bf9d724af..a3e58ca8cc32 100644 --- a/source/isaaclab_tasks/test/rendering_test_utils.py +++ b/source/isaaclab_tasks/test/rendering_test_utils.py @@ -588,6 +588,108 @@ def maybe_save_stage( os.unlink(stage_path) +def _camera_output_to_pil_image(tensor: torch.Tensor, data_type: str) -> Image.Image: + """Convert one camera output tensor to a display-ready PIL image.""" + condition = torch.logical_or(torch.isinf(tensor), torch.isnan(tensor)) + corrected = torch.where(condition, torch.zeros_like(tensor), tensor) + normalized = normalize_camera_output_for_display(corrected, data_type) + grid = make_camera_output_grid(normalized) + ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy() + return Image.fromarray(ndarr) + + +def _pour_frame_save_dir() -> str | None: + """Return the pour-frame output directory when frame saving is enabled.""" + out_dir = os.environ.get("ISAAC_LAB_SAVE_STAGES") + if not out_dir: + return None + return os.path.join(out_dir, "pour_frames") + + +def _pour_frame_count() -> int: + """Return the number of zero-action pour steps to simulate.""" + raw = os.environ.get("ISAAC_LAB_POUR_FRAMES", "120") + try: + count = int(raw) + except ValueError as error: + raise ValueError(f"ISAAC_LAB_POUR_FRAMES must be an integer, got {raw!r}.") from error + if count < 0: + raise ValueError(f"ISAAC_LAB_POUR_FRAMES must be nonnegative, got {count}.") + return count + + +def maybe_save_camera_frames( + test_name: str, + physics_backend: str, + renderer: str, + data_type: str, + camera_outputs: dict[str, ProxyArray], + frame_index: int, + *, + frame_images: list[Image.Image] | None = None, +) -> list[Image.Image]: + """If ``ISAAC_LAB_SAVE_STAGES`` is set, save the current camera frame and accumulate GIF frames. + + Args: + test_name: Rendering test label used in output filenames. + physics_backend: Physics backend label. + renderer: Renderer nickname. + data_type: Camera data type under test. + camera_outputs: Camera sensor output dictionary. + frame_index: Zero-based frame index in the pour sequence. + frame_images: Optional accumulator for animated GIF assembly. + + Returns: + The updated frame-image accumulator (empty when saving is disabled). + """ + out_dir = _pour_frame_save_dir() + if out_dir is None or data_type not in camera_outputs: + return frame_images or [] + + output = camera_outputs[data_type] + tensor = output if isinstance(output, torch.Tensor) else output.torch + image = _camera_output_to_pil_image(tensor, data_type) + + os.makedirs(out_dir, exist_ok=True) + safe_test_name = test_name.replace("/", "_") + prefix = f"{safe_test_name}-{physics_backend}-{renderer}-{data_type}" + frame_path = os.path.join(out_dir, f"{prefix}-frame{frame_index:03d}.png") + image.save(frame_path, format="PNG") + + accumulated = list(frame_images or []) + accumulated.append(image) + return accumulated + + +def maybe_write_pour_frame_gif( + test_name: str, + physics_backend: str, + renderer: str, + data_type: str, + frame_images: list[Image.Image], +) -> None: + """Write an animated GIF from accumulated pour frames when saving is enabled.""" + if not frame_images: + return + + out_dir = _pour_frame_save_dir() + if out_dir is None: + return + + os.makedirs(out_dir, exist_ok=True) + safe_test_name = test_name.replace("/", "_") + prefix = f"{safe_test_name}-{physics_backend}-{renderer}-{data_type}" + gif_path = os.path.join(out_dir, f"{prefix}.gif") + frame_images[0].save( + gif_path, + save_all=True, + append_images=frame_images[1:], + duration=100, + loop=0, + ) + print(f"[ISAAC_LAB_SAVE_STAGES] wrote {gif_path} ({len(frame_images)} frames)") + + def _apply_overrides_to_env_cfg(env_cfg: Any, override_args: list[str]) -> Any: """Apply override args to env_cfg using parse_overrides and apply_overrides.""" from isaaclab_tasks.utils.hydra import apply_overrides, collect_presets, parse_overrides @@ -1092,10 +1194,7 @@ def validate_camera_outputs( failed_data_types[data_type] = f"Camera output '{data_type}' has no non-zero pixels." continue - normalized = normalize_camera_output_for_display(corrected, data_type) - grid = make_camera_output_grid(normalized) - ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy() - result_image = Image.fromarray(ndarr) + result_image = _camera_output_to_pil_image(corrected, data_type) golden_path = os.path.join(golden_image_dir, f"{physics_backend}-{renderer}-{data_type}.png") if not os.path.exists(golden_path): @@ -1972,12 +2071,12 @@ class TestFrankaPourCameraEnvCfg(FrankaPourEnvCfg): def __post_init__(self) -> None: super().__post_init__() self.seed = 42 - full_stage = self.curriculum_stage_names.index("full") - self.curriculum_start_stage = full_stage + # Match the training task default: start from the held, deeply tilted drain pose. + self.curriculum_start_stage = self.curriculum_stage_names.index("drain") self.curriculum_freeze = True self.decimation = 1 self.physics_substeps = 1 - self.mpm_iterations = 2 + self.mpm_iterations = 6 self.use_cuda_graph = False self.sim.render_interval = 1 self.sim.device = "cuda:0" @@ -2001,14 +2100,12 @@ def rendering_test_franka_pour( if not wp.is_cuda_available(): pytest.skip("FrankaPour rendering tests require a CUDA device.") - from isaaclab_tasks.contrib.franka_pour.pour_env import FrankaPourEnv - import isaaclab.sim as sim_utils + from isaaclab_tasks.contrib.franka_pour.pour_env import FrankaPourEnv + env_cfg = _make_franka_pour_camera_env_cfg(data_type) - env_cfg = _apply_overrides_to_env_cfg( - env_cfg, [f"presets={_physics_preset_name(physics_backend)},{renderer}"] - ) + env_cfg = _apply_overrides_to_env_cfg(env_cfg, [f"presets={_physics_preset_name(physics_backend)},{renderer}"]) env_cfg.scene.num_envs = 4 @@ -2030,19 +2127,31 @@ def rendering_test_franka_pour( maybe_save_stage(test_name, physics_backend, renderer, data_type) - zero_actions = torch.zeros(env.num_envs, env.action_manager.total_action_dim, device=env.device) - maybe_step_env_for_motion(env, data_type, num_steps=2, action_value=0.0) - if data_type != "motion_vectors": - env.step(zero_actions) - + camera_outputs = env.scene.sensors["tiled_camera"].data.output validate_camera_outputs( test_name, physics_backend, renderer, - env.scene.sensors["tiled_camera"].data.output, + camera_outputs, max_different_pixels_percentage=MAX_DIFFERENT_PIXELS_PERCENTAGE_BY_ENV_NAME[test_name], comparison_scores=comparison_scores, ) + + zero_actions = torch.zeros(env.num_envs, env.action_manager.total_action_dim, device=env.device) + pour_frames = _pour_frame_count() + frame_images: list[Image.Image] = [] + for step_index in range(pour_frames): + env.step(zero_actions) + frame_images = maybe_save_camera_frames( + test_name, + physics_backend, + renderer, + data_type, + env.scene.sensors["tiled_camera"].data.output, + step_index, + frame_images=frame_images, + ) + maybe_write_pour_frame_gif(test_name, physics_backend, renderer, data_type, frame_images) finally: if env is not None: env.close() From 16866ca122512e41f495f7fb1863c69122542672 Mon Sep 17 00:00:00 2001 From: HuiDong Chen Date: Tue, 21 Jul 2026 21:51:04 +0800 Subject: [PATCH 05/14] adapt the latest of `develop` - to be reviewed by Maximilian --- .../changelog.d/mpm-task-support.minor.rst | 2 - .../isaaclab_newton/physics/mpm_manager.py | 13 ++- .../physics/mpm_manager_cfg.py | 27 +----- .../isaaclab_newton/physics/newton_manager.py | 39 ++++++++ .../test_newton_manager_abstraction.py | 96 +++++-------------- .../contrib/franka_pour/pour_env_cfg.py | 22 ++--- .../test/contrib/test_franka_pour_env_cfg.py | 29 +++--- 7 files changed, 97 insertions(+), 131 deletions(-) diff --git a/source/isaaclab_newton/changelog.d/mpm-task-support.minor.rst b/source/isaaclab_newton/changelog.d/mpm-task-support.minor.rst index e0f554aebeee..45980c890693 100644 --- a/source/isaaclab_newton/changelog.d/mpm-task-support.minor.rst +++ b/source/isaaclab_newton/changelog.d/mpm-task-support.minor.rst @@ -3,8 +3,6 @@ Added * Added scoped Newton builder-world hooks and independent clone-source builder copies for tasks that extend replicated Newton worlds. -* Added isolated-world and bounded sparse-grid capacity options to - :class:`~isaaclab_newton.physics.MPMSolverCfg`. * Added :meth:`~isaaclab_newton.physics.NewtonManager.reset_solver_state` for clearing solver-owned history after selective simulation-state rewrites. diff --git a/source/isaaclab_newton/isaaclab_newton/physics/mpm_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/mpm_manager.py index 20b6099a0bd5..4cc6f6798270 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/mpm_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/mpm_manager.py @@ -28,10 +28,6 @@ def _make_solver_config(solver_cfg: MPMSolverCfg) -> SolverImplicitMPM.Config: grid_type=solver_cfg.grid_type, grid_padding=solver_cfg.grid_padding, max_active_cell_count=solver_cfg.max_active_cell_count, - max_leaf_node_count=solver_cfg.max_leaf_node_count, - max_lower_node_count=solver_cfg.max_lower_node_count, - max_upper_node_count=solver_cfg.max_upper_node_count, - separate_worlds=solver_cfg.separate_worlds, transfer_scheme=solver_cfg.transfer_scheme, integration_scheme=solver_cfg.integration_scheme, critical_fraction=solver_cfg.critical_fraction, @@ -124,6 +120,15 @@ def _build_solver(cls, model: Model, solver_cfg: MPMSolverCfg) -> None: NewtonManager._supports_rigid_body_force_input = False cls._project_outside_colliders = solver_cfg.project_outside_colliders + @classmethod + def _supports_cuda_graph_capture(cls) -> bool: + """Return ``True`` only for fixed-grid MPM. + + Sparse and dense grids reallocate as particles move, which is not + capturable in a CUDA graph; the fixed grid keeps a static topology. + """ + return cls._solver.grid_type == "fixed" + @classmethod def _step_solver( cls, state_0: State, state_1: State, control: Control, contacts: Contacts | None, substep_dt: float diff --git a/source/isaaclab_newton/isaaclab_newton/physics/mpm_manager_cfg.py b/source/isaaclab_newton/isaaclab_newton/physics/mpm_manager_cfg.py index 44b31375afa5..9b8c49489b49 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/mpm_manager_cfg.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/mpm_manager_cfg.py @@ -64,32 +64,7 @@ class MPMSolverCfg(NewtonSolverCfg): """Number of empty cells to add around particles when allocating the grid.""" max_active_cell_count: int = -1 - """Maximum active cell count shared by all worlds. - - A positive value also enables capacity-bounded sparse-grid rebuilding. - ``-1`` retains allocating sparse behavior and leaves other grids unbounded. - """ - - max_leaf_node_count: int = -1 - """Maximum sparse-grid leaf-node count shared by all worlds. - - ``-1`` derives the capacity from :attr:`max_active_cell_count`. - """ - - max_lower_node_count: int = -1 - """Maximum sparse-grid lower internal-node count shared by all worlds. - - ``-1`` derives the capacity from the initial topology. - """ - - max_upper_node_count: int = -1 - """Maximum sparse-grid upper internal-node count shared by all worlds. - - ``-1`` derives the capacity from the initial topology. - """ - - separate_worlds: bool = False - """Whether each Newton world uses an independent local MPM grid environment.""" + """Maximum active cell count for dense-grid active subsets. ``-1`` means unlimited.""" transfer_scheme: Literal["apic", "pic"] = "apic" """Particle-grid transfer scheme.""" diff --git a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py index a18cea3961d4..f18050aae9a4 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py @@ -69,6 +69,7 @@ def _paused_gc(): ModelFlags, ShapeFlags, State, + StateFlags, eval_fk, ) from newton.sensors import SensorContact as NewtonContactSensor @@ -2040,6 +2041,44 @@ def _reset_solver_internals(cls, world_mask: wp.array | None) -> None: return cls._solver.reset(cls._state_0, world_mask=world_mask, flags=0) + @classmethod + def reset_solver_state( + cls, + state: State | None = None, + world_mask: wp.array(dtype=wp.bool) | None = None, + flags: StateFlags | int | None = None, + ) -> None: + """Reset solver-private history after simulation state is rewritten. + + When :paramref:`state` is omitted, both distinct manager state buffers + are reset so a later buffer swap cannot restore stale solver history. + + Args: + state: State whose solver-private history should be reset. If + omitted, reset both manager states. + world_mask: Optional mask selecting Newton worlds to reset. + flags: State components whose solver-private history should reset. + + Raises: + RuntimeError: If the solver or a usable state is not initialized. + """ + if cls._solver is None: + raise RuntimeError("Newton solver is not initialized; cannot reset solver state.") + + candidates = (state,) if state is not None else (cls._state_1, cls._state_0) + states: list[State] = [] + seen: set[int] = set() + for candidate in candidates: + if candidate is None or id(candidate) in seen: + continue + seen.add(id(candidate)) + states.append(candidate) + if not states: + raise RuntimeError("Newton state is not initialized; provide an explicit state to reset.") + + for candidate in states: + cls._solver.reset(candidate, world_mask=world_mask, flags=flags) + # ----- Lifecycle orchestration ---------------------------------------- @classmethod diff --git a/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py b/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py index ecd7525f06cd..832440d0a03f 100644 --- a/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py +++ b/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py @@ -31,7 +31,6 @@ import numpy as np import pytest import warp as wp -from isaaclab_newton.cloner import copy_newton_source_builder, newton_builder_world_hook from isaaclab_newton.physics import ( FeatherstoneSolverCfg, KaminoSolverCfg, @@ -374,9 +373,6 @@ def test_mpm_solver_cfg_maps_only_newton_solver_fields(): ("grid_type", "dense"), ("grid_padding", 4), ("max_active_cell_count", 1024), - ("max_leaf_node_count", 512), - ("max_lower_node_count", 128), - ("max_upper_node_count", 32), ("transfer_scheme", "pic"), ("integration_scheme", "gimp"), ("critical_fraction", 0.25), @@ -385,7 +381,6 @@ def test_mpm_solver_cfg_maps_only_newton_solver_fields(): ("collider_basis", "Q1"), ("strain_basis", "P1d"), ("velocity_basis", "B2"), - ("separate_worlds", True), ] @@ -600,39 +595,24 @@ def counting_project(*args, **kwargs): assert calls["n"] == 0 -def test_mpm_solver_cfg_preserves_shared_world_default(): - """World-isolated MPM remains opt-in for backward compatibility.""" - - assert MPMSolverCfg().separate_worlds is False - - @pytest.mark.parametrize( - "grid_type, advertised, expected", + "grid_type, expected", [ - ("fixed", True, True), - ("sparse", True, True), - ("dense", False, False), + ("fixed", True), + ("sparse", False), + ("dense", False), ], ) -def test_mpm_cuda_graph_capture_uses_solver_capability(monkeypatch, grid_type, advertised, expected): - """The manager delegates graph safety to Newton's resolved solver configuration.""" +def test_mpm_cuda_graph_capture_supports_only_fixed_grid(monkeypatch, grid_type, expected): + """Newton implicit MPM is CUDA-graph capturable only with a fixed grid.""" - solver = SimpleNamespace(grid_type=grid_type, supports_graph_capture=advertised) - monkeypatch.setattr(NewtonManager, "_solver", solver, raising=False) + monkeypatch.setattr(NewtonManager, "_solver", SimpleNamespace(grid_type=grid_type), raising=False) assert NewtonMPMManager._supports_cuda_graph_capture() is expected -def test_cuda_graph_capture_keeps_legacy_solver_support(monkeypatch): - """Solvers without the optional capability property retain their existing capture path.""" - - monkeypatch.setattr(NewtonManager, "_solver", SimpleNamespace(), raising=False) - - assert NewtonManager._supports_cuda_graph_capture() is True - - -def test_solver_advertised_unsupported_cuda_graph_capture_uses_eager_execution(monkeypatch): - """A solver capability rejection prevents the manager from entering capture.""" +def test_mpm_unsupported_cuda_graph_capture_uses_eager_execution(monkeypatch): + """Sparse/dense MPM should not enter a CUDA graph capture window.""" from isaaclab.physics import PhysicsManager monkeypatch.setattr( @@ -642,70 +622,46 @@ def test_solver_advertised_unsupported_cuda_graph_capture_uses_eager_execution(m raising=False, ) monkeypatch.setattr(PhysicsManager, "_device", "cuda:0", raising=False) - monkeypatch.setattr(NewtonManager, "_solver", SimpleNamespace(supports_graph_capture=False), raising=False) + monkeypatch.setattr(NewtonManager, "_solver", SimpleNamespace(grid_type="sparse"), raising=False) monkeypatch.setattr(NewtonManager, "_graph", object(), raising=False) - monkeypatch.setattr(NewtonManager, "_graph_capture_pending", "standard", raising=False) + monkeypatch.setattr(NewtonManager, "_graph_capture_pending", True, raising=False) - NewtonManager._capture_or_defer_graph() + NewtonMPMManager._capture_or_defer_graph() assert NewtonManager._graph is None - assert NewtonManager._graph_capture_pending is None + assert NewtonManager._graph_capture_pending is False -@pytest.mark.parametrize("usdrt_stage, expected_mode", [(None, "standard"), (object(), "relaxed")]) -def test_cuda_graph_capture_is_deferred_with_explicit_mode(monkeypatch, usdrt_stage, expected_mode): - """Headless and RTX runs schedule their respective capture modes until the first step.""" +def test_cuda_graph_capture_uses_simulation_device(monkeypatch): + """CUDA graph capture should use the simulation device instead of Warp's default device.""" from isaaclab.physics import PhysicsManager - monkeypatch.setattr(PhysicsManager, "_cfg", SimpleNamespace(use_cuda_graph=True), raising=False) - monkeypatch.setattr(PhysicsManager, "_device", "cuda:1", raising=False) - monkeypatch.setattr(NewtonManager, "_usdrt_stage", usdrt_stage, raising=False) - monkeypatch.setattr(NewtonManager, "_solver", SimpleNamespace(supports_graph_capture=True), raising=False) - monkeypatch.setattr(NewtonManager, "_graph", object(), raising=False) - monkeypatch.setattr(NewtonManager, "_graph_capture_pending", None, raising=False) - - NewtonManager._capture_or_defer_graph() - - assert NewtonManager._graph is None - assert NewtonManager._graph_capture_pending == expected_mode - - -def test_standard_cuda_graph_capture_prepares_solver_before_recording(monkeypatch): - """Solver-owned persistent resources are prepared before the standard capture window.""" - events = [] - contacts = object() + captured_devices = [] captured_graph = object() class FakeScopedCapture: - def __init__(self, device=None, **_kwargs): - assert device == "cuda:0" + def __init__(self, device=None): + captured_devices.append(device) self.graph = captured_graph def __enter__(self): - events.append(("capture", None)) return self def __exit__(self, exc_type, exc_value, traceback): return False - solver = SimpleNamespace( - supports_graph_capture=True, - prepare_graph_capture=lambda received: events.append(("prepare", received)), - ) - monkeypatch.setattr(NewtonManager, "_solver", solver, raising=False) - monkeypatch.setattr(NewtonManager, "_contacts", contacts, raising=False) + monkeypatch.setattr(PhysicsManager, "_cfg", SimpleNamespace(use_cuda_graph=True), raising=False) + monkeypatch.setattr(PhysicsManager, "_device", "cuda:1", raising=False) + monkeypatch.setattr(NewtonManager, "_usdrt_stage", None, raising=False) + monkeypatch.setattr(NewtonManager, "_solver", None, raising=False) monkeypatch.setattr(NewtonManager, "_is_all_graphable", classmethod(lambda cls: False)) - monkeypatch.setattr( - NewtonManager, - "_simulate_physics_only", - classmethod(lambda cls: events.append(("simulate", None))), - ) + monkeypatch.setattr(NewtonManager, "_simulate_physics_only", classmethod(lambda cls: None)) monkeypatch.setattr(wp, "ScopedCapture", FakeScopedCapture) - graph = NewtonManager._capture_standard_graph("cuda:0") + NewtonManager._capture_or_defer_graph() - assert events == [("prepare", contacts), ("capture", None), ("simulate", None)] - assert graph is captured_graph + assert captured_devices == ["cuda:1"] + assert NewtonManager._graph is captured_graph def test_relaxed_cuda_graph_capture_prepares_solver_before_warmup(monkeypatch): diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/pour_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/pour_env_cfg.py index ed571bcaa87c..dd2611b1fb99 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/pour_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/pour_env_cfg.py @@ -84,6 +84,10 @@ } ) SPILL_FLOOR_LABEL_PATTERN = r".*/SpillFloor$" +# Coupler body selectors use full Newton body-label regexes (not SceneEntityCfg). +ROBOT_BODY_LABEL_PATTERN = r"/World/envs/env_.*/Robot" +SOURCE_CUP_BODY_LABEL_PATTERN = r"/World/envs/env_.*/SourceCup" +TARGET_CUP_BODY_LABEL_PATTERN = r"/World/envs/env_.*/TargetCup" GRASP_APPROACH_STAGE_NAMES = ( "approach_1", "approach_2", @@ -998,7 +1002,6 @@ def __post_init__(self): self.sim.physics = NewtonCfg( solver_cfg=CouplerProxyCfg( - scene_cfg=self.scene, entries=[ CouplerEntryCfg( name=RIGID_ENTRY, @@ -1010,9 +1013,9 @@ def __post_init__(self): use_mujoco_contacts=False, integrator="implicitfast", njmax=510, nconmax=400 ), bodies=[ - SceneEntityCfg("robot"), - SceneEntityCfg("source_cup"), - SceneEntityCfg("target_cup"), + ROBOT_BODY_LABEL_PATTERN, + SOURCE_CUP_BODY_LABEL_PATTERN, + TARGET_CUP_BODY_LABEL_PATTERN, ], include_static_shapes=True, substeps=self.rigid_entry_substeps, @@ -1036,7 +1039,6 @@ def __post_init__(self): # Keep the task's validated nonlinear solve while sparse topology is # rebuilt eagerly around the physically separated environments. solver="jacobi", - separate_worlds=True, ), all_particles=True, bodies=[SPILL_FLOOR_LABEL_PATTERN], @@ -1050,7 +1052,7 @@ def __post_init__(self): CouplerProxyMappingCfg( source=RIGID_ENTRY, destination=MPM_ENTRY, - bodies=[SceneEntityCfg("source_cup"), SceneEntityCfg("target_cup")], + bodies=[SOURCE_CUP_BODY_LABEL_PATTERN, TARGET_CUP_BODY_LABEL_PATTERN], mass_scale=self.proxy_mass_scale, mode="lagged", # Implicit MPM resolves its proxy colliders internally; the shared outer @@ -1933,14 +1935,6 @@ def finalize(self) -> FrankaPourEnvCfg: ) mpm_solver_cfg = _mpm_solver_cfg(resolved) mpm_solver_cfg.max_active_cell_count = _resolve_mpm_cell_cap(resolved) - if mpm_solver_cfg.grid_type == "sparse" and resolved.sim.physics.use_cuda_graph: - # The compact initial fill underestimates hierarchy nodes needed after a particle moves - # into a different NanoVDB region. Reserve the task's workspace-derived headroom per - # independent world; this changes capacity only, not MPM stepping or physics. - world_count = int(resolved.scene.num_envs) - mpm_solver_cfg.max_lower_node_count = max(32, 16 * world_count) - mpm_solver_cfg.max_upper_node_count = max(32, (world_count + 1) // 2) - resolved.sim.physics.solver_cfg.scene_cfg = resolved.scene return resolved diff --git a/source/isaaclab_tasks/test/contrib/test_franka_pour_env_cfg.py b/source/isaaclab_tasks/test/contrib/test_franka_pour_env_cfg.py index 3595e9059887..4e1f6ae8618a 100644 --- a/source/isaaclab_tasks/test/contrib/test_franka_pour_env_cfg.py +++ b/source/isaaclab_tasks/test/contrib/test_franka_pour_env_cfg.py @@ -23,7 +23,7 @@ from isaaclab.sim.schemas import MassCfg, UsdPhysicsCollisionCfg, UsdPhysicsRigidBodyCfg from isaaclab.sim.spawners.materials import RigidBodyMaterialBaseCfg -from isaaclab_contrib.coupling import CouplerProxyCfg, NewtonCoupler +from isaaclab_contrib.coupling import CouplerProxyCfg, NewtonCouplerManager import isaaclab_tasks.contrib.franka_pour as franka_pour import isaaclab_tasks.contrib.franka_pour.config.franka # noqa: F401 @@ -252,7 +252,6 @@ def test_finalize_builds_scene_assets_without_mutating_the_caller(): assert isinstance(resolved.scene.media, MPMObjectCfg) assert resolved is not original assert resolved.scene is not original.scene - assert resolved.sim.physics.solver_cfg.scene_cfg is resolved.scene assert _media_capacity(resolved) == 8 * 512 @@ -426,22 +425,24 @@ def test_cfg_routes_each_body_to_exactly_one_solver(): media = entries["media"] assert arm.solver_cfg.integrator == "implicitfast" assert arm.solver_cfg.use_mujoco_contacts is False - assert solver.scene_cfg is cfg.scene assert cfg.sim.physics.collision_cfg.soft_contact_max == 0 assert arm.include_static_shapes is True - assert arm.bodies == [SceneEntityCfg("robot"), SceneEntityCfg("source_cup"), SceneEntityCfg("target_cup")] + assert arm.bodies == [ + pour_env_cfg.ROBOT_BODY_LABEL_PATTERN, + pour_env_cfg.SOURCE_CUP_BODY_LABEL_PATTERN, + pour_env_cfg.TARGET_CUP_BODY_LABEL_PATTERN, + ] assert cfg.sim.physics.num_substeps == cfg.physics_substeps == 1 assert arm.substeps == cfg.rigid_entry_substeps == 4 assert media.substeps == cfg.mpm_entry_substeps == 2 assert cfg.sim.physics.num_substeps * media.substeps == 2 assert media.all_particles is True - assert media.bodies == [r".*/SpillFloor$"] + assert media.bodies == [pour_env_cfg.SPILL_FLOOR_LABEL_PATTERN] assert media.include_static_shapes is False assert media.in_place is True assert media.solver_cfg.grid_type == "sparse" assert media.solver_cfg.grid_padding == 0 assert media.solver_cfg.max_active_cell_count == 2 * 512 - assert media.solver_cfg.separate_worlds is True assert media.solver_cfg.solver == "jacobi" assert media.solver_cfg.warmstart_mode == "none" assert media.solver_cfg.max_iterations == 24 @@ -450,7 +451,10 @@ def test_cfg_routes_each_body_to_exactly_one_solver(): proxies = solver.proxies assert len(proxies) == 1 assert proxies[0].source == "arm" and proxies[0].destination == "media" - assert proxies[0].bodies == [SceneEntityCfg("source_cup"), SceneEntityCfg("target_cup")] + assert proxies[0].bodies == [ + pour_env_cfg.SOURCE_CUP_BODY_LABEL_PATTERN, + pour_env_cfg.TARGET_CUP_BODY_LABEL_PATTERN, + ] assert proxies[0].collision_pipeline is not None assert proxies[0].collision_pipeline(None) is None assert proxies[0].mass_scale == pytest.approx(cfg.proxy_mass_scale) @@ -1300,7 +1304,6 @@ def test_reset_dataset_play_preserves_policy_abi_with_captured_sparse_grid(): assert play_solver_cfg.grid_type == "sparse" assert play_solver_cfg.grid_padding == 0 assert play_solver_cfg.max_active_cell_count == 512 - assert play_solver_cfg.separate_worlds is True assert play_cfg.sim.physics.use_cuda_graph is True assert play_cfg.terminations.success.func is mdp.immediate_pour_success assert play_cfg.terminations.success.params == {} @@ -1802,9 +1805,7 @@ def test_sparse_training_reserves_capturable_isolated_grid_capacity(num_envs): assert _media_entry(resolved).solver_cfg.grid_type == "sparse" assert _media_capacity(resolved) == 512 * num_envs solver_cfg = _media_entry(resolved).solver_cfg - assert solver_cfg.separate_worlds is True - assert solver_cfg.max_lower_node_count == max(32, 16 * num_envs) - assert solver_cfg.max_upper_node_count == max(32, (num_envs + 1) // 2) + assert solver_cfg.max_active_cell_count == 512 * num_envs assert resolved.scene.env_spacing == pytest.approx(cfg.scene.env_spacing) assert resolved.sim.physics.use_cuda_graph is True @@ -1841,12 +1842,10 @@ def test_mpm_uses_captured_sparse_training_and_play_configs(): assert solver_cfg.project_outside_colliders is False assert solver_cfg.grid_type == "sparse" assert solver_cfg.max_active_cell_count == 200 * 512 - assert solver_cfg.separate_worlds is True assert play_solver_cfg.collider_basis == "pic27" assert play_solver_cfg.grid_type == "sparse" assert play_solver_cfg.grid_padding == 0 assert play_solver_cfg.max_active_cell_count == 512 - assert play_solver_cfg.separate_worlds is True assert play_cfg.sim.physics.use_cuda_graph is True @@ -1941,7 +1940,7 @@ def test_media_selector_includes_spill_floor_without_unrelated_shapes(): particle_count=3, ) - resolved = NewtonCoupler._resolve_entry(model, media, cfg.scene) + resolved = NewtonCouplerManager._resolve_entry(model, media) assert resolved.bodies == [1] assert resolved.shapes == [1] @@ -2017,7 +2016,7 @@ def test_task_source_does_not_traverse_private_solver_state(): if isinstance(node, ast.Attribute) and node.attr.startswith("_") and isinstance(node.value, ast.Name) - and node.value.id in {"NewtonManager", "NewtonCoupler"} + and node.value.id in {"NewtonManager", "NewtonCouplerManager"} } assert private_manager_attrs == set() From 97f2ae0a910a18979758de72d4ce336b696d1702 Mon Sep 17 00:00:00 2001 From: HuiDong Chen Date: Wed, 22 Jul 2026 19:36:52 +0800 Subject: [PATCH 06/14] Fix OVRTX particle and camera bindings Prime MPM points before asynchronous GPU updates and bind cloned render products to per-environment cameras. Expand Franka Pour rendering coverage and add Newton renderer golden images. --- .../test/test_ovrtx_deformable_bindings.py | 3 + .../newton-isaacsim_rtx_renderer-albedo.png | 3 + .../newton-isaacsim_rtx_renderer-depth.png | 3 + ...aacsim_rtx_renderer-distance_to_camera.png | 3 + ...m_rtx_renderer-distance_to_image_plane.png | 3 + ...renderer-instance_id_segmentation_fast.png | 3 + ...tx_renderer-instance_segmentation_fast.png | 3 + .../newton-isaacsim_rtx_renderer-normals.png | 3 + .../newton-isaacsim_rtx_renderer-rgb.png | 3 + .../newton-isaacsim_rtx_renderer-rgba.png | 3 + ...sim_rtx_renderer-semantic_segmentation.png | 3 + ...nderer-simple_shading_constant_diffuse.png | 3 + ...tx_renderer-simple_shading_diffuse_mdl.png | 3 + ...m_rtx_renderer-simple_shading_full_mdl.png | 3 + .../newton-newton_renderer-depth.png | 3 + ...ton-newton_renderer-distance_to_camera.png | 3 + ...ewton_renderer-distance_to_image_plane.png | 3 + .../newton-newton_renderer-normals.png | 3 + .../newton-newton_renderer-rgb.png | 3 + .../newton-newton_renderer-rgba.png | 3 + .../newton-ovrtx_renderer-albedo.png | 3 + .../newton-ovrtx_renderer-depth.png | 3 + ...wton-ovrtx_renderer-distance_to_camera.png | 3 + ...ovrtx_renderer-distance_to_image_plane.png | 3 + ...tx_renderer-instance_segmentation_fast.png | 3 + .../newton-ovrtx_renderer-normals.png | 3 + .../franka_pour/newton-ovrtx_renderer-rgb.png | 3 + .../newton-ovrtx_renderer-rgba.png | 3 + ...n-ovrtx_renderer-semantic_segmentation.png | 3 + ...nderer-simple_shading_constant_diffuse.png | 3 + ...tx_renderer-simple_shading_diffuse_mdl.png | 3 + ...ovrtx_renderer-simple_shading_full_mdl.png | 3 + .../test/rendering_test_utils.py | 79 ++++++++++++------- 33 files changed, 145 insertions(+), 30 deletions(-) create mode 100644 source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-albedo.png create mode 100644 source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-depth.png create mode 100644 source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-distance_to_camera.png create mode 100644 source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-distance_to_image_plane.png create mode 100644 source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-instance_id_segmentation_fast.png create mode 100644 source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-instance_segmentation_fast.png create mode 100644 source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-normals.png create mode 100644 source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-rgb.png create mode 100644 source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-rgba.png create mode 100644 source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-semantic_segmentation.png create mode 100644 source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-simple_shading_constant_diffuse.png create mode 100644 source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-simple_shading_diffuse_mdl.png create mode 100644 source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-simple_shading_full_mdl.png create mode 100644 source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-depth.png create mode 100644 source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-distance_to_camera.png create mode 100644 source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-distance_to_image_plane.png create mode 100644 source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-normals.png create mode 100644 source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-rgb.png create mode 100644 source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-rgba.png create mode 100644 source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-albedo.png create mode 100644 source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-depth.png create mode 100644 source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-distance_to_camera.png create mode 100644 source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-distance_to_image_plane.png create mode 100644 source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-instance_segmentation_fast.png create mode 100644 source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-normals.png create mode 100644 source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-rgb.png create mode 100644 source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-rgba.png create mode 100644 source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-semantic_segmentation.png create mode 100644 source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-simple_shading_constant_diffuse.png create mode 100644 source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-simple_shading_diffuse_mdl.png create mode 100644 source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-simple_shading_full_mdl.png diff --git a/source/isaaclab_ov/test/test_ovrtx_deformable_bindings.py b/source/isaaclab_ov/test/test_ovrtx_deformable_bindings.py index 871730c398b0..3dcb3aa4a350 100644 --- a/source/isaaclab_ov/test/test_ovrtx_deformable_bindings.py +++ b/source/isaaclab_ov/test/test_ovrtx_deformable_bindings.py @@ -89,6 +89,9 @@ def query_prims(self, **kwargs): # noqa: ARG002 def write_attribute(self, **kwargs): self.writes.append(kwargs) + def write_array_attribute(self, **kwargs): + self.writes.append(kwargs) + def _make_renderer_without_backend(device: str = "cpu") -> tuple[OVRTXRenderer, _FakeOVRTXBackend]: renderer = OVRTXRenderer.__new__(OVRTXRenderer) diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-albedo.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-albedo.png new file mode 100644 index 000000000000..812d169d74bb --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-albedo.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:06c1648e4cf9323882caeae4c3db1eb91dd5323943bd1779bb70c34e6ed1c7d6 +size 5915 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-depth.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-depth.png new file mode 100644 index 000000000000..b45b349f1ba1 --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-depth.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5cd056b828384c5dc985db16a927278c17a9b0574470f33cb81b7f959ae9766a +size 8818 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-distance_to_camera.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-distance_to_camera.png new file mode 100644 index 000000000000..061f4b8c285b --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-distance_to_camera.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:31614a911c1439ad1e66d6a73001c67f0d9a31e1004c112a13293413f8d77714 +size 10822 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-distance_to_image_plane.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-distance_to_image_plane.png new file mode 100644 index 000000000000..561bd8508317 --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-distance_to_image_plane.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6493b29b2fe01b89c9fdb8543a2ff268a4c0b8b87962475362751afe6e65f46b +size 8828 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-instance_id_segmentation_fast.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-instance_id_segmentation_fast.png new file mode 100644 index 000000000000..06a1c603e731 --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-instance_id_segmentation_fast.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3f7148d1ca5409457cdba5e368a4a95a3dca31e6c8b804bbc184ecc10103fcb0 +size 8800 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-instance_segmentation_fast.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-instance_segmentation_fast.png new file mode 100644 index 000000000000..4c4c737e6499 --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-instance_segmentation_fast.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:786b50054bb5e55c201d58b00f495e0bb7b5e9e34428d6cc5ef14a62dabf5ea9 +size 2110 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-normals.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-normals.png new file mode 100644 index 000000000000..f9ba1a9d7c82 --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-normals.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:39d46ab76b27b601d8c9c4253192834b6c10f13d924b17eab6b6878116c8b084 +size 20075 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-rgb.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-rgb.png new file mode 100644 index 000000000000..ee816e94c278 --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-rgb.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:514aceb37032f2930718c2134be2f366268878bf106e4fd901f9457e486a29b4 +size 69842 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-rgba.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-rgba.png new file mode 100644 index 000000000000..04134baa0bdc --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-rgba.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ca6266b26059357ebb12d1be38f81b656ed1c48dbb67132a94f7e62b13260f88 +size 78092 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-semantic_segmentation.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-semantic_segmentation.png new file mode 100644 index 000000000000..4c4c737e6499 --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-semantic_segmentation.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:786b50054bb5e55c201d58b00f495e0bb7b5e9e34428d6cc5ef14a62dabf5ea9 +size 2110 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-simple_shading_constant_diffuse.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-simple_shading_constant_diffuse.png new file mode 100644 index 000000000000..caff703abe07 --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-simple_shading_constant_diffuse.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:bedcbf5216709c415d2f0d942eec7b15cc9da7a1a290dc320bff14b971b0c371 +size 13345 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-simple_shading_diffuse_mdl.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-simple_shading_diffuse_mdl.png new file mode 100644 index 000000000000..8b135d9a572c --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-simple_shading_diffuse_mdl.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:54843fb067153f344485c16db96014fb2e06a9a94ce07ede3c8153b33ef1302b +size 20118 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-simple_shading_full_mdl.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-simple_shading_full_mdl.png new file mode 100644 index 000000000000..294fd96fb8d7 --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-isaacsim_rtx_renderer-simple_shading_full_mdl.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:828cdc85c77edb76050b93e925c8fc1146e1dcb06858b3348c786a8bc2ca1191 +size 27153 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-depth.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-depth.png new file mode 100644 index 000000000000..8abbdcd02669 --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-depth.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8d3fc3d970aa13d2618321c356ebb2fa2c77429dc291d91ddddac5af1d698fc0 +size 8885 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-distance_to_camera.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-distance_to_camera.png new file mode 100644 index 000000000000..c9414e987719 --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-distance_to_camera.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6ede1eb901f12d2c6820ee3b8b22b70b175f01791cd646cb6fc90be1b9b889be +size 10888 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-distance_to_image_plane.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-distance_to_image_plane.png new file mode 100644 index 000000000000..1f335e203698 --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-distance_to_image_plane.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:be3d18079bb5bd8a4a3e97207127dce9d5ad8bc82967f1b317e37bac79260a4f +size 8888 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-normals.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-normals.png new file mode 100644 index 000000000000..bd35deca4508 --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-normals.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8d3ed415530b4a55344741a2279467ac503bd6a39221111b90a07cd36529f83c +size 15296 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-rgb.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-rgb.png new file mode 100644 index 000000000000..b96aad2b601e --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-rgb.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7e6f85aa16cfa86b897fe3425a345ead3ae61cfacd85316f7970086762ba9302 +size 13383 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-rgba.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-rgba.png new file mode 100644 index 000000000000..841baa89eb56 --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-rgba.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:08ce608f7daa4e37552442cf98ec089476af8715ee7639e99f20f5ac498b21c8 +size 14532 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-albedo.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-albedo.png new file mode 100644 index 000000000000..5cf5e01789e8 --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-albedo.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1f2bb3d714ad4174cd84132e66569b370a967d90d0dc2994142363b9b1f32739 +size 11018 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-depth.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-depth.png new file mode 100644 index 000000000000..f4db92893809 --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-depth.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:003435c01d3e71204c16c1b1a076e0b30b21356b1f534cf8c81e3b2cdac56b2f +size 8914 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-distance_to_camera.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-distance_to_camera.png new file mode 100644 index 000000000000..86219a6307d8 --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-distance_to_camera.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2dd3f07cefc4eaf78a32e4d8ac4e23d2417c9906b5e0d01ee3874c53675bac85 +size 11022 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-distance_to_image_plane.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-distance_to_image_plane.png new file mode 100644 index 000000000000..7db92effb4b8 --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-distance_to_image_plane.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:273764d912f188e303fdc02e6f81b7bc52328d3b199b01a23358a83f63218ad8 +size 8914 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-instance_segmentation_fast.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-instance_segmentation_fast.png new file mode 100644 index 000000000000..667eb49d4b1e --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-instance_segmentation_fast.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ffb95913b24bf5beb02b8ab00f37b7e2f8af52762958dffdfd631947b9e57014 +size 2050 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-normals.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-normals.png new file mode 100644 index 000000000000..34735d340e10 --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-normals.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ea0402058dea4091439eaae32b40bf4cf38c8cfee6466e0752b729613660a232 +size 17842 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-rgb.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-rgb.png new file mode 100644 index 000000000000..7f972f5a88ab --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-rgb.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:cbe40bfa5f4dd030a8757fcc30aeaf2ae03917f2ab28009d313ff8a76cc59525 +size 60860 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-rgba.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-rgba.png new file mode 100644 index 000000000000..a9af7ae67fb7 --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-rgba.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ddc24ea597740c41757261fb597694c0be7e16fa658d3624795944447e489a46 +size 68618 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-semantic_segmentation.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-semantic_segmentation.png new file mode 100644 index 000000000000..667eb49d4b1e --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-semantic_segmentation.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ffb95913b24bf5beb02b8ab00f37b7e2f8af52762958dffdfd631947b9e57014 +size 2050 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-simple_shading_constant_diffuse.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-simple_shading_constant_diffuse.png new file mode 100644 index 000000000000..40111f152688 --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-simple_shading_constant_diffuse.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8ef976b613a6e7c427d420ef34120cfa487a3948588377679bc5fd31b7f788bf +size 11613 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-simple_shading_diffuse_mdl.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-simple_shading_diffuse_mdl.png new file mode 100644 index 000000000000..4384df2b9b25 --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-simple_shading_diffuse_mdl.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8cff6b6e6610daf3f64f21674e6c17916753c003101fcadb45b961447388f8eb +size 13171 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-simple_shading_full_mdl.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-simple_shading_full_mdl.png new file mode 100644 index 000000000000..b3b4f3b9f73c --- /dev/null +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-ovrtx_renderer-simple_shading_full_mdl.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ae7b3e185809859781a2362e2a90e83cfd3a357b3067b16683ccd6f77b2481f3 +size 19627 diff --git a/source/isaaclab_tasks/test/rendering_test_utils.py b/source/isaaclab_tasks/test/rendering_test_utils.py index a3e58ca8cc32..94ba412ad723 100644 --- a/source/isaaclab_tasks/test/rendering_test_utils.py +++ b/source/isaaclab_tasks/test/rendering_test_utils.py @@ -782,6 +782,27 @@ def _physics_preset_name_deformable(physics_backend: str) -> str: return "newton_mjwarp_vbd_proxy" if physics_backend == "newton" else physics_backend +def _supported_presets_for_concrete_physics(physics_cfg: Any) -> set[str]: + """Map a fixed (non-``PresetCfg``) physics cfg to Hydra preset names it can satisfy. + + Returns an empty set when the concrete physics type is unrecognized. + """ + supported_presets = set() + + name = type(physics_cfg).__name__ + if name == "NewtonCfg": + # Fixed Newton backends accept any newton_* Hydra label used by these tests. + supported_presets.add("newton") + supported_presets.add("newton_mjwarp") + supported_presets.add("newton_mjwarp_vbd") + if name == "PhysxCfg": + supported_presets.add("physx") + if name == "OvPhysxCfg": + supported_presets.add("ovphysx") + + return supported_presets + + def _skip_if_physics_preset_unsupported(env_cfg: Any, physics_preset_name: str) -> None: """Skip the test when the env does not support the given physics preset. @@ -2090,22 +2111,18 @@ def rendering_test_franka_pour( data_type: str, comparison_scores: list[dict], ) -> None: - if physics_backend in ("physx", "ovphysx"): - pytest.skip("FrankaPour env cfg is Newton-only.") - _skip_if_newton_motion_vectors(physics_backend, data_type) - import warp as wp - - if not wp.is_cuda_available(): - pytest.skip("FrankaPour rendering tests require a CUDA device.") - - import isaaclab.sim as sim_utils - from isaaclab_tasks.contrib.franka_pour.pour_env import FrankaPourEnv env_cfg = _make_franka_pour_camera_env_cfg(data_type) - env_cfg = _apply_overrides_to_env_cfg(env_cfg, [f"presets={_physics_preset_name(physics_backend)},{renderer}"]) + + # Skip if the physics preset is not supported by the env cfg. + physics_preset_name = _physics_preset_name(physics_backend) + if physics_preset_name not in _supported_presets_for_concrete_physics(env_cfg.sim.physics): + pytest.skip(f"FrankaPour env cfg does not support '{physics_preset_name}'.") + + env_cfg = _apply_overrides_to_env_cfg(env_cfg, [f"presets={physics_preset_name},{renderer}"]) env_cfg.scene.num_envs = 4 @@ -2115,11 +2132,9 @@ def rendering_test_franka_pour( test_name = "franka_pour" env = None + FRAMES_BEFORE_POURING_PARTICLES = 40 + try: - # Kit-based RTX tests need a fresh stage between parametrized cases. Kitless OVRTX uses a - # different USD bootstrap path and can segfault if we force-create a stage here. - if renderer == "isaacsim_rtx_renderer": - sim_utils.create_new_stage() env = FrankaPourEnv(env_cfg) env.sim._app_control_on_stop_handle = None @@ -2127,6 +2142,10 @@ def rendering_test_franka_pour( maybe_save_stage(test_name, physics_backend, renderer, data_type) + zero_actions = torch.zeros(env.num_envs, env.action_manager.total_action_dim, device=env.device) + for _ in range(FRAMES_BEFORE_POURING_PARTICLES): + env.step(zero_actions) + camera_outputs = env.scene.sensors["tiled_camera"].data.output validate_camera_outputs( test_name, @@ -2137,21 +2156,21 @@ def rendering_test_franka_pour( comparison_scores=comparison_scores, ) - zero_actions = torch.zeros(env.num_envs, env.action_manager.total_action_dim, device=env.device) - pour_frames = _pour_frame_count() - frame_images: list[Image.Image] = [] - for step_index in range(pour_frames): - env.step(zero_actions) - frame_images = maybe_save_camera_frames( - test_name, - physics_backend, - renderer, - data_type, - env.scene.sensors["tiled_camera"].data.output, - step_index, - frame_images=frame_images, - ) - maybe_write_pour_frame_gif(test_name, physics_backend, renderer, data_type, frame_images) + # zero_actions = torch.zeros(env.num_envs, env.action_manager.total_action_dim, device=env.device) + # pour_frames = 40 + # frame_images: list[Image.Image] = [] + # for step_index in range(pour_frames): + # env.step(zero_actions) + # frame_images = maybe_save_camera_frames( + # test_name, + # physics_backend, + # renderer, + # data_type, + # env.scene.sensors["tiled_camera"].data.output, + # step_index, + # frame_images=frame_images, + # ) + # maybe_write_pour_frame_gif(test_name, physics_backend, renderer, data_type, frame_images) finally: if env is not None: env.close() From 9561e382947797480a226db78959c48c6c36ec48 Mon Sep 17 00:00:00 2001 From: HuiDong Chen Date: Thu, 23 Jul 2026 08:58:51 +0800 Subject: [PATCH 07/14] Remove the unused environment count argument because registered particle paths already identify each visual instance. Align tests and backend stubs with the exercised API surface. --- source/isaaclab_ov/test/test_ovrtx_deformable_bindings.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/source/isaaclab_ov/test/test_ovrtx_deformable_bindings.py b/source/isaaclab_ov/test/test_ovrtx_deformable_bindings.py index 3dcb3aa4a350..871730c398b0 100644 --- a/source/isaaclab_ov/test/test_ovrtx_deformable_bindings.py +++ b/source/isaaclab_ov/test/test_ovrtx_deformable_bindings.py @@ -89,9 +89,6 @@ def query_prims(self, **kwargs): # noqa: ARG002 def write_attribute(self, **kwargs): self.writes.append(kwargs) - def write_array_attribute(self, **kwargs): - self.writes.append(kwargs) - def _make_renderer_without_backend(device: str = "cpu") -> tuple[OVRTXRenderer, _FakeOVRTXBackend]: renderer = OVRTXRenderer.__new__(OVRTXRenderer) From 1da0a7ade11197fb942a808f11408455104a05e1 Mon Sep 17 00:00:00 2001 From: HuiDong Chen Date: Thu, 23 Jul 2026 13:04:24 +0800 Subject: [PATCH 08/14] newton warp golden images --- .../franka_pour/newton-newton_renderer-depth.png | 4 ++-- .../franka_pour/newton-newton_renderer-distance_to_camera.png | 4 ++-- .../newton-newton_renderer-distance_to_image_plane.png | 4 ++-- .../franka_pour/newton-newton_renderer-normals.png | 4 ++-- .../golden_images/franka_pour/newton-newton_renderer-rgb.png | 4 ++-- .../golden_images/franka_pour/newton-newton_renderer-rgba.png | 4 ++-- 6 files changed, 12 insertions(+), 12 deletions(-) diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-depth.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-depth.png index 8abbdcd02669..486f9aa45084 100644 --- a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-depth.png +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-depth.png @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:8d3fc3d970aa13d2618321c356ebb2fa2c77429dc291d91ddddac5af1d698fc0 -size 8885 +oid sha256:f052269d4a9d8b97b05bdf9cc6ee6969c536a63afe5d74cc375678d2e5b1ad07 +size 9537 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-distance_to_camera.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-distance_to_camera.png index c9414e987719..c947d6686e26 100644 --- a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-distance_to_camera.png +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-distance_to_camera.png @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:6ede1eb901f12d2c6820ee3b8b22b70b175f01791cd646cb6fc90be1b9b889be -size 10888 +oid sha256:fad45f724f6bb14cfa4189df80350cf59697c6d5eea222940c191ecd8c72982c +size 11556 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-distance_to_image_plane.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-distance_to_image_plane.png index 1f335e203698..053a936eb860 100644 --- a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-distance_to_image_plane.png +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-distance_to_image_plane.png @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:be3d18079bb5bd8a4a3e97207127dce9d5ad8bc82967f1b317e37bac79260a4f -size 8888 +oid sha256:6b05a932ba9b128f869b8a79b7bd22b6ae897fa5af1f6ee3666c5be7938659e3 +size 9587 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-normals.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-normals.png index bd35deca4508..a45a7cd2c487 100644 --- a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-normals.png +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-normals.png @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:8d3ed415530b4a55344741a2279467ac503bd6a39221111b90a07cd36529f83c -size 15296 +oid sha256:6209e4b69f9961d94652db1836ffd5cc4e81771dc34590dfaf36c65bef7d58a0 +size 17144 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-rgb.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-rgb.png index b96aad2b601e..67055ebab7b9 100644 --- a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-rgb.png +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-rgb.png @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:7e6f85aa16cfa86b897fe3425a345ead3ae61cfacd85316f7970086762ba9302 -size 13383 +oid sha256:0fe118b596810d2adb063ad4223fced75514bd10ba5c2751b4ba6d878975f8b8 +size 14967 diff --git a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-rgba.png b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-rgba.png index 841baa89eb56..1f3b4d7cdc6d 100644 --- a/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-rgba.png +++ b/source/isaaclab_tasks/test/golden_images/franka_pour/newton-newton_renderer-rgba.png @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:08ce608f7daa4e37552442cf98ec089476af8715ee7639e99f20f5ac498b21c8 -size 14532 +oid sha256:00cbc54fafb70e6ce512d9e28e81975e063daf361b4496d51ad930c00df6dcdc +size 16259 From 9e32537e2dc7f8e41a61b99b4da3ed0065dc9a24 Mon Sep 17 00:00:00 2001 From: HuiDong Chen Date: Thu, 23 Jul 2026 16:27:05 +0800 Subject: [PATCH 09/14] tiny cleanup --- source/isaaclab_tasks/test/contrib/test_franka_pour_env_cfg.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/source/isaaclab_tasks/test/contrib/test_franka_pour_env_cfg.py b/source/isaaclab_tasks/test/contrib/test_franka_pour_env_cfg.py index 4e1f6ae8618a..9bcce7f5c8b0 100644 --- a/source/isaaclab_tasks/test/contrib/test_franka_pour_env_cfg.py +++ b/source/isaaclab_tasks/test/contrib/test_franka_pour_env_cfg.py @@ -19,7 +19,7 @@ from isaaclab_newton.sim.schemas import MujocoJointCfg from isaaclab.assets import RigidObjectCfg -from isaaclab.managers import CurriculumTermCfg, RewardTermCfg, SceneEntityCfg, TerminationTermCfg +from isaaclab.managers import CurriculumTermCfg, RewardTermCfg, TerminationTermCfg from isaaclab.sim.schemas import MassCfg, UsdPhysicsCollisionCfg, UsdPhysicsRigidBodyCfg from isaaclab.sim.spawners.materials import RigidBodyMaterialBaseCfg From ee735aebb65255c68a8c12dc9f72224f9963b670 Mon Sep 17 00:00:00 2001 From: HuiDong Chen Date: Fri, 24 Jul 2026 15:36:58 +0800 Subject: [PATCH 10/14] Fix Franka pour early pxr import - To be reviewed by Maximilian Franka pour training could warn or crash because task configuration resolution imported pxr before SimulationApp launched. The eager MDP exports loaded runtime action implementations through Articulation and the USD cloner. Split action configuration from runtime implementations, use resolvable class strings, and lazily export MDP symbols so config construction remains Kit-free. --- .../franka-pour-early-pxr-import.rst | 6 + .../contrib/franka_pour/mdp/__init__.py | 102 +---------- .../contrib/franka_pour/mdp/__init__.pyi | 164 ++++++++++++++++++ .../contrib/franka_pour/mdp/actions.py | 117 +------------ .../contrib/franka_pour/mdp/actions_cfg.py | 128 ++++++++++++++ .../test_franka_pour_spacemouse_teleop.py | 6 +- 6 files changed, 311 insertions(+), 212 deletions(-) create mode 100644 source/isaaclab_tasks/changelog.d/franka-pour-early-pxr-import.rst create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/__init__.pyi create mode 100644 source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/actions_cfg.py diff --git a/source/isaaclab_tasks/changelog.d/franka-pour-early-pxr-import.rst b/source/isaaclab_tasks/changelog.d/franka-pour-early-pxr-import.rst new file mode 100644 index 000000000000..68b493707ae7 --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/franka-pour-early-pxr-import.rst @@ -0,0 +1,6 @@ +Fixed +^^^^^ + +* Fixed Franka pour env-config construction importing USD ``pxr`` before Kit launch by + splitting pour MDP action configs into :mod:`isaaclab_tasks.contrib.franka_pour.mdp.actions_cfg` + and converting the pour MDP package to lazy exports. diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/__init__.py index c1fc60d3e6c4..3c3a6d3a79f7 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/__init__.py @@ -5,104 +5,6 @@ """MDP terms for the Franka pour task (grasp a dynamic cup of MPM media and pour).""" -from isaaclab.envs.mdp import ( # noqa: F401 - AbsBinaryJointPositionActionCfg, - BinaryJointPositionActionCfg, - DifferentialInverseKinematicsActionCfg, - JointPositionActionCfg, - RelativeJointPositionActionCfg, - action_l2, - action_rate_l2, - joint_pos_rel, - joint_vel_l2, - joint_vel_rel, - last_action, - time_out, -) +from isaaclab.utils.module import lazy_export -from .actions import ( # noqa: F401 - CurriculumGripperPositionAction, - CurriculumGripperPositionActionCfg, - CurriculumJointPositionAction, - CurriculumJointPositionActionCfg, - TrajectoryJointPositionAction, - TrajectoryJointPositionActionCfg, -) -from .curriculums import PourCurriculum # noqa: F401 -from .events import reset_pour_scene # noqa: F401 -from .observations import ( # noqa: F401 - arm_reference_error_obs, - arm_reference_phase_obs, - cup_pose_obs, - cup_to_target_obs, - cup_velocity_obs, - ee_pose_obs, - finger_position_obs, - finger_velocity_obs, - grasp_to_tcp_quat_obs, - gripper_contact_obs, - gripper_target_obs, - gripper_width_obs, - held_delivery_history_obs, - lost_grasp_dwell_obs, - particle_fractions_obs, - particle_transfer_obs, - pour_target_fraction_obs, - success_dwell_obs, - target_position_c_obs, - target_pose_obs, - tcp_pose_obs, - tcp_to_grasp_obs, - tcp_to_grasp_position_c_obs, - time_remaining_obs, - trajectory_status_obs, -) -from .rewards import ( # noqa: F401 - AlignProgress, - ApproachProgress, - GraspLiftProgress, - HeldDeliveryProgress, - LiftProgress, - NewlyDeliveredParticles, - NewlySpilledParticles, - PourReferenceProgress, - PourTaskProgress, - PourTiltProgress, - align_command_progress, - align_cup_over_target, - finite_joint_velocity_l2, - grasp_cup, - lift_command_progress, - lift_cup, - media_target_distance_tanh, - particles_in_source, - particles_in_target, - pour_success_bonus, - sustained_pour_success, - reach_cup, - spilled_particles, - terminal_failure, - tcp_cup_distance_tanh, - tilt_command_progress, - tilt_over_target, -) -from .reset_dataset import ( # noqa: F401 - PourResetDatasetCurriculum, - reset_dataset_difficulty, -) -from .reset_mixture import ( # noqa: F401 - RESET_MIXTURE_REGION_NAMES, - RESET_MIXTURE_STAGE_NAMES, - PourResetMixture, -) -from .terminations import ( # noqa: F401 - excessive_spill, - extreme_rigid_state, - immediate_pour_success, - lost_lifted_grasp, - nonterminating_stable_pour_success, - nonfinite_failure, - particle_out_of_bounds, - stable_pour_success, - unsuccessful_time_out, -) +lazy_export() diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/__init__.pyi b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/__init__.pyi new file mode 100644 index 000000000000..5431210b2011 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/__init__.pyi @@ -0,0 +1,164 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +__all__ = [ + "AlignProgress", + "ApproachProgress", + "CurriculumGripperPositionAction", + "CurriculumGripperPositionActionCfg", + "CurriculumJointPositionAction", + "CurriculumJointPositionActionCfg", + "GraspLiftProgress", + "HeldDeliveryProgress", + "LiftProgress", + "NewlyDeliveredParticles", + "NewlySpilledParticles", + "PourCurriculum", + "PourReferenceProgress", + "PourResetDatasetCurriculum", + "PourResetMixture", + "PourTaskProgress", + "PourTiltProgress", + "RESET_MIXTURE_REGION_NAMES", + "RESET_MIXTURE_STAGE_NAMES", + "TrajectoryJointPositionAction", + "TrajectoryJointPositionActionCfg", + "align_command_progress", + "align_cup_over_target", + "arm_reference_error_obs", + "arm_reference_phase_obs", + "cup_pose_obs", + "cup_to_target_obs", + "cup_velocity_obs", + "ee_pose_obs", + "excessive_spill", + "extreme_rigid_state", + "finite_joint_velocity_l2", + "finger_position_obs", + "finger_velocity_obs", + "grasp_cup", + "grasp_to_tcp_quat_obs", + "gripper_contact_obs", + "gripper_target_obs", + "gripper_width_obs", + "held_delivery_history_obs", + "immediate_pour_success", + "lift_command_progress", + "lift_cup", + "lost_grasp_dwell_obs", + "lost_lifted_grasp", + "media_target_distance_tanh", + "nonfinite_failure", + "nonterminating_stable_pour_success", + "particle_fractions_obs", + "particle_out_of_bounds", + "particle_transfer_obs", + "particles_in_source", + "particles_in_target", + "pour_success_bonus", + "pour_target_fraction_obs", + "reach_cup", + "reset_dataset_difficulty", + "reset_pour_scene", + "spilled_particles", + "stable_pour_success", + "success_dwell_obs", + "sustained_pour_success", + "target_position_c_obs", + "target_pose_obs", + "tcp_cup_distance_tanh", + "tcp_pose_obs", + "tcp_to_grasp_obs", + "tcp_to_grasp_position_c_obs", + "terminal_failure", + "tilt_command_progress", + "tilt_over_target", + "time_remaining_obs", + "trajectory_status_obs", + "unsuccessful_time_out", +] + +from .actions import ( + CurriculumGripperPositionAction, + CurriculumJointPositionAction, + TrajectoryJointPositionAction, +) +from .actions_cfg import ( + CurriculumGripperPositionActionCfg, + CurriculumJointPositionActionCfg, + TrajectoryJointPositionActionCfg, +) +from .curriculums import PourCurriculum +from .events import reset_pour_scene +from .observations import ( + arm_reference_error_obs, + arm_reference_phase_obs, + cup_pose_obs, + cup_to_target_obs, + cup_velocity_obs, + ee_pose_obs, + finger_position_obs, + finger_velocity_obs, + grasp_to_tcp_quat_obs, + gripper_contact_obs, + gripper_target_obs, + gripper_width_obs, + held_delivery_history_obs, + lost_grasp_dwell_obs, + particle_fractions_obs, + particle_transfer_obs, + pour_target_fraction_obs, + success_dwell_obs, + target_position_c_obs, + target_pose_obs, + tcp_pose_obs, + tcp_to_grasp_obs, + tcp_to_grasp_position_c_obs, + time_remaining_obs, + trajectory_status_obs, +) +from .reset_dataset import PourResetDatasetCurriculum, reset_dataset_difficulty +from .reset_mixture import RESET_MIXTURE_REGION_NAMES, RESET_MIXTURE_STAGE_NAMES, PourResetMixture +from .rewards import ( + AlignProgress, + ApproachProgress, + GraspLiftProgress, + HeldDeliveryProgress, + LiftProgress, + NewlyDeliveredParticles, + NewlySpilledParticles, + PourReferenceProgress, + PourTaskProgress, + PourTiltProgress, + align_command_progress, + align_cup_over_target, + finite_joint_velocity_l2, + grasp_cup, + lift_command_progress, + lift_cup, + media_target_distance_tanh, + particles_in_source, + particles_in_target, + pour_success_bonus, + reach_cup, + spilled_particles, + sustained_pour_success, + tcp_cup_distance_tanh, + terminal_failure, + tilt_command_progress, + tilt_over_target, +) +from .terminations import ( + excessive_spill, + extreme_rigid_state, + immediate_pour_success, + lost_lifted_grasp, + nonfinite_failure, + nonterminating_stable_pour_success, + particle_out_of_bounds, + stable_pour_success, + unsuccessful_time_out, +) +from isaaclab.envs.mdp import * diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/actions.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/actions.py index 20cc8109a077..6cb1b6598899 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/actions.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/actions.py @@ -9,14 +9,17 @@ import math from collections.abc import Sequence -from dataclasses import MISSING import torch -from isaaclab.envs.mdp.actions.actions_cfg import JointPositionActionCfg from isaaclab.envs.mdp.actions.joint_actions import JointPositionAction -from isaaclab.managers import ActionTerm, ActionTermCfg -from isaaclab.utils.configclass import configclass +from isaaclab.managers import ActionTerm + +from .actions_cfg import ( + CurriculumGripperPositionActionCfg, + CurriculumJointPositionActionCfg, + TrajectoryJointPositionActionCfg, +) _GRIPPER_POSITION_TOLERANCE = 1.0e-6 @@ -157,28 +160,6 @@ def reset(self, env_ids: Sequence[int] | torch.Tensor | slice | None = None) -> self._previous_target[selected] = self.action_offset[selected] -@configclass -class CurriculumJointPositionActionCfg(JointPositionActionCfg): - """Configuration for :class:`CurriculumJointPositionAction`.""" - - alpha: float = 0.2 - """Weight of the current joint target in the exponential moving average.""" - - reference_target: tuple[float, ...] = () - """Validated absolute joint target used by the early reference projection.""" - - project_reference_through_stage: int = -1 - """Last curriculum stage projected onto the reset-to-reference segment, or ``-1`` to disable.""" - - reference_action_magnitude: float = 1.0 - """Policy-space scalar magnitude that commands the complete reference segment.""" - - reference_action_index: int = 0 - """Fixed policy-action coordinate used as the early-stage reference phase.""" - - class_type: type[CurriculumJointPositionAction] = CurriculumJointPositionAction - - class TrajectoryJointPositionAction(ActionTerm): """Filtered joint-position residuals around a monotonic per-environment reference trajectory. @@ -553,43 +534,6 @@ def reset(self, env_ids: Sequence[int] | torch.Tensor | slice | None = None) -> self._raw_actions[selected] = 0.0 -@configclass -class TrajectoryJointPositionActionCfg(ActionTermCfg): - """Configuration for :class:`TrajectoryJointPositionAction`.""" - - joint_names: list[str] = MISSING - preserve_order: bool = False - waypoint_count: int = 6 - residual_scale: float | tuple[float, ...] = 0.05 - alpha: float = 0.2 - """Exponential smoothing weight applied only to policy joint residuals.""" - phase_rate: float = 1.0 / 3.5 - approach_phase_rate: float = 1.0 / 3.5 - """Phase rate before the grasp waypoint, kept slower for a contact-safe approach [1/s].""" - transport_phase_rate: float = 0.5 - """Phase rate after a validated grasp and before the receiver-aligned pour [1/s].""" - waypoint_phases: tuple[float, ...] = (0.0, 0.12, 0.24, 0.40, 0.62, 1.0) - approach_waypoint: int = 1 - grasp_waypoint: int = 2 - lift_waypoint: int = 3 - align_waypoint: int = 4 - grasp_gate_stage: int = 3 - approach_max_lateral_distance: float = 0.01 - """Maximum cross-track TCP error perpendicular to the grasp-approach axis [m].""" - approach_max_joint_error: float = 0.08 - approach_dwell_steps: int = 10 - """Consecutive centered, stationary steps required before the guarded grasp approach.""" - approach_max_linear_velocity: float = 0.01 - approach_max_angular_velocity: float = 0.1 - align_max_distance: float = 0.06 - """Maximum source-grasp-point error from the receiver-side pour pose [m].""" - grasp_dwell_steps: int = 15 - grasp_max_tcp_distance: float = 0.01 - grasp_max_linear_velocity: float = 0.05 - grasp_max_angular_velocity: float = 0.5 - class_type: type[TrajectoryJointPositionAction] = TrajectoryJointPositionAction - - class CurriculumGripperPositionAction(ActionTerm): """Filtered symmetric finger-position command with residual, incremental, and binary modes.""" @@ -857,50 +801,3 @@ def reset(self, env_ids: Sequence[int] | torch.Tensor | slice | None = None) -> else: self._capture_unlocked[selected] = self._env.curriculum_stage[selected] < self._force_open_stage self._capture_dwell_count[selected] = 0 - - -@configclass -class CurriculumGripperPositionActionCfg(ActionTermCfg): - """Configuration for :class:`CurriculumGripperPositionAction`.""" - - joint_names: list[str] = MISSING - scale: float = 0.04 - """Per-finger residual or incremental delta per policy-action unit [m]; unused in binary mode.""" - alpha: float = 0.2 - """Interpolation weight applied to the selected finger target.""" - use_incremental_target: bool = False - """Whether actions increment the target by ``alpha * scale`` so zero action holds its position.""" - binary_threshold: float | None = None - """Optional threshold selecting filtered close/maximum targets; values below it close.""" - close_position: float = 0.0 - neutral_position: float = 0.025 - """Largest per-finger command accepted from the action [m].""" - open_position: float = 0.04 - default_position: float | None = None - """Per-finger residual-mode zero command and initial target [m]. ``None`` uses ``close_position``.""" - limit_to_preload: bool = True - """Whether task validation restricts action targets to the contact-safe preload interval.""" - contact_min_deflection: float = 0.001 - """Minimum settled position-drive deflection required on each finger [m].""" - contact_max_velocity: float = 0.005 - """Maximum absolute finger speed accepted as settled bilateral contact [m/s].""" - force_open_before_phase_stage: int = -1 - """First stage that holds the hand open during the approach phase, or ``-1`` to disable.""" - force_open_before_phase: float = 0.25 - """Reference phase below which configured approach stages force the hand open.""" - capture_max_lateral_distance: float = 0.005 - """Maximum cross-track TCP error perpendicular to the grasp-approach axis [m].""" - capture_max_vertical_distance: float = 0.008 - """Maximum absolute TCP error along the grasp-approach axis [m]. - - The field retains its historical name for configuration compatibility. - """ - capture_max_joint_error: float = 0.08 - """Maximum reference-to-physical arm-joint error that releases the interlock [rad].""" - capture_dwell_steps: int = 5 - """Consecutive centered, stationary steps required before finger closure is enabled.""" - capture_max_linear_velocity: float = 0.02 - """Maximum source-cup linear speed during capture qualification [m/s].""" - capture_max_angular_velocity: float = 0.2 - """Maximum source-cup angular speed during capture qualification [rad/s].""" - class_type: type[CurriculumGripperPositionAction] = CurriculumGripperPositionAction diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/actions_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/actions_cfg.py new file mode 100644 index 000000000000..d2176fbffe85 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/actions_cfg.py @@ -0,0 +1,128 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Kit-free configuration for Franka pour reset-relative arm and gripper actions.""" + +from __future__ import annotations + +from dataclasses import MISSING +from typing import TYPE_CHECKING + +from isaaclab.envs.mdp.actions.actions_cfg import JointPositionActionCfg +from isaaclab.managers import ActionTermCfg +from isaaclab.utils.configclass import configclass + +if TYPE_CHECKING: + from .actions import ( + CurriculumGripperPositionAction, + CurriculumJointPositionAction, + TrajectoryJointPositionAction, + ) + + +@configclass +class CurriculumJointPositionActionCfg(JointPositionActionCfg): + """Configuration for :class:`CurriculumJointPositionAction`.""" + + alpha: float = 0.2 + """Weight of the current joint target in the exponential moving average.""" + + reference_target: tuple[float, ...] = () + """Validated absolute joint target used by the early reference projection.""" + + project_reference_through_stage: int = -1 + """Last curriculum stage projected onto the reset-to-reference segment, or ``-1`` to disable.""" + + reference_action_magnitude: float = 1.0 + """Policy-space scalar magnitude that commands the complete reference segment.""" + + reference_action_index: int = 0 + """Fixed policy-action coordinate used as the early-stage reference phase.""" + + class_type: type[CurriculumJointPositionAction] | str = "{DIR}.actions:CurriculumJointPositionAction" + + +@configclass +class TrajectoryJointPositionActionCfg(ActionTermCfg): + """Configuration for :class:`TrajectoryJointPositionAction`.""" + + joint_names: list[str] = MISSING + preserve_order: bool = False + waypoint_count: int = 6 + residual_scale: float | tuple[float, ...] = 0.05 + alpha: float = 0.2 + """Exponential smoothing weight applied only to policy joint residuals.""" + phase_rate: float = 1.0 / 3.5 + approach_phase_rate: float = 1.0 / 3.5 + """Phase rate before the grasp waypoint, kept slower for a contact-safe approach [1/s].""" + transport_phase_rate: float = 0.5 + """Phase rate after a validated grasp and before the receiver-aligned pour [1/s].""" + waypoint_phases: tuple[float, ...] = (0.0, 0.12, 0.24, 0.40, 0.62, 1.0) + approach_waypoint: int = 1 + grasp_waypoint: int = 2 + lift_waypoint: int = 3 + align_waypoint: int = 4 + grasp_gate_stage: int = 3 + approach_max_lateral_distance: float = 0.01 + """Maximum cross-track TCP error perpendicular to the grasp-approach axis [m].""" + approach_max_joint_error: float = 0.08 + approach_dwell_steps: int = 10 + """Consecutive centered, stationary steps required before the guarded grasp approach.""" + approach_max_linear_velocity: float = 0.01 + approach_max_angular_velocity: float = 0.1 + align_max_distance: float = 0.06 + """Maximum source-grasp-point error from the receiver-side pour pose [m].""" + grasp_dwell_steps: int = 15 + grasp_max_tcp_distance: float = 0.01 + grasp_max_linear_velocity: float = 0.05 + grasp_max_angular_velocity: float = 0.5 + class_type: type[TrajectoryJointPositionAction] | str = "{DIR}.actions:TrajectoryJointPositionAction" + + +@configclass +class CurriculumGripperPositionActionCfg(ActionTermCfg): + """Configuration for :class:`CurriculumGripperPositionAction`.""" + + joint_names: list[str] = MISSING + scale: float = 0.04 + """Per-finger residual or incremental delta per policy-action unit [m]; unused in binary mode.""" + alpha: float = 0.2 + """Interpolation weight applied to the selected finger target.""" + use_incremental_target: bool = False + """Whether actions increment the target by ``alpha * scale`` so zero action holds its position.""" + binary_threshold: float | None = None + """Optional threshold selecting filtered close/maximum targets; values below it close.""" + close_position: float = 0.0 + neutral_position: float = 0.025 + """Largest per-finger command accepted from the action [m].""" + open_position: float = 0.04 + default_position: float | None = None + """Per-finger residual-mode zero command and initial target [m]. ``None`` uses ``close_position``.""" + limit_to_preload: bool = True + """Whether task validation restricts action targets to the contact-safe preload interval.""" + contact_min_deflection: float = 0.001 + """Minimum settled position-drive deflection required on each finger [m].""" + contact_max_velocity: float = 0.005 + """Maximum absolute finger speed accepted as settled bilateral contact [m/s].""" + force_open_before_phase_stage: int = -1 + """First stage that holds the hand open during the approach phase, or ``-1`` to disable.""" + force_open_before_phase: float = 0.25 + """Reference phase below which configured approach stages force the hand open.""" + capture_max_lateral_distance: float = 0.005 + """Maximum cross-track TCP error perpendicular to the grasp-approach axis [m].""" + capture_max_vertical_distance: float = 0.008 + """Maximum absolute TCP error along the grasp-approach axis [m]. + + The field retains its historical name for configuration compatibility. + """ + capture_max_joint_error: float = 0.08 + """Maximum reference-to-physical arm-joint error that releases the interlock [rad].""" + capture_dwell_steps: int = 5 + """Consecutive centered, stationary steps required before finger closure is enabled.""" + capture_max_linear_velocity: float = 0.02 + """Maximum source-cup linear speed during capture qualification [m/s].""" + capture_max_angular_velocity: float = 0.2 + """Maximum source-cup angular speed during capture qualification [rad/s].""" + class_type: type[CurriculumGripperPositionAction] | str = "{DIR}.actions:CurriculumGripperPositionAction" diff --git a/source/isaaclab_tasks/test/contrib/test_franka_pour_spacemouse_teleop.py b/source/isaaclab_tasks/test/contrib/test_franka_pour_spacemouse_teleop.py index 41b9c6ef0621..3062f3f5317f 100644 --- a/source/isaaclab_tasks/test/contrib/test_franka_pour_spacemouse_teleop.py +++ b/source/isaaclab_tasks/test/contrib/test_franka_pour_spacemouse_teleop.py @@ -113,10 +113,12 @@ def test_teleop_config_finalizes_without_trajectory_controller_or_rl_distributio import isaaclab_tasks # noqa: F401 - cfg = FrankaPourEnvCfg_TELEOP().finalize() + teleop_cfg = FrankaPourEnvCfg_TELEOP() + assert str(teleop_cfg.actions.arm_action.class_type).endswith(":CurriculumJointPositionAction") + cfg = teleop_cfg.finalize() task_spec = gym.spec("Isaac-Pour-Franka-Teleop-v0") - assert cfg.actions.arm_action.class_type.__name__ == "CurriculumJointPositionAction" + assert str(cfg.actions.arm_action.class_type).endswith(":CurriculumJointPositionAction") assert cfg.actions.gripper_action.force_open_before_phase_stage == -1 assert cfg.actions.gripper_action.limit_to_preload is False assert cfg.actions.gripper_action.default_position == pytest.approx(cfg.gripper_open_pos) From 5b4cb35dc3be5c034a356c982f18d29709ed535e Mon Sep 17 00:00:00 2001 From: HuiDong Chen Date: Mon, 27 Jul 2026 20:46:12 +0800 Subject: [PATCH 11/14] Fix Implicit MPM reset masks Adapt coupled and standalone Implicit MPM reset masks to the solver's world-mask contract, and skip unsupported selective shared-world resets. Add regression tests and changelog fragments. --- .../changelog.d/mpm-coupled-reset-mask.rst | 11 ++ .../isaaclab_contrib/coupling/coupler.py | 55 +++++- .../test/coupling/test_coupler.py | 38 ++++ .../changelog.d/mpm-world-mask-shape.rst | 8 + .../isaaclab_newton/physics/mpm_manager.py | 109 +++++++++++- .../test_newton_manager_abstraction.py | 39 ++++ .../test/rendering_test_utils.py | 168 +++++------------- 7 files changed, 307 insertions(+), 121 deletions(-) create mode 100644 source/isaaclab_contrib/changelog.d/mpm-coupled-reset-mask.rst create mode 100644 source/isaaclab_newton/changelog.d/mpm-world-mask-shape.rst diff --git a/source/isaaclab_contrib/changelog.d/mpm-coupled-reset-mask.rst b/source/isaaclab_contrib/changelog.d/mpm-coupled-reset-mask.rst new file mode 100644 index 000000000000..330490532731 --- /dev/null +++ b/source/isaaclab_contrib/changelog.d/mpm-coupled-reset-mask.rst @@ -0,0 +1,11 @@ +Fixed +^^^^^ + +* Fixed coupled Implicit MPM environment resets raising + ``ValueError: world_mask has shape ...`` or + ``RuntimeError: Masked reset cannot selectively clear grid-backed warm + starts`` when :meth:`~isaaclab_newton.physics.NewtonManager.reset_solver_state` + forwarded Isaac Lab's ``(world_count,)`` mask. MPM entry ``reset`` now receives + the ``(world_count + 1,)`` mask required by + :meth:`newton.solvers.SolverImplicitMPM.reset` (or skips selective shared-world + resets), while MJWarp entries keep the original parent mask. diff --git a/source/isaaclab_contrib/isaaclab_contrib/coupling/coupler.py b/source/isaaclab_contrib/isaaclab_contrib/coupling/coupler.py index f88fae66d9da..96bcfc901425 100644 --- a/source/isaaclab_contrib/isaaclab_contrib/coupling/coupler.py +++ b/source/isaaclab_contrib/isaaclab_contrib/coupling/coupler.py @@ -18,9 +18,14 @@ NewtonCollisionPipelineCfg, NewtonSolverCfg, ) -from isaaclab_newton.physics.mpm_manager import NewtonMPMManager +from isaaclab_newton.physics.mpm_manager import ( + NewtonMPMManager, + adapt_world_mask_for_implicit_mpm, + should_skip_implicit_mpm_masked_reset, +) from isaaclab_newton.physics.newton_manager import NewtonManager from newton import CollisionPipeline, Model, ModelBuilder, ShapeFlags +from newton.solvers import SolverImplicitMPM from newton.solvers.experimental.coupled import SolverCoupled, SolverCoupledADMM, SolverCoupledProxy from isaaclab.physics import PhysicsManager @@ -114,6 +119,54 @@ def _build_solver(cls, model: Model, solver_cfg: CouplerCfg) -> None: NewtonManager._supports_contact_sensors = False NewtonManager._needs_collision_pipeline = needs_collision_pipeline NewtonManager._supports_rigid_body_force_input = True + cls._install_implicit_mpm_reset_mask_adapter( + NewtonManager._solver, + [entry.config.name for entry in resolved_entries], + ) + + @classmethod + def _install_implicit_mpm_reset_mask_adapter(cls, solver: object, entry_names: list[str]) -> None: + """Adapt coupled-reset masks for Implicit MPM's world-mask contract. + + :class:`~newton.solvers.experimental.coupled.SolverCoupled` forwards the + parent ``(world_count,)`` mask to every entry. MJWarp expects that shape, + while :class:`~newton.solvers.SolverImplicitMPM` expects one extra bit for + global world ``-1``. Shared multi-world Implicit MPM + (``separate_worlds=False``) also rejects selective masks when clearing + grid-backed warm starts; those resets are skipped. Wrapping only MPM + entry ``reset`` methods keeps the coupled parent API unchanged. + """ + # SolverCoupled exposes sub-solvers via ``solver(name)`` (not ``entry_solver``). + get_entry_solver = getattr(solver, "solver", None) + if not callable(get_entry_solver): + return + + for entry_name in entry_names: + try: + entry_solver = get_entry_solver(entry_name) + except (KeyError, TypeError, ValueError): + continue + if not isinstance(entry_solver, SolverImplicitMPM): + continue + original_reset = entry_solver.reset + + def _reset( + state, + world_mask=None, + flags=None, + *, + _original_reset=original_reset, + _entry_solver=entry_solver, + ): + if should_skip_implicit_mpm_masked_reset(_entry_solver, world_mask): + return None + return _original_reset( + state, + world_mask=adapt_world_mask_for_implicit_mpm(_entry_solver, world_mask), + flags=flags, + ) + + entry_solver.reset = _reset # type: ignore[method-assign] @classmethod def _validate_config(cls, solver_cfg: CouplerCfg) -> None: diff --git a/source/isaaclab_contrib/test/coupling/test_coupler.py b/source/isaaclab_contrib/test/coupling/test_coupler.py index 4a9b06ee76e1..63da905d9792 100644 --- a/source/isaaclab_contrib/test/coupling/test_coupler.py +++ b/source/isaaclab_contrib/test/coupling/test_coupler.py @@ -864,3 +864,41 @@ def test_admm_build_auto_detects_symmetric_contact_pairs_by_default(monkeypatch) cfg.contact_pairs = [] solver = NewtonCouplerManager._build_admm_coupled_solver(model, entries, cfg) assert list(solver.coupling.contact_pairs) == [] + + +def test_install_implicit_mpm_reset_mask_adapter_pads_entry_masks(monkeypatch): + """Coupled MPM entry resets pad full masks and skip selective shared-world ones.""" + import warp as wp + + recorded: list[list[bool] | None] = [] + + class _FakeImplicitMPM: + def __init__(self): + self.model = SimpleNamespace(world_count=4) + self._separate_worlds = False + + def reset(self, state, world_mask=None, flags=None): + del state, flags + recorded.append(None if world_mask is None else world_mask.numpy().tolist()) + + monkeypatch.setattr(coupler, "SolverImplicitMPM", _FakeImplicitMPM) + mpm_solver = _FakeImplicitMPM() + + class _FakeCoupled: + def solver(self, name: str): + assert name == "mpm" + return mpm_solver + + NewtonCouplerManager._install_implicit_mpm_reset_mask_adapter(_FakeCoupled(), ["mpm"]) + + # Shared multi-world selective masks are skipped (no call into original reset). + mpm_solver.reset(object(), world_mask=wp.array([True, False, True, False], dtype=wp.bool, device="cpu")) + assert recorded == [] + + mpm_solver.reset(object(), world_mask=wp.array([True, True, True, True], dtype=wp.bool, device="cpu")) + assert recorded == [None] + + # With separate worlds, selective masks are forwarded as (N+1,). + mpm_solver._separate_worlds = True + mpm_solver.reset(object(), world_mask=wp.array([True, False, True, False], dtype=wp.bool, device="cpu")) + assert recorded == [None, [True, False, True, False, False]] diff --git a/source/isaaclab_newton/changelog.d/mpm-world-mask-shape.rst b/source/isaaclab_newton/changelog.d/mpm-world-mask-shape.rst new file mode 100644 index 000000000000..e0e6f45e161e --- /dev/null +++ b/source/isaaclab_newton/changelog.d/mpm-world-mask-shape.rst @@ -0,0 +1,8 @@ +Fixed +^^^^^ + +* Fixed :meth:`~isaaclab_newton.physics.NewtonManager.reset_solver_state` under + Implicit MPM rejecting Isaac Lab ``(world_count,)`` masks. The manager now + expands them to the ``(world_count + 1,)`` shape required by + :meth:`newton.solvers.SolverImplicitMPM.reset`, and skips selective masks when + shared multi-world Implicit MPM cannot clear grid-backed warm starts per world. diff --git a/source/isaaclab_newton/isaaclab_newton/physics/mpm_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/mpm_manager.py index 4cc6f6798270..cbb3b0a8c924 100644 --- a/source/isaaclab_newton/isaaclab_newton/physics/mpm_manager.py +++ b/source/isaaclab_newton/isaaclab_newton/physics/mpm_manager.py @@ -7,8 +7,9 @@ from __future__ import annotations +import numpy as np import warp as wp -from newton import BodyFlags, Contacts, Control, GeoType, Model, ModelBuilder, State +from newton import BodyFlags, Contacts, Control, GeoType, Model, ModelBuilder, State, StateFlags from newton.solvers import SolverImplicitMPM from warp.fem import TemporaryStore @@ -16,6 +17,84 @@ from .newton_manager import NewtonManager +def adapt_world_mask_for_implicit_mpm( + solver: SolverImplicitMPM, + world_mask: wp.array | None, +) -> wp.array | None: + """Adapt an Isaac Lab world-reset mask to Implicit MPM's mask contract. + + Isaac Lab and most Newton solvers use a per-world mask of shape + ``(world_count,)``. :meth:`SolverImplicitMPM.reset` instead expects + ``(world_count + 1,)``, where the trailing entry selects global objects + whose world index is ``-1``. + + Args: + solver: Implicit MPM solver whose model defines ``world_count``. + world_mask: Optional Isaac Lab / Newton mask of shape ``(world_count,)`` + or an already-adapted Implicit MPM mask of shape + ``(world_count + 1,)``. + + Returns: + ``None`` when no mask is provided or every world is selected (a full + reset, including global world ``-1``). Otherwise a boolean Warp array + of shape ``(world_count + 1,)`` with the global bit left ``False``. + + Raises: + ValueError: If ``world_mask`` has neither ``(world_count,)`` nor + ``(world_count + 1,)`` shape. + """ + if world_mask is None: + return None + + local_selected = _local_world_selection(solver, world_mask) + if bool(np.all(local_selected)): + # Full world selection is equivalent to an unmasked reset and also + # covers global (world index -1) particle/collider history. + return None + + world_count = int(solver.model.world_count) + if tuple(world_mask.shape) == (world_count + 1,): + return world_mask + + padded = np.zeros(world_count + 1, dtype=bool) + padded[:-1] = local_selected + return wp.array(padded, dtype=wp.bool, device=world_mask.device) + + +def should_skip_implicit_mpm_masked_reset( + solver: SolverImplicitMPM, + world_mask: wp.array | None, +) -> bool: + """Whether a masked Implicit MPM reset must be skipped. + + Shared multi-world Implicit MPM (``separate_worlds=False``) rejects selective + masks when clearing grid-backed warm starts. Matching + :meth:`NewtonMPMManager._reset_solver_internals`, Isaac Lab skips those + resets instead of raising. Full-world selection still proceeds as an + unmasked reset. + """ + if world_mask is None: + return False + if bool(getattr(solver, "_separate_worlds", False)) or int(solver.model.world_count) <= 1: + return False + return not bool(np.all(_local_world_selection(solver, world_mask))) + + +def _local_world_selection(solver: SolverImplicitMPM, world_mask: wp.array) -> np.ndarray: + """Return the per-world selection bits from an Isaac Lab or Implicit MPM mask.""" + world_count = int(solver.model.world_count) + shape = tuple(world_mask.shape) + selected = world_mask.numpy() + if shape == (world_count + 1,): + return selected[:-1] + if shape == (world_count,): + return selected + raise ValueError( + f"world_mask has shape {shape}, expected ({world_count},) or ({world_count + 1},) " + "for SolverImplicitMPM.reset." + ) + + def _make_solver_config(solver_cfg: MPMSolverCfg) -> SolverImplicitMPM.Config: """Build Newton's implicit MPM solver config from Isaac Lab's cfg.""" return SolverImplicitMPM.Config( @@ -161,6 +240,34 @@ def _reset_solver_internals(cls, world_mask: wp.array | None) -> None: world_mask: Per-world reset mask, ignored. """ + @classmethod + def reset_solver_state( + cls, + state: State | None = None, + world_mask: wp.array(dtype=wp.bool) | None = None, + flags: StateFlags | int | None = None, + ) -> None: + """Reset Implicit MPM history after simulation state is rewritten. + + Expands Isaac Lab's ``(world_count,)`` mask to the + ``(world_count + 1,)`` shape required by :meth:`SolverImplicitMPM.reset` + before delegating to :meth:`NewtonManager.reset_solver_state`. Shared + multi-world selective masks are skipped; see + :func:`should_skip_implicit_mpm_masked_reset`. + """ + if not isinstance(cls._solver, SolverImplicitMPM): + raise RuntimeError( + f"{cls.__name__}.reset_solver_state requires an active SolverImplicitMPM; " + f"got {type(cls._solver).__name__}." + ) + if should_skip_implicit_mpm_masked_reset(cls._solver, world_mask): + return + super().reset_solver_state( + state=state, + world_mask=adapt_world_mask_for_implicit_mpm(cls._solver, world_mask), + flags=flags, + ) + @classmethod def _solver_specific_clear(cls) -> None: """Reset MPM-specific class state on teardown. diff --git a/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py b/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py index 832440d0a03f..b8861074460e 100644 --- a/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py +++ b/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py @@ -632,6 +632,45 @@ def test_mpm_unsupported_cuda_graph_capture_uses_eager_execution(monkeypatch): assert NewtonManager._graph_capture_pending is False +def test_adapt_world_mask_for_implicit_mpm_pads_selective_masks(): + """Isaac Lab (world_count,) masks expand to Implicit MPM's (world_count + 1,) contract.""" + from isaaclab_newton.physics.mpm_manager import ( + adapt_world_mask_for_implicit_mpm, + should_skip_implicit_mpm_masked_reset, + ) + + solver = SimpleNamespace(model=SimpleNamespace(world_count=4), _separate_worlds=False) + world_mask = wp.array([True, False, True, False], dtype=wp.bool, device="cpu") + + assert should_skip_implicit_mpm_masked_reset(solver, world_mask) is True + adapted = adapt_world_mask_for_implicit_mpm(solver, world_mask) + + assert adapted is not None + assert adapted.numpy().tolist() == [True, False, True, False, False] + + solver._separate_worlds = True + assert should_skip_implicit_mpm_masked_reset(solver, world_mask) is False + + +def test_adapt_world_mask_for_implicit_mpm_full_selection_becomes_unmasked(): + """Selecting every world is equivalent to an unmasked Implicit MPM reset.""" + from isaaclab_newton.physics.mpm_manager import ( + adapt_world_mask_for_implicit_mpm, + should_skip_implicit_mpm_masked_reset, + ) + + solver = SimpleNamespace(model=SimpleNamespace(world_count=3), _separate_worlds=False) + world_mask = wp.array([True, True, True], dtype=wp.bool, device="cpu") + + assert should_skip_implicit_mpm_masked_reset(solver, world_mask) is False + assert adapt_world_mask_for_implicit_mpm(solver, world_mask) is None + assert adapt_world_mask_for_implicit_mpm(solver, None) is None + assert should_skip_implicit_mpm_masked_reset(solver, None) is False + + empty = wp.array([False, False, False], dtype=wp.bool, device="cpu") + assert should_skip_implicit_mpm_masked_reset(solver, empty) is True + + def test_cuda_graph_capture_uses_simulation_device(monkeypatch): """CUDA graph capture should use the simulation device instead of Warp's default device.""" from isaaclab.physics import PhysicsManager diff --git a/source/isaaclab_tasks/test/rendering_test_utils.py b/source/isaaclab_tasks/test/rendering_test_utils.py index 94ba412ad723..357cfd22cf7a 100644 --- a/source/isaaclab_tasks/test/rendering_test_utils.py +++ b/source/isaaclab_tasks/test/rendering_test_utils.py @@ -598,98 +598,6 @@ def _camera_output_to_pil_image(tensor: torch.Tensor, data_type: str) -> Image.I return Image.fromarray(ndarr) -def _pour_frame_save_dir() -> str | None: - """Return the pour-frame output directory when frame saving is enabled.""" - out_dir = os.environ.get("ISAAC_LAB_SAVE_STAGES") - if not out_dir: - return None - return os.path.join(out_dir, "pour_frames") - - -def _pour_frame_count() -> int: - """Return the number of zero-action pour steps to simulate.""" - raw = os.environ.get("ISAAC_LAB_POUR_FRAMES", "120") - try: - count = int(raw) - except ValueError as error: - raise ValueError(f"ISAAC_LAB_POUR_FRAMES must be an integer, got {raw!r}.") from error - if count < 0: - raise ValueError(f"ISAAC_LAB_POUR_FRAMES must be nonnegative, got {count}.") - return count - - -def maybe_save_camera_frames( - test_name: str, - physics_backend: str, - renderer: str, - data_type: str, - camera_outputs: dict[str, ProxyArray], - frame_index: int, - *, - frame_images: list[Image.Image] | None = None, -) -> list[Image.Image]: - """If ``ISAAC_LAB_SAVE_STAGES`` is set, save the current camera frame and accumulate GIF frames. - - Args: - test_name: Rendering test label used in output filenames. - physics_backend: Physics backend label. - renderer: Renderer nickname. - data_type: Camera data type under test. - camera_outputs: Camera sensor output dictionary. - frame_index: Zero-based frame index in the pour sequence. - frame_images: Optional accumulator for animated GIF assembly. - - Returns: - The updated frame-image accumulator (empty when saving is disabled). - """ - out_dir = _pour_frame_save_dir() - if out_dir is None or data_type not in camera_outputs: - return frame_images or [] - - output = camera_outputs[data_type] - tensor = output if isinstance(output, torch.Tensor) else output.torch - image = _camera_output_to_pil_image(tensor, data_type) - - os.makedirs(out_dir, exist_ok=True) - safe_test_name = test_name.replace("/", "_") - prefix = f"{safe_test_name}-{physics_backend}-{renderer}-{data_type}" - frame_path = os.path.join(out_dir, f"{prefix}-frame{frame_index:03d}.png") - image.save(frame_path, format="PNG") - - accumulated = list(frame_images or []) - accumulated.append(image) - return accumulated - - -def maybe_write_pour_frame_gif( - test_name: str, - physics_backend: str, - renderer: str, - data_type: str, - frame_images: list[Image.Image], -) -> None: - """Write an animated GIF from accumulated pour frames when saving is enabled.""" - if not frame_images: - return - - out_dir = _pour_frame_save_dir() - if out_dir is None: - return - - os.makedirs(out_dir, exist_ok=True) - safe_test_name = test_name.replace("/", "_") - prefix = f"{safe_test_name}-{physics_backend}-{renderer}-{data_type}" - gif_path = os.path.join(out_dir, f"{prefix}.gif") - frame_images[0].save( - gif_path, - save_all=True, - append_images=frame_images[1:], - duration=100, - loop=0, - ) - print(f"[ISAAC_LAB_SAVE_STAGES] wrote {gif_path} ({len(frame_images)} frames)") - - def _apply_overrides_to_env_cfg(env_cfg: Any, override_args: list[str]) -> Any: """Apply override args to env_cfg using parse_overrides and apply_overrides.""" from isaaclab_tasks.utils.hydra import apply_overrides, collect_presets, parse_overrides @@ -835,6 +743,35 @@ def _save_comparison_image(img: Image.Image, filename: str) -> str: return path +def _camera_outputs_to_pil(camera_outputs: dict[str, ProxyArray]) -> Image.Image: + """Convert the first camera AOV in ``camera_outputs`` to a displayable PIL image.""" + data_type, output = next(iter(camera_outputs.items())) + tensor = output if isinstance(output, torch.Tensor) else output.torch + condition = torch.logical_or(torch.isinf(tensor), torch.isnan(tensor)) + corrected = torch.where(condition, torch.zeros_like(tensor), tensor) + normalized = normalize_camera_output_for_display(corrected, data_type) + grid = make_camera_output_grid(normalized) + ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy() + return Image.fromarray(ndarr) + + +def _save_image_sequence_gif(frames: list[Image.Image], filename: str, duration_ms: int = 100) -> str: + """Save a PIL image sequence as an animated GIF under the current working directory (repo root).""" + if not frames: + raise ValueError("No frames to save as GIF.") + path = os.path.join(os.getcwd(), filename) + rgb_frames = [frame.convert("RGB") for frame in frames] + rgb_frames[0].save( + path, + save_all=True, + append_images=rgb_frames[1:], + duration=duration_ms, + loop=0, + ) + logger.info("Wrote GIF with %d frames to %s", len(frames), path) + return path + + def _format_bcompare_command(actual_path: str, golden_path: str) -> str: """Build a shell command that opens actual and golden images in Beyond Compare.""" return f"bcompare \\\n {actual_path} \\\n {golden_path}" @@ -2111,8 +2048,6 @@ def rendering_test_franka_pour( data_type: str, comparison_scores: list[dict], ) -> None: - _skip_if_newton_motion_vectors(physics_backend, data_type) - from isaaclab_tasks.contrib.franka_pour.pour_env import FrankaPourEnv env_cfg = _make_franka_pour_camera_env_cfg(data_type) @@ -2136,41 +2071,36 @@ def rendering_test_franka_pour( try: env = FrankaPourEnv(env_cfg) - env.sim._app_control_on_stop_handle = None + env.reset() - _maybe_disable_instancing_for_current_stage(physics_backend, renderer, data_type) + env.sim._app_control_on_stop_handle = None maybe_save_stage(test_name, physics_backend, renderer, data_type) zero_actions = torch.zeros(env.num_envs, env.action_manager.total_action_dim, device=env.device) - for _ in range(FRAMES_BEFORE_POURING_PARTICLES): + pour_frames = 180 + frames: list[Image.Image] = [] + for _ in range(pour_frames): env.step(zero_actions) - - camera_outputs = env.scene.sensors["tiled_camera"].data.output - validate_camera_outputs( - test_name, - physics_backend, - renderer, - camera_outputs, - max_different_pixels_percentage=MAX_DIFFERENT_PIXELS_PERCENTAGE_BY_ENV_NAME[test_name], - comparison_scores=comparison_scores, + frames.append(_camera_outputs_to_pil(env.scene.sensors["tiled_camera"].data.output)) + _save_image_sequence_gif( + frames, f"{test_name}-{physics_backend}-{renderer}-{data_type}.gif" ) # zero_actions = torch.zeros(env.num_envs, env.action_manager.total_action_dim, device=env.device) - # pour_frames = 40 - # frame_images: list[Image.Image] = [] - # for step_index in range(pour_frames): + # for _ in range(FRAMES_BEFORE_POURING_PARTICLES): # env.step(zero_actions) - # frame_images = maybe_save_camera_frames( - # test_name, - # physics_backend, - # renderer, - # data_type, - # env.scene.sensors["tiled_camera"].data.output, - # step_index, - # frame_images=frame_images, - # ) - # maybe_write_pour_frame_gif(test_name, physics_backend, renderer, data_type, frame_images) + + # camera_outputs = env.scene.sensors["tiled_camera"].data.output + # validate_camera_outputs( + # test_name, + # physics_backend, + # renderer, + # camera_outputs, + # max_different_pixels_percentage=MAX_DIFFERENT_PIXELS_PERCENTAGE_BY_ENV_NAME[test_name], + # comparison_scores=comparison_scores, + # ) + finally: if env is not None: env.close() From fa2c57201e4c16d04474dd6e47e1b541e7607e05 Mon Sep 17 00:00:00 2001 From: HuiDong Chen Date: Mon, 10 Aug 2026 20:17:47 +0800 Subject: [PATCH 12/14] Fix rebasing errors --- .../test_newton_manager_abstraction.py | 246 +----------------- .../contrib/franka_pour/pour_env.py | 4 +- .../franka_pour/reset_dataset_generator.py | 10 +- 3 files changed, 8 insertions(+), 252 deletions(-) diff --git a/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py b/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py index b8861074460e..2403e39bb72d 100644 --- a/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py +++ b/source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py @@ -31,6 +31,7 @@ import numpy as np import pytest import warp as wp +from isaaclab_newton.cloner import copy_newton_source_builder, newton_builder_world_hook from isaaclab_newton.physics import ( FeatherstoneSolverCfg, KaminoSolverCfg, @@ -703,251 +704,6 @@ def __exit__(self, exc_type, exc_value, traceback): assert NewtonManager._graph is captured_graph -def test_relaxed_cuda_graph_capture_prepares_solver_before_warmup(monkeypatch): - """The RTX-compatible path prepares solver resources before its eager allocation warmup.""" - import isaaclab_newton.physics.newton_manager as newton_manager_module - - events = [] - contacts = object() - solver = SimpleNamespace(prepare_graph_capture=lambda received: events.append(("prepare", received))) - fake_cudart = SimpleNamespace(cudaStreamCreateWithFlags=lambda *_args: 1) - - monkeypatch.setattr(newton_manager_module, "_cudart", fake_cudart) - monkeypatch.setattr(NewtonManager, "_solver", solver, raising=False) - monkeypatch.setattr(NewtonManager, "_contacts", contacts, raising=False) - monkeypatch.setattr(NewtonManager, "_is_all_graphable", classmethod(lambda cls: False)) - monkeypatch.setattr( - NewtonManager, - "_simulate_physics_only", - classmethod(lambda cls: events.append(("warmup", None))), - ) - monkeypatch.setattr(wp, "get_stream", lambda *_args, **_kwargs: object()) - monkeypatch.setattr(wp, "synchronize_stream", lambda *_args, **_kwargs: None) - - assert NewtonManager._capture_relaxed_graph("cpu") == (None, True) - assert events == [("prepare", contacts), ("warmup", None)] - - -@pytest.mark.parametrize("all_graphable", [True, False]) -def test_manager_checks_solver_status_after_graph_replay(monkeypatch, all_graphable): - """Asynchronous solver failures are inspected after either manager graph path replays.""" - from isaaclab.physics import PhysicsManager - - events = [] - mask = SimpleNamespace(zero_=lambda: None) - solver = SimpleNamespace(check_status=lambda: events.append("status")) - monkeypatch.setattr(PhysicsManager, "_sim", SimpleNamespace(is_playing=lambda: True), raising=False) - monkeypatch.setattr(PhysicsManager, "_cfg", SimpleNamespace(use_cuda_graph=True), raising=False) - monkeypatch.setattr(PhysicsManager, "_device", "cuda:0", raising=False) - monkeypatch.setattr(PhysicsManager, "_sim_time", 0.0, raising=False) - monkeypatch.setattr(NewtonManager, "_model_changes", set(), raising=False) - monkeypatch.setattr(NewtonManager, "_solver_reset_pending", False, raising=False) - monkeypatch.setattr(NewtonManager, "_graph_capture_pending", None, raising=False) - monkeypatch.setattr(NewtonManager, "_graph", object(), raising=False) - monkeypatch.setattr(NewtonManager, "_solver", solver, raising=False) - monkeypatch.setattr(NewtonManager, "_world_reset_mask", mask, raising=False) - monkeypatch.setattr(NewtonManager, "_fk_reset_mask", mask, raising=False) - monkeypatch.setattr(NewtonManager, "_needs_collision_pipeline", False, raising=False) - monkeypatch.setattr(NewtonManager, "_needs_fk_before_step", False, raising=False) - monkeypatch.setattr(NewtonManager, "_solver_dt", 1.0 / 120.0, raising=False) - monkeypatch.setattr(NewtonManager, "_num_substeps", 1, raising=False) - monkeypatch.setattr(NewtonManager, "_decimation", 1, raising=False) - monkeypatch.setattr(NewtonManager, "_adapter", None, raising=False) - monkeypatch.setattr(NewtonManager, "_post_actuator_callbacks", [], raising=False) - monkeypatch.setattr(NewtonManager, "_usdrt_stage", None, raising=False) - monkeypatch.setattr(NewtonManager, "_particle_visual_prims", {}, raising=False) - monkeypatch.setattr(NewtonManager, "_is_all_graphable", classmethod(lambda cls: all_graphable)) - monkeypatch.setattr(NewtonManager, "_reset_solver_internals", classmethod(lambda cls, world_mask: None)) - monkeypatch.setattr(wp, "capture_launch", lambda graph: events.append("replay")) - monkeypatch.setattr(NewtonManager, "_log_solver_debug", classmethod(lambda cls: events.append("debug"))) - - NewtonManager.step() - - assert events == ["replay", "status", "debug"] - - -@pytest.mark.parametrize("mode", ["standard", "relaxed"]) -@pytest.mark.parametrize("all_graphable", [True, False]) -def test_deferred_graph_capture_runs_after_reset_setup_then_replays_once(monkeypatch, mode, all_graphable): - """The first reset is consumed before capture, then the new graph advances exactly one step.""" - from isaaclab.physics import PhysicsManager - - events = [] - graph = object() - - class Mask: - def __init__(self, name): - self.name = name - - def zero_(self): - events.append(f"zero_{self.name}") - - adapter = SimpleNamespace(step=lambda *_args: events.append("actuator")) if not all_graphable else None - solver = SimpleNamespace(check_status=lambda: events.append("status")) - monkeypatch.setattr(PhysicsManager, "_sim", SimpleNamespace(is_playing=lambda: True), raising=False) - monkeypatch.setattr(PhysicsManager, "_cfg", SimpleNamespace(use_cuda_graph=True), raising=False) - monkeypatch.setattr(PhysicsManager, "_device", "cuda:0", raising=False) - monkeypatch.setattr(PhysicsManager, "_sim_time", 0.0, raising=False) - monkeypatch.setattr(NewtonManager, "_model_changes", set(), raising=False) - monkeypatch.setattr(NewtonManager, "_solver_reset_pending", True, raising=False) - monkeypatch.setattr(NewtonManager, "_graph_capture_pending", mode, raising=False) - monkeypatch.setattr(NewtonManager, "_graph", None, raising=False) - monkeypatch.setattr(NewtonManager, "_solver", solver, raising=False) - monkeypatch.setattr(NewtonManager, "_world_reset_mask", Mask("world"), raising=False) - monkeypatch.setattr(NewtonManager, "_fk_reset_mask", Mask("fk"), raising=False) - monkeypatch.setattr(NewtonManager, "_needs_collision_pipeline", True, raising=False) - monkeypatch.setattr(NewtonManager, "_needs_fk_before_step", False, raising=False) - monkeypatch.setattr(NewtonManager, "_solver_dt", 1.0 / 120.0, raising=False) - monkeypatch.setattr(NewtonManager, "_num_substeps", 1, raising=False) - monkeypatch.setattr(NewtonManager, "_decimation", 1, raising=False) - monkeypatch.setattr(NewtonManager, "_adapter", adapter, raising=False) - monkeypatch.setattr(NewtonManager, "_post_actuator_callbacks", [], raising=False) - monkeypatch.setattr(NewtonManager, "_usdrt_stage", None, raising=False) - monkeypatch.setattr(NewtonManager, "_particle_visual_prims", {}, raising=False) - monkeypatch.setattr(NewtonManager, "_is_all_graphable", classmethod(lambda cls: all_graphable)) - monkeypatch.setattr( - NewtonManager, - "_reset_solver_internals", - classmethod(lambda cls, world_mask: events.append("reset")), - ) - monkeypatch.setattr( - NewtonManager, - "_eval_fk", - classmethod(lambda cls, world_mask, fk_mask: events.append("fk")), - ) - monkeypatch.setattr( - NewtonManager, - "_capture_standard_graph", - classmethod(lambda cls, device: events.append("capture_standard") or graph), - ) - monkeypatch.setattr( - NewtonManager, - "_capture_relaxed_graph", - classmethod(lambda cls, device: events.append("capture_relaxed") or (graph, True)), - ) - monkeypatch.setattr(wp, "capture_launch", lambda captured: events.append("replay")) - monkeypatch.setattr(NewtonManager, "_log_solver_debug", classmethod(lambda cls: events.append("debug"))) - monkeypatch.setattr( - NewtonManager, - "_simulate_full", - classmethod(lambda cls: pytest.fail("captured step must not also execute eagerly")), - ) - monkeypatch.setattr( - NewtonManager, - "_simulate_physics_only", - classmethod(lambda cls: pytest.fail("captured step must not also execute eagerly")), - ) - - NewtonManager.step() - - expected = ["reset", "fk", "zero_world", "zero_fk"] - if not all_graphable: - expected.append("actuator") - expected.append(f"capture_{mode}") - if mode == "standard": - expected.extend(["replay", "status"]) - expected.append("debug") - assert events == expected - assert NewtonManager._graph_capture_pending is None - - -def test_reset_solver_state_resets_distinct_buffers_and_deduplicates_aliases(monkeypatch): - """Selective resets cannot revive stale history after a state-buffer swap.""" - calls = [] - state_0 = object() - state_1 = object() - world_mask = SimpleNamespace(zero_=lambda: None) - solver = SimpleNamespace(reset=lambda state, *, world_mask, flags: calls.append((state, world_mask, flags))) - monkeypatch.setattr(NewtonManager, "_solver", solver, raising=False) - monkeypatch.setattr(NewtonManager, "_model", SimpleNamespace(world_count=2), raising=False) - monkeypatch.setattr(NewtonManager, "_state_0", state_0, raising=False) - monkeypatch.setattr(NewtonManager, "_state_1", state_1, raising=False) - - NewtonManager.reset_solver_state(world_mask=world_mask, flags=17) - assert calls == [(state_1, world_mask, 17), (state_0, world_mask, 17)] - - calls.clear() - monkeypatch.setattr(NewtonManager, "_state_1", state_0, raising=False) - NewtonManager.reset_solver_state() - assert calls == [(state_0, None, None)] - - -def test_solver_internal_reset_is_event_gated_and_single_world_uses_full_reset(monkeypatch): - """Absent invalidation does no work; a dirty single world uses a full reset.""" - calls = [] - state = object() - state_1 = object() - world_mask = SimpleNamespace(zero_=lambda: None) - solver = SimpleNamespace(reset=lambda *args, **kwargs: calls.append((args, kwargs))) - monkeypatch.setattr(NewtonManager, "_solver", solver, raising=False) - monkeypatch.setattr(NewtonManager, "_model", SimpleNamespace(world_count=1), raising=False) - monkeypatch.setattr(NewtonManager, "_state_0", state, raising=False) - monkeypatch.setattr(NewtonManager, "_state_1", state_1, raising=False) - monkeypatch.setattr(NewtonManager, "_solver_reset_pending", False, raising=False) - monkeypatch.setattr(NewtonManager, "_world_reset_mask", world_mask, raising=False) - monkeypatch.setattr(NewtonManager, "_fk_reset_mask", SimpleNamespace(zero_=lambda: None), raising=False) - monkeypatch.setattr(NewtonManager, "_eval_fk", classmethod(lambda cls, *_args: None)) - monkeypatch.setattr(NewtonManager, "_usdrt_stage", None, raising=False) - - NewtonManager.forward() - assert calls == [] - - NewtonManager._solver_reset_pending = True - NewtonManager.forward() - assert calls == [((state,), {"world_mask": None, "flags": 0})] - assert NewtonManager._solver_reset_pending is False - - calls.clear() - NewtonManager.reset_solver_state(world_mask=world_mask, flags=17) - assert calls == [ - ((state_1,), {"world_mask": world_mask, "flags": 17}), - ((state,), {"world_mask": world_mask, "flags": 17}), - ] - - -def test_forward_publishes_reset_fk_to_fabric(monkeypatch): - """A public forward call must publish reset joint poses before returning.""" - events = [] - world_mask = SimpleNamespace(zero_=lambda: events.append("clear_world")) - fk_mask = SimpleNamespace(zero_=lambda: events.append("clear_fk")) - - monkeypatch.setattr(NewtonManager, "_world_reset_mask", world_mask, raising=False) - monkeypatch.setattr(NewtonManager, "_fk_reset_mask", fk_mask, raising=False) - monkeypatch.setattr(NewtonManager, "_solver_reset_pending", True, raising=False) - monkeypatch.setattr(NewtonManager, "_usdrt_stage", object(), raising=False) - monkeypatch.setattr( - NewtonManager, - "_reset_solver_internals", - classmethod(lambda cls, mask: events.append(("reset", mask))), - ) - monkeypatch.setattr( - NewtonManager, - "_eval_fk", - classmethod(lambda cls, worlds, articulations: events.append(("fk", worlds, articulations))), - ) - monkeypatch.setattr( - NewtonManager, - "_mark_transforms_dirty", - classmethod(lambda cls: events.append("mark_transforms")), - ) - monkeypatch.setattr( - NewtonManager, - "sync_transforms_to_usd", - classmethod(lambda cls: events.append("sync_transforms")), - ) - - NewtonManager.forward() - - assert events == [ - ("reset", world_mask), - ("fk", world_mask, fk_mask), - "mark_transforms", - "sync_transforms", - "clear_fk", - "clear_world", - ] - - def test_newton_builder_world_hook_is_scoped_and_preserves_existing_registration(monkeypatch): def existing(*args): pass diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/pour_env.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/pour_env.py index 6cb157a04c40..18037f1ebef4 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/pour_env.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/pour_env.py @@ -36,7 +36,7 @@ from isaaclab_newton.physics import NewtonManager import isaaclab.sim as sim_utils -from isaaclab.cloner import resolve_clone_plan_source +from isaaclab.cloner import query as cloner_query from isaaclab.envs import ManagerBasedRLEnv from isaaclab.utils import math as math_utils @@ -773,7 +773,7 @@ def _validate_loaded_reset_dataset(self, payload: dict) -> None: def _build_randomized_reset_bank(self) -> None: """Build a small Newton-IK bank for collision-safe randomized pre-grasp resets.""" plan = sim_utils.SimulationContext.instance().get_clone_plan() - resolved = resolve_clone_plan_source(self._robot.cfg.prim_path, plan) if plan is not None else None + resolved = cloner_query.path_to_source(plan, self._robot.cfg.prim_path) if plan is not None else None if resolved is None: raise RuntimeError(f"Could not resolve clone-plan source for {self._robot.cfg.prim_path!r}.") source_path = resolved[0] diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/reset_dataset_generator.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/reset_dataset_generator.py index b514f68f8649..f516a678aacc 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/reset_dataset_generator.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/reset_dataset_generator.py @@ -418,11 +418,11 @@ def _derive_tabletop_support_bounds( from pxr import Usd, UsdGeom import isaaclab.sim as sim_utils - from isaaclab.cloner import resolve_clone_plan_source + from isaaclab.cloner import query as cloner_query if source_env_path is None: plan = sim_utils.SimulationContext.instance().get_clone_plan() - resolved = resolve_clone_plan_source(env._robot.cfg.prim_path, plan) if plan is not None else None + resolved = cloner_query.path_to_source(plan, env._robot.cfg.prim_path) if plan is not None else None if resolved is None: raise RuntimeError(f"Could not resolve clone-plan source for {env._robot.cfg.prim_path!r}.") source_env_path = resolved[0] @@ -1157,10 +1157,10 @@ def _build_ik_context(self) -> None: from isaaclab_newton.ik.newton_ik_solver_cfg import NewtonIKSolverCfg import isaaclab.sim as sim_utils - from isaaclab.cloner import resolve_clone_plan_source + from isaaclab.cloner import query as cloner_query plan = sim_utils.SimulationContext.instance().get_clone_plan() - resolved = resolve_clone_plan_source(self.env._robot.cfg.prim_path, plan) if plan is not None else None + resolved = cloner_query.path_to_source(plan, self.env._robot.cfg.prim_path) if plan is not None else None if resolved is None: raise RuntimeError(f"Could not resolve clone-plan source for {self.env._robot.cfg.prim_path!r}.") source_builder = copy_newton_source_builder(resolved[0]) @@ -1172,7 +1172,7 @@ def _build_ik_context(self) -> None: "/Table/" in str(label) or str(label).endswith("/Table") for label in self._prototype_builder.shape_label ): table_prim_path = self.env.scene["table"].cfg.prim_path - table_resolved = resolve_clone_plan_source(table_prim_path, plan) + table_resolved = cloner_query.path_to_source(plan, table_prim_path) if plan is not None else None if table_resolved is None: raise RuntimeError(f"Could not resolve clone-plan source for {table_prim_path!r}.") table_builder = copy_newton_source_builder(table_resolved[0]) From c162eaef993fca8e1601bfda361a775f52a34b86 Mon Sep 17 00:00:00 2001 From: HuiDong Chen Date: Mon, 10 Aug 2026 20:31:54 +0800 Subject: [PATCH 13/14] uniform amplication of motion vectors (temp change) --- source/isaaclab/isaaclab/utils/images.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/source/isaaclab/isaaclab/utils/images.py b/source/isaaclab/isaaclab/utils/images.py index bb02548d55cf..cbf4635b9466 100644 --- a/source/isaaclab/isaaclab/utils/images.py +++ b/source/isaaclab/isaaclab/utils/images.py @@ -113,9 +113,7 @@ def normalize_camera_output_for_display(tensor: torch.Tensor, data_type: str) -> # peak magnitude to map into [-1, 1], remap to [0, 1], and pack the two channels into an RGB image # (u -> R, v -> G, unused B -> 0) so the result can be composed into a grid and saved as an image. uv = normalized[..., :2] - max_mag = uv.abs().max() - if max_mag > 0: - uv = uv / max_mag + uv = uv * 1000 uv = (uv + 1.0) * 0.5 blue = torch.zeros_like(uv[..., :1]) normalized = torch.cat([uv, blue], dim=-1) From e7e30fbed31a07746a4499c265920f613b9950a1 Mon Sep 17 00:00:00 2001 From: HuiDong Chen Date: Wed, 12 Aug 2026 07:20:34 +0800 Subject: [PATCH 14/14] clamp motion vector to [-1, 1] --- source/isaaclab/isaaclab/utils/images.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/source/isaaclab/isaaclab/utils/images.py b/source/isaaclab/isaaclab/utils/images.py index cbf4635b9466..4e5c9361216e 100644 --- a/source/isaaclab/isaaclab/utils/images.py +++ b/source/isaaclab/isaaclab/utils/images.py @@ -112,8 +112,7 @@ def normalize_camera_output_for_display(tensor: torch.Tensor, data_type: str) -> # Motion vectors are per-pixel (u, v) offsets that can be positive or negative. Normalize by the # peak magnitude to map into [-1, 1], remap to [0, 1], and pack the two channels into an RGB image # (u -> R, v -> G, unused B -> 0) so the result can be composed into a grid and saved as an image. - uv = normalized[..., :2] - uv = uv * 1000 + uv = normalized[..., :2].clamp(-1.0, 1.0) uv = (uv + 1.0) * 0.5 blue = torch.zeros_like(uv[..., :1]) normalized = torch.cat([uv, blue], dim=-1)