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/isaaclab/utils/images.py b/source/isaaclab/isaaclab/utils/images.py index bb02548d55cf..4e5c9361216e 100644 --- a/source/isaaclab/isaaclab/utils/images.py +++ b/source/isaaclab/isaaclab/utils/images.py @@ -112,10 +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] - max_mag = uv.abs().max() - if max_mag > 0: - uv = uv / max_mag + 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) 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/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..45980c890693 --- /dev/null +++ b/source/isaaclab_newton/changelog.d/mpm-task-support.minor.rst @@ -0,0 +1,17 @@ +Added +^^^^^ + +* Added scoped Newton builder-world hooks and independent clone-source builder + copies for tasks that extend replicated Newton worlds. +* 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/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/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/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/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..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/isaaclab_newton/physics/newton_manager.py b/source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py index fdcaf956a403..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 @@ -1588,22 +1589,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() @@ -2032,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/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() 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..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, @@ -632,6 +633,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 @@ -664,6 +704,70 @@ def __exit__(self, exc_type, exc_value, traceback): assert NewtonManager._graph is captured_graph +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) + + with pytest.raises(RuntimeError, match="/World/missing.*Available: /World/known"): + copy_newton_source_builder("/World/missing") + + # --------------------------------------------------------------------------- # Manager state-refresh boundaries (no SimulationContext required) # --------------------------------------------------------------------------- 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_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-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/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..3c3a6d3a79f7 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/__init__.py @@ -0,0 +1,10 @@ +# 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.utils.module import lazy_export + +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 new file mode 100644 index 000000000000..6cb1b6598899 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/actions.py @@ -0,0 +1,803 @@ +# 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 + +import torch + +from isaaclab.envs.mdp.actions.joint_actions import JointPositionAction +from isaaclab.managers import ActionTerm + +from .actions_cfg import ( + CurriculumGripperPositionActionCfg, + CurriculumJointPositionActionCfg, + TrajectoryJointPositionActionCfg, +) + +_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] + + +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 + + +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 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/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..18037f1ebef4 --- /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 query as cloner_query +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 = 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] + 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..dd2611b1fb99 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/pour_env_cfg.py @@ -0,0 +1,2133 @@ +# 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$" +# 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", + "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( + 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=[ + ROBOT_BODY_LABEL_PATTERN, + SOURCE_CUP_BODY_LABEL_PATTERN, + TARGET_CUP_BODY_LABEL_PATTERN, + ], + 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", + ), + 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=[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 + # 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) + 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..f516a678aacc --- /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 query as cloner_query + + if source_env_path is None: + plan = sim_utils.SimulationContext.instance().get_clone_plan() + 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] + + 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 query as cloner_query + + plan = sim_utils.SimulationContext.instance().get_clone_plan() + 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]) + 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 = 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]) + 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..9bcce7f5c8b0 --- /dev/null +++ b/source/isaaclab_tasks/test/contrib/test_franka_pour_env_cfg.py @@ -0,0 +1,2094 @@ +# 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, TerminationTermCfg +from isaaclab.sim.schemas import MassCfg, UsdPhysicsCollisionCfg, UsdPhysicsRigidBodyCfg +from isaaclab.sim.spawners.materials import RigidBodyMaterialBaseCfg + +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 +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 _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 cfg.sim.physics.collision_cfg.soft_contact_max == 0 + assert arm.include_static_shapes is True + 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 == [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.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 == [ + 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) + 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_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.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 + + +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 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_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 = NewtonCouplerManager._resolve_entry(model, media) + + 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", "NewtonCouplerManager"} + } + 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..3062f3f5317f --- /dev/null +++ b/source/isaaclab_tasks/test/contrib/test_franka_pour_spacemouse_teleop.py @@ -0,0 +1,130 @@ +# 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 + + 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 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) + 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/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/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..486f9aa45084 --- /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: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 new file mode 100644 index 000000000000..c947d6686e26 --- /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: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 new file mode 100644 index 000000000000..053a936eb860 --- /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: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 new file mode 100644 index 000000000000..a45a7cd2c487 --- /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: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 new file mode 100644 index 000000000000..67055ebab7b9 --- /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: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 new file mode 100644 index 000000000000..1f3b4d7cdc6d --- /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:00cbc54fafb70e6ce512d9e28e81975e063daf361b4496d51ad930c00df6dcdc +size 16259 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 cccaa7eb359f..357cfd22cf7a 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. @@ -586,6 +588,16 @@ 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 _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 @@ -678,6 +690,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. @@ -710,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}" @@ -1090,10 +1152,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): @@ -1908,3 +1967,144 @@ 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 + # 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 = 6 + 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: + from isaaclab_tasks.contrib.franka_pour.pour_env import FrankaPourEnv + + env_cfg = _make_franka_pour_camera_env_cfg(data_type) + + # 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 + + if renderer == "ovrtx_renderer": + _redirect_ovrtx_renderer_log_to_stdout(env_cfg) + + test_name = "franka_pour" + env = None + + FRAMES_BEFORE_POURING_PARTICLES = 40 + + try: + env = FrankaPourEnv(env_cfg) + env.reset() + + 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) + pour_frames = 180 + frames: list[Image.Image] = [] + for _ in range(pour_frames): + env.step(zero_actions) + 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) + # 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, + # 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() + + # 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 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.