diff --git a/autodata_examples/envs/__init__.py b/autodata_examples/envs/__init__.py new file mode 100644 index 0000000..d829881 --- /dev/null +++ b/autodata_examples/envs/__init__.py @@ -0,0 +1,65 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Example simulation environments used by AutoData workflows.""" + +import argparse +from typing import Any + + +def register_environment_for_run( + *, + env_name: str, + enable_cameras: bool, + num_envs: int, + device: str, + seed: int, + presets: str | None = None, +) -> dict[str, dict[str, Any]]: + """Register the AutoData-owned environment selected for the current run. + + Args: + env_name: Gym ID of the environment selected for this run. + enable_cameras: Whether Arena environments should include their camera sensors. + num_envs: Number of parallel environments in this run. + device: Simulation device for this run. + seed: Environment seed for this run. + presets: Optional Arena physics backend preset; None preserves the environment default. + + Returns: + Gym constructor kwargs keyed by registered environment ID. + """ + + from .isaac_lab_arena import build_and_register_arena_environment, is_arena_environment + + if not is_arena_environment(env_name): + return {} + return build_and_register_arena_environment( + enable_cameras=enable_cameras, + num_envs=num_envs, + device=device, + seed=seed, + presets=presets, + ) + + +def register_environments() -> list[str]: + """External callback used by unmodified Isaac Lab scripts. + + Isaac Lab invokes this function without arguments after starting Isaac Sim. The + selected ``--task`` determines whether the callback needs to compile an AutoData + Arena environment or only import AutoData's regular Isaac Lab registrations. + + Returns: + Command-line arguments not consumed by Arena environment registration. + """ + + from .isaac_lab_arena import is_arena_environment, register_environment_from_cli + + parser = argparse.ArgumentParser(add_help=False) + parser.add_argument("--task", type=str) + args, _ = parser.parse_known_args() + env_name = args.task.split(":")[-1] if args.task is not None else None + if env_name is not None and not is_arena_environment(env_name): + return [] + return register_environment_from_cli() diff --git a/autodata_examples/envs/isaac_lab_arena/__init__.py b/autodata_examples/envs/isaac_lab_arena/__init__.py new file mode 100644 index 0000000..d33c1b6 --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from .registration import build_and_register_arena_environment, is_arena_environment, register_environment_from_cli + +__all__ = ["build_and_register_arena_environment", "is_arena_environment", "register_environment_from_cli"] diff --git a/autodata_examples/envs/isaac_lab_arena/franka_rope/__init__.py b/autodata_examples/envs/isaac_lab_arena/franka_rope/__init__.py new file mode 100644 index 0000000..2574837 --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/__init__.py @@ -0,0 +1,14 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from __future__ import annotations + +from . import embodiment as _embodiment # noqa: F401 +from .environment import FRANKA_ROPE_ARENA_ENV_ID, FrankaRopeArenaEnvironment, FrankaRopeArenaEnvironmentCfg + +__all__ = [ + "FRANKA_ROPE_ARENA_ENV_ID", + "FrankaRopeArenaEnvironment", + "FrankaRopeArenaEnvironmentCfg", +] diff --git a/autodata_examples/envs/isaac_lab_arena/franka_rope/assets/Rope.usd b/autodata_examples/envs/isaac_lab_arena/franka_rope/assets/Rope.usd new file mode 100644 index 0000000..fac1654 Binary files /dev/null and b/autodata_examples/envs/isaac_lab_arena/franka_rope/assets/Rope.usd differ diff --git a/autodata_examples/envs/isaac_lab_arena/franka_rope/embodiment.py b/autodata_examples/envs/isaac_lab_arena/franka_rope/embodiment.py new file mode 100644 index 0000000..e75c49a --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/embodiment.py @@ -0,0 +1,186 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from __future__ import annotations + +from collections.abc import Callable + +import isaaclab.envs.mdp as mdp_isaac_lab +import isaaclab.sim as sim_utils +from isaaclab.managers import EventTermCfg as EventTerm +from isaaclab.managers import ObservationGroupCfg as ObsGroup +from isaaclab.managers import ObservationTermCfg as ObsTerm +from isaaclab.managers import SceneEntityCfg +from isaaclab.sensors import CameraCfg +from isaaclab.utils.configclass import configclass +from isaaclab_arena.assets.register import register_asset, register_retargeter +from isaaclab_arena.assets.retargeter_library import RetargetterBase +from isaaclab_arena.embodiments.common.arm_mode import ArmMode +from isaaclab_arena.embodiments.franka.franka import FrankaIKEmbodiment +from isaaclab_arena.utils.cameras import ArenaCameraCfg +from isaaclab_arena.utils.pose import Pose +from isaaclab_physx.renderers import IsaacRtxRendererCfg, IsaacRtxRendererGlobalSettingsCfg + +from . import mdp + +_FRANKA_ROPE_INIT_JOINT_POS: dict[str, float] = { + "panda_joint1": 0.0444, + "panda_joint2": -0.1894, + "panda_joint3": -0.1107, + "panda_joint4": -2.5148, + "panda_joint5": 0.0044, + "panda_joint6": 2.3775, + "panda_joint7": 2.43, + "panda_finger_joint.*": 0.04, +} + +_FRANKA_ROPE_INIT_JOINT_POSE = [ + 0.0444, + -0.1894, + -0.1107, + -2.5148, + 0.0044, + 2.3775, + 2.43, + 0.04, + 0.04, +] + + +@configclass +class FrankaRopeCameraCfg(ArenaCameraCfg): + """Wrist and fixed-view cameras matching the Isaac Lab rope environment.""" + + robot0_eye_in_hand_image = CameraCfg( + prim_path="{ENV_REGEX_NS}/Robot/panda_hand/robot0_eye_in_hand_image", + update_period=0.0, + height=128, + width=128, + data_types=["rgb"], + renderer_cfg=IsaacRtxRendererCfg(global_settings=IsaacRtxRendererGlobalSettingsCfg(antialiasing_mode="DLSS")), + spawn=sim_utils.PinholeCameraCfg( + focal_length=24.0, + focus_distance=400.0, + horizontal_aperture=20.955, + clipping_range=(0.1, 2.0), + ), + offset=CameraCfg.OffsetCfg( + pos=(0.13, 0.0, -0.15), + rot=(0.03701, 0.03701, -0.70614, -0.70614), + convention="ros", + ), + ) + + agentview_image = CameraCfg( + prim_path="{ENV_REGEX_NS}/agentview_image", + update_period=0.0, + height=128, + width=128, + data_types=["rgb"], + renderer_cfg=IsaacRtxRendererCfg(global_settings=IsaacRtxRendererGlobalSettingsCfg(antialiasing_mode="DLSS")), + spawn=sim_utils.PinholeCameraCfg( + focal_length=14.0, + focus_distance=400.0, + horizontal_aperture=20.955, + clipping_range=(0.1, 2.0), + ), + offset=CameraCfg.OffsetCfg( + pos=(0.5, 0.5, 0.4), + rot=(0.0, 0.38268, 0.92388, 0.0), + convention="opengl", + ), + ) + + +@configclass +class FrankaRopeObservationsCfg: + """Arena-standard Franka observations plus deformable rope nodal state.""" + + @configclass + class PolicyCfg(ObsGroup): + """Policy observations recorded into demonstration datasets.""" + + actions = ObsTerm(func=mdp_isaac_lab.last_action) + joint_pos = ObsTerm(func=mdp_isaac_lab.joint_pos_rel) + joint_vel = ObsTerm(func=mdp_isaac_lab.joint_vel_rel) + eef_pos = ObsTerm(func=mdp.ee_frame_pos) + eef_quat = ObsTerm(func=mdp.ee_frame_quat) + gripper_pos = ObsTerm(func=mdp.gripper_pos) + object_nodal_pos = ObsTerm(func=mdp.object_nodal_pos) + + def __post_init__(self) -> None: + self.enable_corruption = False + self.concatenate_terms = False + + policy: PolicyCfg = PolicyCfg() + + +@configclass +class FrankaRopeEventsCfg: + """Reset the Franka exactly as the equivalent Isaac Lab rope environment.""" + + reset_all = EventTerm(func=mdp_isaac_lab.reset_scene_to_default, mode="reset") + randomize_franka_joint_state = EventTerm( + func=mdp_isaac_lab.reset_joints_by_offset, + mode="reset", + params={ + "position_range": (-0.02, 0.02), + "velocity_range": (0.0, 0.0), + "asset_cfg": SceneEntityCfg("robot", joint_names=["panda_joint.*"]), + }, + ) + + +@register_asset +class FrankaRopeIKEmbodiment(FrankaIKEmbodiment): + """Arena Franka IK embodiment with rope-specific pose, observations, and cameras.""" + + name = "franka_rope_ik" + + def __init__( + self, + enable_cameras: bool = False, + initial_pose: Pose | None = None, + concatenate_observation_terms: bool = False, + arm_mode: ArmMode | None = None, + ) -> None: + super().__init__( + enable_cameras=enable_cameras, + initial_pose=initial_pose, + initial_joint_pose=list(_FRANKA_ROPE_INIT_JOINT_POSE), + concatenate_observation_terms=concatenate_observation_terms, + arm_mode=arm_mode, + ) + self.scene_config.robot.init_state.joint_pos = dict(_FRANKA_ROPE_INIT_JOINT_POS) + self.observation_config = FrankaRopeObservationsCfg() + self.observation_config.policy.concatenate_terms = concatenate_observation_terms + self.camera_config = FrankaRopeCameraCfg() + self.event_config = FrankaRopeEventsCfg() + self.reward_config = None + + +@register_retargeter +class FrankaRopeKeyboardRetargeter(RetargetterBase): + """Use the standard relative-pose keyboard pipeline for the rope embodiment.""" + + device = "keyboard" + embodiment = FrankaRopeIKEmbodiment.name + + def get_pipeline_builder(self, embodiment: object) -> Callable | None: + """Return no custom pipeline; Isaac Lab handles relative-pose keyboard input.""" + + del embodiment + + +@register_retargeter +class FrankaRopeSpaceMouseRetargeter(RetargetterBase): + """Use the standard relative-pose SpaceMouse pipeline for the rope embodiment.""" + + device = "spacemouse" + embodiment = FrankaRopeIKEmbodiment.name + + def get_pipeline_builder(self, embodiment: object) -> Callable | None: + """Return no custom pipeline; Isaac Lab handles relative-pose SpaceMouse input.""" + + del embodiment diff --git a/autodata_examples/envs/isaac_lab_arena/franka_rope/environment.py b/autodata_examples/envs/isaac_lab_arena/franka_rope/environment.py new file mode 100644 index 0000000..1781e57 --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/environment.py @@ -0,0 +1,107 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +from isaaclab.devices import DevicesCfg, Se3KeyboardCfg, Se3SpaceMouseCfg +from isaaclab_arena.assets.register import register_environment +from isaaclab_arena.environments.arena_environment_factory import ArenaEnvironmentCfg, ArenaEnvironmentFactory +from isaaclab_arena.utils.pose import Pose + +from .rope_asset import FrankaRopeAsset +from .task import FrankaRopeTask + +if TYPE_CHECKING: + from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment + +FRANKA_ROPE_ARENA_ENV_ID = "Isaac-Rope-Franka-Arena-IK-Rel-v0" + + +@dataclass +class FrankaRopeArenaEnvironmentCfg(ArenaEnvironmentCfg): + """Configure the AutoData-owned Arena Franka rope environment.""" + + enable_cameras: bool = False + embodiment: str = "franka_rope_ik" + teleop_device: str | None = None + + +def _configure_franka_rope_env(env_cfg: Any) -> Any: + """Apply rope timing and asset settings for PhysX.""" + + env_cfg.decimation = 5 + env_cfg.sim.dt = 0.01 + env_cfg.sim.render_interval = 2 + env_cfg.num_rerenders_on_reset = 1 + env_cfg.scene.replicate_physics = False + env_cfg.teleop_devices = DevicesCfg( + devices={ + "keyboard": Se3KeyboardCfg( + pos_sensitivity=0.05, + rot_sensitivity=0.2, + ), + "spacemouse": Se3SpaceMouseCfg( + pos_sensitivity=0.2, + rot_sensitivity=0.5, + ), + }, + ) + return env_cfg + + +@register_environment +class FrankaRopeArenaEnvironment(ArenaEnvironmentFactory[FrankaRopeArenaEnvironmentCfg]): + """Compose the rope scene, Franka embodiment, and rope task through Arena.""" + + name = FRANKA_ROPE_ARENA_ENV_ID + _legacy_argparse_cfg_type = FrankaRopeArenaEnvironmentCfg + + def build(self, cfg: FrankaRopeArenaEnvironmentCfg) -> IsaacLabArenaEnvironment: + """Build an Arena environment description from ``cfg``.""" + + from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment + from isaaclab_arena.scene.scene import Scene + + table = self.asset_registry.get_asset_by_name("table")() + table.set_initial_pose( + Pose( + position_xyz=(0.5, 0.0, 0.0), + rotation_xyzw=(0.0, 0.0, 0.707, 0.707), + ), + ) + rope = FrankaRopeAsset() + ground_plane = self.asset_registry.get_asset_by_name("ground_plane")() + ground_plane.set_initial_pose(Pose(position_xyz=(0.0, 0.0, -1.05))) + light = self.asset_registry.get_asset_by_name("light")() + light.set_intensity(3000.0) + light.set_color((0.75, 0.75, 0.75)) + + embodiment = self.asset_registry.get_asset_by_name(cfg.embodiment)( + enable_cameras=cfg.enable_cameras, + ) + teleop_device = ( + self.device_registry.get_device_by_name(cfg.teleop_device)() if cfg.teleop_device is not None else None + ) + + scene = Scene( + assets=[ + table, + rope, + ground_plane, + light, + ], + ) + task = FrankaRopeTask() + + return IsaacLabArenaEnvironment( + name=self.name, + embodiment=embodiment, + scene=scene, + task=task, + teleop_device=teleop_device, + env_cfg_callback=_configure_franka_rope_env, + ) diff --git a/autodata_examples/envs/isaac_lab_arena/franka_rope/mdp/__init__.py b/autodata_examples/envs/isaac_lab_arena/franka_rope/mdp/__init__.py new file mode 100644 index 0000000..5185858 --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/mdp/__init__.py @@ -0,0 +1,19 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from .events import reset_rope_end_tracking, reset_rope_nodal_state +from .observations import ee_frame_pos, ee_frame_quat, gripper_pos, object_grasped, object_nodal_pos +from .terminations import rope_below_minimum, rope_ends_close_tracked + +__all__ = [ + "ee_frame_pos", + "ee_frame_quat", + "gripper_pos", + "object_grasped", + "object_nodal_pos", + "reset_rope_end_tracking", + "reset_rope_nodal_state", + "rope_below_minimum", + "rope_ends_close_tracked", +] diff --git a/autodata_examples/envs/isaac_lab_arena/franka_rope/mdp/events.py b/autodata_examples/envs/isaac_lab_arena/franka_rope/mdp/events.py new file mode 100644 index 0000000..2297c31 --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/mdp/events.py @@ -0,0 +1,106 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +import isaaclab.utils.math as math_utils +from isaaclab.assets import DeformableObject +from isaaclab.managers import SceneEntityCfg + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv, ManagerBasedRLEnv + + +# Mesh-resolution-specific nodes covering one half of the 549-node rope simulation mesh. +# fmt: off +_ROPE_PARTIAL_NODE_IDS = ( + 222, 223, 224, 225, 226, 227, 228, 229, 230, 231, 232, 233, + 270, 271, 272, 273, 274, 275, 276, 277, 278, 279, 280, 281, + 282, 283, 284, 285, 286, 287, 288, 289, 290, 291, 296, 297, + 298, 299, 300, 301, 302, 303, 304, 305, 306, 307, 308, 309, + 310, 311, 312, 313, 314, 315, 316, 317, 318, 319, 320, 321, + 322, 323, 324, 325, 326, 327, 328, 329, 331, 332, 333, 334, + 335, 336, 338, 339, 340, 341, 342, 343, 344, 345, 346, 347, + 348, 349, 350, 351, 352, 353, 354, 355, 356, 357, 358, 359, + 360, 361, 362, 363, 364, 365, 366, 367, 368, 369, 370, 371, + 372, 373, 374, 375, 376, 377, 378, 379, 380, 381, 382, 383, + 384, 385, 386, 387, 388, 389, 390, 391, 392, 393, 394, 395, + 396, 397, 398, 399, 400, 401, 402, 403, 404, 405, 406, 407, + 408, 409, 410, 411, 412, 413, 414, 415, 416, 417, 418, 419, + 420, 421, 422, 423, 424, 425, 426, 427, 428, 429, 430, 431, + 432, 433, 434, 435, 436, 437, 438, 439, 440, 441, 442, 443, + 444, 445, 446, 447, 448, 449, 450, 451, 452, 453, 454, 455, + 456, 457, 458, 459, 460, 461, 462, 463, 464, 465, 466, 467, + 468, 469, 470, 471, 472, 473, 474, 475, 476, 477, 478, 479, + 480, 481, 482, 483, 484, 485, 486, 487, 488, 489, 490, 491, + 492, 493, 494, 495, 496, 497, 498, 499, 500, 501, 502, 503, + 504, 505, 506, 507, 508, 509, 510, 511, 512, 513, 514, 515, + 516, 517, 518, 519, 520, 528, 529, 530, 531, 532, 533, 534, + 535, 536, 537, 538, 539, 540, 541, 542, 543, 544, 545, 546, + 547, 548, +) +# fmt: on + + +def _sample_pose_delta( + pose_range: dict[str, tuple[float, float]], + count: int, + device: str, +) -> tuple[torch.Tensor, torch.Tensor]: + """Sample translation and rotation deltas from named pose ranges.""" + + range_list = [pose_range.get(key, (0.0, 0.0)) for key in ("x", "y", "z", "roll", "pitch", "yaw")] + ranges = torch.tensor(range_list, device=device) + samples = math_utils.sample_uniform(ranges[:, 0], ranges[:, 1], (count, 6), device=device) + quaternion = math_utils.quat_from_euler_xyz(samples[:, 3], samples[:, 4], samples[:, 5]) + return samples[:, :3], quaternion + + +def reset_rope_end_tracking(env: ManagerBasedRLEnv, env_ids: torch.Tensor | None = None) -> None: + """Discard cached rope endpoint indices for reset environments.""" + + if not hasattr(env, "_rope_end_indices"): + return + if env_ids is None: + env._rope_end_indices.clear() + return + for env_id in env_ids.tolist(): + env._rope_end_indices.pop(int(env_id), None) + + +def reset_rope_nodal_state( + env: ManagerBasedEnv, + env_ids: torch.Tensor, + pose_range: dict[str, tuple[float, float]], + partial_pose_range: dict[str, tuple[float, float]], + asset_cfg: SceneEntityCfg = SceneEntityCfg("object"), +) -> None: + """Reset and randomize the rope, including a transform of one rope half. + + Args: + env: Environment containing the deformable rope. + env_ids: Environments to reset. + pose_range: Translation [m] and rotation [rad] ranges applied to all nodes. + partial_pose_range: Additional translation [m] and rotation [rad] ranges applied to one rope half. + asset_cfg: Deformable rope entity. + """ + + rope: DeformableObject = env.scene[asset_cfg.name] + nodal_state = rope.data.default_nodal_state_w.torch[env_ids].clone() + + position, quaternion = _sample_pose_delta(pose_range, len(env_ids), rope.device) + nodal_state[..., :3] = rope.transform_nodal_pos(nodal_state[..., :3], position, quaternion) + + assert max(_ROPE_PARTIAL_NODE_IDS) < nodal_state.shape[1], ( + f"Rope mesh has {nodal_state.shape[1]} simulation nodes, but the reset mapping expects at least " + f"{max(_ROPE_PARTIAL_NODE_IDS) + 1}." + ) + node_ids = torch.tensor(_ROPE_PARTIAL_NODE_IDS, device=rope.device) + position, quaternion = _sample_pose_delta(partial_pose_range, len(env_ids), rope.device) + nodal_state[..., node_ids, :3] = rope.transform_nodal_pos(nodal_state[..., node_ids, :3], position, quaternion) + + rope.write_nodal_state_to_sim_index(nodal_state, env_ids=env_ids) diff --git a/autodata_examples/envs/isaac_lab_arena/franka_rope/mdp/observations.py b/autodata_examples/envs/isaac_lab_arena/franka_rope/mdp/observations.py new file mode 100644 index 0000000..98900ca --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/mdp/observations.py @@ -0,0 +1,93 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +from isaaclab.assets import Articulation, DeformableObject +from isaaclab.managers import SceneEntityCfg +from isaaclab.sensors import FrameTransformer + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + + +def ee_frame_pos( + env: ManagerBasedRLEnv, + ee_frame_cfg: SceneEntityCfg = SceneEntityCfg("ee_frame"), +) -> torch.Tensor: + """Return the end-effector position in the environment frame [m].""" + + ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] + return ee_frame.data.target_pos_w.torch[:, 0, :] - env.scene.env_origins + + +def ee_frame_quat( + env: ManagerBasedRLEnv, + ee_frame_cfg: SceneEntityCfg = SceneEntityCfg("ee_frame"), +) -> torch.Tensor: + """Return the end-effector orientation quaternion in the world frame.""" + + ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] + return ee_frame.data.target_quat_w.torch[:, 0, :] + + +def gripper_pos( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntityCfg = SceneEntityCfg("robot", joint_names=["panda_finger_joint.*"]), +) -> torch.Tensor: + """Return the two Franka finger positions using the SoftMimicGen sign convention [m].""" + + robot: Articulation = env.scene[robot_cfg.name] + joint_pos = robot.data.joint_pos.torch[:, robot_cfg.joint_ids] + return torch.cat((joint_pos[:, 0:1], -joint_pos[:, 1:2]), dim=1) + + +def object_grasped( + env: ManagerBasedRLEnv, + robot_cfg: SceneEntityCfg, + ee_frame_cfg: SceneEntityCfg, + object_cfg: SceneEntityCfg, + distance_threshold: float = 0.015, + gripper_open_value: float = 0.04, + gripper_threshold: float = 0.005, +) -> torch.Tensor: + """Return whether the gripper is closed around a nearby rope node. + + Args: + env: Environment containing the robot and rope. + robot_cfg: Franka articulation and finger-joint selection. + ee_frame_cfg: End-effector frame sensor. + object_cfg: Deformable rope entity. + distance_threshold: Maximum end-effector-to-node distance [m]. + gripper_open_value: Finger position when fully open [m]. + gripper_threshold: Required deviation from the open position [m]. + + Returns: + Boolean tensor with one value per environment. + """ + + robot: Articulation = env.scene[robot_cfg.name] + ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name] + rope: DeformableObject = env.scene[object_cfg.name] + + nodal_pos_w = rope.data.nodal_pos_w.torch + end_effector_pos_w = ee_frame.data.target_pos_w.torch[:, 0, :] + minimum_distance = torch.linalg.vector_norm(nodal_pos_w - end_effector_pos_w.unsqueeze(1), dim=-1).amin(dim=1) + + finger_pos = robot.data.joint_pos.torch[:, robot_cfg.joint_ids] + gripper_closed = (torch.abs(finger_pos - gripper_open_value) > gripper_threshold).all(dim=1) + return (minimum_distance < distance_threshold) & gripper_closed + + +def object_nodal_pos( + env: ManagerBasedRLEnv, + object_cfg: SceneEntityCfg = SceneEntityCfg("object"), +) -> torch.Tensor: + """Return rope simulation-node positions in the environment frame [m].""" + + rope: DeformableObject = env.scene[object_cfg.name] + return rope.data.nodal_pos_w.torch - env.scene.env_origins.unsqueeze(1) diff --git a/autodata_examples/envs/isaac_lab_arena/franka_rope/mdp/terminations.py b/autodata_examples/envs/isaac_lab_arena/franka_rope/mdp/terminations.py new file mode 100644 index 0000000..ac42349 --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/mdp/terminations.py @@ -0,0 +1,78 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from __future__ import annotations + +import torch +from typing import TYPE_CHECKING + +from isaaclab.assets import DeformableObject +from isaaclab.managers import SceneEntityCfg + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedRLEnv + + +def rope_below_minimum( + env: ManagerBasedRLEnv, + minimum_height: float, + object_cfg: SceneEntityCfg = SceneEntityCfg("object"), +) -> torch.Tensor: + """Return whether the rope center is below a minimum world height [m].""" + + rope: DeformableObject = env.scene[object_cfg.name] + return rope.data.root_pos_w.torch[:, 2] < minimum_height + + +def rope_ends_close_tracked( + env: ManagerBasedRLEnv, + object_cfg: SceneEntityCfg = SceneEntityCfg("object"), + distance_threshold: float = 0.08, + distance_threshold_min: float = 0.02, + num_end_nodes: int = 3, +) -> torch.Tensor: + """Return whether the two tracked rope ends are close enough for success. + + Endpoint nodes are discovered from the farthest node pair at the start of each + episode and then tracked by simulation-node index. + + Args: + env: Environment containing the rope. + object_cfg: Deformable rope entity. + distance_threshold: Maximum distance between endpoint neighborhoods [m]. + distance_threshold_min: Minimum separation used to reject collapsed false positives [m]. + num_end_nodes: Number of nearby nodes used for each endpoint neighborhood. + + Returns: + Boolean tensor with one value per environment. + """ + + rope: DeformableObject = env.scene[object_cfg.name] + nodal_positions = rope.data.nodal_pos_w.torch + num_envs, num_nodes, _ = nodal_positions.shape + assert num_end_nodes <= num_nodes, f"Requested {num_end_nodes} endpoint nodes from a {num_nodes}-node rope." + + if not hasattr(env, "_rope_end_indices"): + env._rope_end_indices = {} + + result = torch.zeros(num_envs, dtype=torch.bool, device=env.device) + for env_id in range(num_envs): + env_nodal_pos = nodal_positions[env_id] + if env_id not in env._rope_end_indices: + farthest_flat_index = torch.cdist(env_nodal_pos, env_nodal_pos).argmax() + env._rope_end_indices[env_id] = ( + int((farthest_flat_index // num_nodes).item()), + int((farthest_flat_index % num_nodes).item()), + ) + + end_1_id, end_2_id = env._rope_end_indices[env_id] + end_1_distances = torch.linalg.vector_norm(env_nodal_pos - env_nodal_pos[end_1_id], dim=1) + end_2_distances = torch.linalg.vector_norm(env_nodal_pos - env_nodal_pos[end_2_id], dim=1) + end_1_node_ids = torch.topk(end_1_distances, num_end_nodes, largest=False).indices + end_2_node_ids = torch.topk(end_2_distances, num_end_nodes, largest=False).indices + + minimum_end_distance = torch.cdist(env_nodal_pos[end_1_node_ids], env_nodal_pos[end_2_node_ids]).amin() + result[env_id] = distance_threshold_min < minimum_end_distance < distance_threshold + + return result diff --git a/autodata_examples/envs/isaac_lab_arena/franka_rope/rope_asset.py b/autodata_examples/envs/isaac_lab_arena/franka_rope/rope_asset.py new file mode 100644 index 0000000..1eb21cd --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/rope_asset.py @@ -0,0 +1,39 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from __future__ import annotations + +from pathlib import Path + +from isaaclab.assets import DeformableObjectCfg +from isaaclab.sim.spawners.from_files.from_files_cfg import UsdFileCfg +from isaaclab_arena.assets.asset import Asset + +_ROPE_USD_PATH = Path(__file__).resolve().parent / "assets" / "Rope.usd" + + +class FrankaRopeAsset(Asset): + """Deformable rope asset that can participate in an Arena :class:`Scene`.""" + + def __init__(self, name: str = "object") -> None: + super().__init__(name=name, tags=["deformable"]) + self.object_cfg = DeformableObjectCfg( + prim_path="{ENV_REGEX_NS}/Object", + init_state=DeformableObjectCfg.InitialStateCfg( + pos=(0.5, 0.0, 0.02), + rot=(0.0, 0.0, 0.707, 0.707), + ), + spawn=UsdFileCfg(usd_path=str(_ROPE_USD_PATH)), + debug_vis=False, + ) + + def get_object_cfg(self) -> tuple[str, DeformableObjectCfg]: + """Return the scene key and Isaac Lab deformable-object configuration.""" + + return self.name, self.object_cfg + + def get_event_cfg(self) -> tuple[str, None]: + """Return no asset-level event; the Arena task owns nodal reset semantics.""" + + return self.name, None diff --git a/autodata_examples/envs/isaac_lab_arena/franka_rope/task.py b/autodata_examples/envs/isaac_lab_arena/franka_rope/task.py new file mode 100644 index 0000000..c7127f9 --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/task.py @@ -0,0 +1,152 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from __future__ import annotations + +import math +from typing import Never + +from isaaclab.envs.common import ViewerCfg +from isaaclab.managers import EventTermCfg as EventTerm +from isaaclab.managers import ObservationGroupCfg as ObsGroup +from isaaclab.managers import ObservationTermCfg as ObsTerm +from isaaclab.managers import SceneEntityCfg +from isaaclab.managers import TerminationTermCfg as DoneTerm +from isaaclab.utils.configclass import configclass +from isaaclab_arena.assets.register import register_task +from isaaclab_arena.embodiments.common.arm_mode import ArmMode +from isaaclab_arena.metrics.metric_base import MetricBase +from isaaclab_arena.metrics.success_rate import SuccessRateMetric +from isaaclab_arena.progress_tracking.completion_criteria import CompletionCriteria +from isaaclab_arena.tasks.task_base import TaskBase +from isaaclab_arena.tasks.task_termination_cfg import TaskTerminationCfg + +from . import mdp + + +@configclass +class FrankaRopeTaskObservationsCfg: + """Automatic SoftMimicGen subtask signals.""" + + @configclass + class SubtaskCfg(ObsGroup): + """Signals consumed by AutoData annotation.""" + + grasp = ObsTerm( + func=mdp.object_grasped, + params={ + "robot_cfg": SceneEntityCfg("robot", joint_names=["panda_finger_joint.*"]), + "ee_frame_cfg": SceneEntityCfg("ee_frame"), + "object_cfg": SceneEntityCfg("object"), + }, + ) + + def __post_init__(self) -> None: + self.enable_corruption = False + self.concatenate_terms = False + + subtask_terms: SubtaskCfg = SubtaskCfg() + + +@configclass +class FrankaRopeTaskEventsCfg: + """Rope randomization and endpoint tracking events.""" + + reset_rope_end_tracking = EventTerm( + func=mdp.reset_rope_end_tracking, + mode="reset", + ) + + reset_object_position = EventTerm( + func=mdp.reset_rope_nodal_state, + mode="reset", + params={ + "pose_range": { + "x": (0.0, 0.1), + "y": (0.0, 0.0), + "z": (0.0, 0.0), + "roll": (0.0, 0.0), + "pitch": (0.0, 0.0), + "yaw": (0.0, 0.0), + }, + "partial_pose_range": { + "x": (0.0, 0.1), + "y": (0.0, 0.1), + "z": (0.0, 0.0), + "roll": (0.0, 0.0), + "pitch": (0.0, 0.0), + "yaw": (-math.pi / 6.0, math.pi / 6.0), + }, + "asset_cfg": SceneEntityCfg("object"), + }, + ) + + +@register_task +class FrankaRopeTask(TaskBase): + """Bring the two ends of a deformable rope together with a Franka arm.""" + + def __init__(self, episode_length_s: float = 10.0) -> None: + super().__init__( + episode_length_s=episode_length_s, + task_description="Grasp and manipulate the deformable rope until its ends meet.", + ) + self.observation_cfg = FrankaRopeTaskObservationsCfg() + self.events_cfg = FrankaRopeTaskEventsCfg() + self.termination_cfg = TaskTerminationCfg( + timeout_s=self.episode_length_s, + success=[ + CompletionCriteria( + name="rope_ends_close", + predicate_sequence=[mdp.rope_ends_close_tracked], + ), + ], + failures={ + "object_dropping": DoneTerm( + func=mdp.rope_below_minimum, + params={ + "minimum_height": -0.05, + "object_cfg": SceneEntityCfg("object"), + }, + ), + }, + ) + + def get_scene_cfg(self) -> None: + """Return no task scene additions; the environment owns all assets.""" + + def get_observation_cfg(self) -> FrankaRopeTaskObservationsCfg: + """Return automatic subtask-signal observations.""" + + return self.observation_cfg + + def get_termination_cfg(self) -> TaskTerminationCfg: + """Return task termination terms.""" + + return self.termination_cfg + + def get_events_cfg(self) -> FrankaRopeTaskEventsCfg: + """Return deformable reset events.""" + + return self.events_cfg + + def get_mimic_env_cfg(self, arm_mode: ArmMode) -> Never: + """Reject Arena Mimic mode; SoftMimicGen is provided by AutoData.""" + + del arm_mode + raise NotImplementedError("FrankaRopeTask uses AutoData SoftMimicGen, not Arena Mimic mode.") + + def get_metrics(self) -> list[MetricBase]: + """Return success-rate evaluation.""" + + return [SuccessRateMetric()] + + def get_viewer_cfg(self) -> ViewerCfg: + """Return the rope workspace camera framing.""" + + return ViewerCfg( + eye=(1.0, 1.0, 0.5), + lookat=(0.0, 0.0, 0.0), + origin_type="env", + ) diff --git a/autodata_examples/envs/isaac_lab_arena/registration.py b/autodata_examples/envs/isaac_lab_arena/registration.py new file mode 100644 index 0000000..4c4a666 --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/registration.py @@ -0,0 +1,131 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from __future__ import annotations + +import argparse +import gymnasium as gym +from functools import partial +from typing import Any + +from isaaclab_arena.assets.registries import EnvironmentRegistry + +from .franka_rope import FRANKA_ROPE_ARENA_ENV_ID + + +def is_arena_environment(env_name: str) -> bool: + """Return whether ``env_name`` is owned by AutoData's Arena package.""" + + return env_name == FRANKA_ROPE_ARENA_ENV_ID + + +def _bind_gym_make_kwargs(env_name: str, env_kwargs: dict[str, Any]) -> None: + """Bind Arena-only constructor arguments for callers that use plain ``gym.make``. + + Isaac Lab's external-callback API can register an environment but cannot return + keyword arguments to the later ``gym.make`` call. Bind those arguments into the + registered entry point instead. + """ + + env_spec = gym.spec(env_name) + entry_point = env_spec.entry_point + assert entry_point is not None, f"Environment {env_name!r} has no Gym entry point." + if isinstance(entry_point, str): + from gymnasium.envs.registration import load_env_creator + + entry_point = load_env_creator(entry_point) + env_spec.entry_point = partial(entry_point, **env_kwargs) + + +def build_and_register_arena_environment( + *, + enable_cameras: bool, + num_envs: int, + device: str, + seed: int, + presets: str | None = None, +) -> dict[str, dict[str, Any]]: + """Build and register AutoData's typed Arena environment for the current run. + + Args: + enable_cameras: Whether to include environment camera sensors. + num_envs: Number of parallel environments in this run. + device: Simulation device for this run. + seed: Environment seed for this run. + presets: Optional physics backend preset; None preserves the environment default. + + Returns: + Gym constructor kwargs keyed by registered environment ID. + """ + + assert ( + FRANKA_ROPE_ARENA_ENV_ID not in gym.registry + ), f"Environment {FRANKA_ROPE_ARENA_ENV_ID!r} is already registered in this process." + + # ArenaEnvBuilder imports simulator modules and must be loaded only after AppLauncher starts Isaac Sim. + from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder + from isaaclab_arena.environments.arena_env_builder_cfg import ArenaEnvBuilderCfg + + environment_registry = EnvironmentRegistry() + environment_factory_type = environment_registry.get_component_by_name(FRANKA_ROPE_ARENA_ENV_ID) + environment_cfg_type = environment_registry.get_environment_cfg_type(environment_factory_type) + arena_environment = environment_factory_type().build(environment_cfg_type(enable_cameras=enable_cameras)) + + builder = ArenaEnvBuilder( + arena_environment, + ArenaEnvBuilderCfg( + num_envs=num_envs, + env_spacing=2.5, + seed=seed, + solve_relations=False, + device=device, + presets=presets, + ), + ) + env_name, _, env_kwargs = builder.build_registered() + return {env_name: env_kwargs} + + +def register_environment_from_cli() -> list[str]: + """Register the Arena rope environment through Isaac Lab's callback API. + + Isaac Lab invokes external registration callbacks without arguments. This adapter + reads the standard Lab and Arena options already present in ``sys.argv``, builds the + typed AutoData environment, and binds Arena's additional constructor arguments for + the script's later plain ``gym.make`` call. + + Returns: + Command-line arguments not consumed by the environment registration parser. + """ + + from isaaclab.app import AppLauncher + from isaaclab_arena.cli.isaaclab_arena_cli import ( + add_isaac_lab_cli_args, + add_isaaclab_arena_cli_args, + arena_env_builder_cfg_from_argparse, + ) + from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder + from isaaclab_arena_environments.cli import add_environment_cli_args, build_environment_from_cli + + parser = argparse.ArgumentParser(add_help=False) + parser.add_argument("--task", type=str, default=FRANKA_ROPE_ARENA_ENV_ID) + AppLauncher.add_app_launcher_args(parser) + add_isaac_lab_cli_args(parser) + add_isaaclab_arena_cli_args(parser) + + environment_factory_type = EnvironmentRegistry().get_component_by_name(FRANKA_ROPE_ARENA_ENV_ID) + add_environment_cli_args(parser, environment_factory_type) + parser.set_defaults(env_spacing=2.5, solve_relations=False) + args, remaining_args = parser.parse_known_args() + + requested_env_name = args.task.split(":")[-1] + assert ( + requested_env_name == FRANKA_ROPE_ARENA_ENV_ID + ), f"Arena callback expected {FRANKA_ROPE_ARENA_ENV_ID!r}, got {requested_env_name!r}." + + arena_environment = build_environment_from_cli(environment_factory_type, args) + builder = ArenaEnvBuilder(arena_environment, arena_env_builder_cfg_from_argparse(args)) + env_name, _, env_kwargs = builder.build_registered() + _bind_gym_make_kwargs(env_name, env_kwargs) + return remaining_args diff --git a/autodata_examples/tasks/franka_rope_softmimicgen.yaml b/autodata_examples/tasks/franka_rope_softmimicgen.yaml new file mode 100644 index 0000000..3794901 --- /dev/null +++ b/autodata_examples/tasks/franka_rope_softmimicgen.yaml @@ -0,0 +1,59 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +name: franka_rope +description: Grasp and manipulate a deformable rope until its ends meet. +algo: softmimicgen + +generation_policy: + name: franka_rope + seed: 1 + num_trials: 10 + guarantee_success: true + keep_failed: false + source_dataset_path: null + generation_path: null + task_name: null + use_skillgen: false + use_navigation_controller: false + reset_settling_steps: 10 + select_src_per_subtask: false + select_src_per_arm: false + transform_first_robot_pose: false + interpolate_from_last_target_pose: true + +subtasks: + franka: + - object_ref: object + description: Approach and grasp the deformable rope. + subtask_term_signal: grasp + selection_strategy: registration_cost + selection_strategy_kwargs: + nn_k: 3 + subtask_term_offset_range: [10, 15] + action_noise: 0.0 + num_interpolation_steps: 10 + num_fixed_steps: 0 + apply_noise_during_interpolation: false + algo_params: + object_soft: true + use_rotation_transform: true + bend_coef: 0.1 + rot_coef: 0.001 + + # Final subtask: no term signal (end-of-trajectory) and no term offset. + - object_ref: object + description: Manipulate the rope until its ends meet. + selection_strategy: registration_cost + selection_strategy_kwargs: + nn_k: 3 + subtask_term_offset_range: [0, 0] + action_noise: 0.0 + num_interpolation_steps: 5 + num_fixed_steps: 0 + apply_noise_during_interpolation: false + algo_params: + object_soft: true + use_rotation_transform: true + bend_coef: 0.1 + rot_coef: 0.001 diff --git a/autodata_interfaces/env/isaaclab_env_interface.py b/autodata_interfaces/env/isaaclab_env_interface.py index 435cc8d..73f8de2 100644 --- a/autodata_interfaces/env/isaaclab_env_interface.py +++ b/autodata_interfaces/env/isaaclab_env_interface.py @@ -23,7 +23,6 @@ import isaaclab.sim as sim_utils from isaaclab.assets import RigidObjectCfg from isaaclab.envs.mdp.actions.pink_task_space_actions import PinkInverseKinematicsAction -from isaaclab.envs.mdp.recorders.recorders_cfg import ActionStateRecorderManagerCfg from isaaclab.managers import DatasetExportMode, EventTermCfg, SceneEntityCfg from isaaclab.managers.recorder_manager import RecorderManagerBaseCfg from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR, ISAACLAB_NUCLEUS_DIR @@ -34,6 +33,7 @@ from autodata_interfaces.env.env_profile import EnvironmentProfile, convert_event_params from autodata_interfaces.env.pink_ik_action import PinkInverseKinematicsActionSharedUrdf +from autodata_interfaces.env.recorders import make_action_state_recorder_manager_cfg from autodata_interfaces.env.reset_request import EnvResetRequest from autodata_interfaces.env.scene_state import get_scene_state from autodata_interfaces.tasks.generation_policy_spec import GenerationPolicy @@ -230,7 +230,7 @@ def setup_env_config( # Setup recorders if recorder_cfg is None: - env_cfg.recorders = ActionStateRecorderManagerCfg() + env_cfg.recorders = make_action_state_recorder_manager_cfg(record_images=enable_cameras) else: env_cfg.recorders = recorder_cfg env_cfg.recorders.dataset_export_dir_path = output_dir diff --git a/autodata_interfaces/env/recorders.py b/autodata_interfaces/env/recorders.py new file mode 100644 index 0000000..a12bc5c --- /dev/null +++ b/autodata_interfaces/env/recorders.py @@ -0,0 +1,29 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +from __future__ import annotations + +from isaaclab.envs.mdp.recorders.recorders_cfg import ActionStateRecorderManagerCfg +from isaaclab.managers.recorder_manager import RecorderTerm, RecorderTermCfg + + +class PreStepCameraObservationsRecorder(RecorderTerm): + """Record image observations exposed in a separate camera group.""" + + def record_pre_step(self): + """Return camera observations alongside the recorded policy observations.""" + return "obs", self._env.obs_buf.get("camera_obs", {}) + + +def make_action_state_recorder_manager_cfg(record_images: bool = False) -> ActionStateRecorderManagerCfg: + """Return Isaac Lab's action-state recorder with optional camera observations. + + Args: + record_images: Include the separate camera observation group when present. + """ + + cfg = ActionStateRecorderManagerCfg() + if record_images: + cfg.record_pre_step_camera_observations = RecorderTermCfg(class_type=PreStepCameraObservationsRecorder) + return cfg diff --git a/autodata_tests/e2e/test_annotation.py b/autodata_tests/e2e/test_annotation.py index f5383d3..e6ff381 100644 --- a/autodata_tests/e2e/test_annotation.py +++ b/autodata_tests/e2e/test_annotation.py @@ -16,9 +16,33 @@ @pytest.mark.with_subprocess -def test_automatic_annotation_from_dataset(tmp_path): +@pytest.mark.parametrize( + ("input_dataset", "task_descriptor", "env_name", "signal_names", "device", "require_all_demos"), + [ + ( + "annotated_dataset_franka_stack_mimicgen.hdf5", + "franka_cube_stack.yaml", + "IsaacContrib-Stack-Cube-Franka-IK-Rel", + {"grasp_1", "stack_1", "grasp_2"}, + "cpu", + True, + ), + ( + "annotated_dataset_franka_rope_softmimicgen.hdf5", + "franka_rope_softmimicgen.yaml", + "Isaac-Rope-Franka-Arena-IK-Rel-v0", + {"grasp"}, + "cuda", + False, + ), + ], + ids=["franka_cube_stack", "franka_rope"], +) +def test_automatic_annotation_from_dataset( + tmp_path, input_dataset, task_descriptor, env_name, signal_names, device, require_all_demos +): """Use the stored task name and export usable termination signals during replay.""" - input_file = Path(TestPaths.test_data_dir) / "annotated_dataset_franka_stack_mimicgen.hdf5" + input_file = Path(TestPaths.test_data_dir) / input_dataset output_file = tmp_path / "annotated.hdf5" run_subprocess([ TestPaths.python_path, @@ -27,9 +51,9 @@ def test_automatic_annotation_from_dataset(tmp_path): "--viz", "none", "--device", - "cpu", + device, "--task_descriptor", - str(Path(TestPaths.tasks_dir) / "franka_cube_stack.yaml"), + str(Path(TestPaths.tasks_dir) / task_descriptor), "--embodiment", str(Path(TestPaths.embodiments_dir) / "franka_ik_rel.yaml"), "--input_file", @@ -39,14 +63,25 @@ def test_automatic_annotation_from_dataset(tmp_path): ]) assert_valid_dataset(str(output_file), min_num_demos=1) with h5py.File(input_file, "r") as source, h5py.File(output_file, "r") as annotated: - assert len(annotated["data"]) == len(source["data"]) - assert json.loads(annotated["data"].attrs["env_args"])["env_name"] == "IsaacContrib-Stack-Cube-Franka-IK-Rel" + if require_all_demos: + assert len(annotated["data"]) == len(source["data"]) + else: + # Deformable replay can diverge across simulator versions. Annotation exports only successful replays. + assert len(annotated["data"]) <= len(source["data"]) + assert json.loads(annotated["data"].attrs["env_args"])["env_name"] == env_name + source_info = next(iter(source["data"].values()))["obs/datagen_info"] for episode in annotated["data"].values(): signals = episode["obs/datagen_info/subtask_term_signals"] - assert set(signals) == {"grasp_1", "stack_1", "grasp_2"} + assert set(signals) == signal_names for signal in signals.values(): assert len(signal) == len(episode["actions"]) assert signal[:].any() + if "object_nodal_position" in source_info: + nodes = episode["obs/datagen_info/object_nodal_position/object"] + assert nodes.shape == ( + len(episode["actions"]), + *source_info["object_nodal_position/object"].shape[1:], + ) @pytest.mark.with_subprocess diff --git a/autodata_tests/e2e/test_recording.py b/autodata_tests/e2e/test_recording.py index 3de6757..b995fd5 100644 --- a/autodata_tests/e2e/test_recording.py +++ b/autodata_tests/e2e/test_recording.py @@ -14,7 +14,7 @@ from autodata_tests.utils.subprocess import run_subprocess -def configure_recording_smoke() -> None: +def configure_recording_smoke() -> list[str]: """Replace physical HID I/O and export three real simulation steps for the test.""" import torch @@ -22,9 +22,15 @@ def configure_recording_smoke() -> None: from isaaclab.devices import Se3SpaceMouse from isaaclab.envs import ManagerBasedRLEnv + from autodata_examples.envs import register_environments + # This assertion observes the actual script's parser after Python startup, # before its task lookup. No task or environment resolution is mocked. - assert recorder.args_cli.task == "IsaacContrib-Stack-Cube-Franka-IK-Rel" + assert recorder.args_cli.task in { + "IsaacContrib-Stack-Cube-Franka-IK-Rel", + "Isaac-Rope-Franka-Arena-IK-Rel-v0", + } + remaining_args = register_environments() Se3SpaceMouse._find_device = lambda self: None Se3SpaceMouse._run_device = lambda self: None original_step = ManagerBasedRLEnv.step @@ -35,6 +41,11 @@ def step_and_finish(env, actions): result = original_step(env, actions) steps += 1 if steps == 3: + if recorder.args_cli.task == "Isaac-Rope-Franka-Arena-IK-Rel-v0": + camera_obs = env.obs_buf["camera_obs"] + assert len(camera_obs) == 2 + for image in camera_obs.values(): + assert image.shape[:3] == (1, 128, 128) # Synthetic input does not solve the task. Explicitly export this # short test episode to verify recording and metadata, then stop. env.recorder_manager.record_pre_reset([0], force_export_or_skip=False) @@ -46,10 +57,19 @@ def step_and_finish(env, actions): return result ManagerBasedRLEnv.step = step_and_finish + return remaining_args @pytest.mark.with_subprocess -def test_lab_record_demos_task_startup(tmp_path): +@pytest.mark.parametrize( + ("task_name", "camera_args"), + [ + ("IsaacContrib-Stack-Cube-Franka-IK-Rel", []), + ("Isaac-Rope-Franka-Arena-IK-Rel-v0", ["--enable_cameras"]), + ], + ids=["franka_cube_stack", "franka_rope"], +) +def test_lab_record_demos_task_startup(tmp_path, task_name, camera_args): """The direct recorder uses the registered task and records real actions and states.""" output_file = tmp_path / "recording.hdf5" script = Path(TestPaths.repo_root) / "submodules/IsaacLab-Arena/submodules/IsaacLab/scripts/tools/record_demos.py" @@ -60,7 +80,7 @@ def test_lab_record_demos_task_startup(tmp_path): "--viz", "kit", "--task", - "IsaacContrib-Stack-Cube-Franka-IK-Rel", + task_name, "--teleop_device", "spacemouse", "--dataset_file", @@ -69,12 +89,15 @@ def test_lab_record_demos_task_startup(tmp_path): "1", "--external_callback", "autodata_tests.e2e.test_recording.configure_recording_smoke", + *camera_args, ], env={**os.environ, "HEADLESS": "1"}, ) with h5py.File(output_file, "r") as dataset: - assert json.loads(dataset["data"].attrs["env_args"])["env_name"] == "IsaacContrib-Stack-Cube-Franka-IK-Rel" + assert json.loads(dataset["data"].attrs["env_args"])["env_name"] == task_name assert len(dataset["data"]) == 1 episode = next(iter(dataset["data"].values())) assert episode["actions"].shape == (3, 7) assert "states" in episode + if task_name == "Isaac-Rope-Franka-Arena-IK-Rel-v0": + assert episode["states/deformable_object/object/nodal_position"].shape == (3, 549, 3) diff --git a/autodata_tests/e2e/test_rope_environment.py b/autodata_tests/e2e/test_rope_environment.py new file mode 100644 index 0000000..63abc7b --- /dev/null +++ b/autodata_tests/e2e/test_rope_environment.py @@ -0,0 +1,181 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Simulate the rope-folding task with the PhysX physics backend.""" + +import os + +import pytest +import warp as wp + +from autodata_tests.utils.constants import TestPaths +from autodata_tests.utils.subprocess import run_subprocess + + +@pytest.mark.with_subprocess +@pytest.mark.parametrize("presets", ["physx"]) +@pytest.mark.parametrize("num_envs", [1, 3]) +def test_rope_environment(presets, num_envs, tmp_path): + result_file = tmp_path / "passed" + run_subprocess( + [ + TestPaths.python_path, + "-m", + __name__, + "--presets", + presets, + "--num_envs", + str(num_envs), + "--viz", + "kit", + "--result_file", + str(result_file), + "--enable_cameras", + ], + env={**os.environ, "HEADLESS": "1"}, + ) + assert result_file.exists(), "Simulation exited before completing the checks" + + +def main() -> None: + """Check spawning, stepping, contacts, observations and indexed resets.""" + import argparse + import tempfile + + wp.config.enable_backward = False + from isaaclab.app import AppLauncher + + parser = argparse.ArgumentParser() + parser.add_argument("--presets", choices=("physx",), required=True) + parser.add_argument("--num_envs", type=int, default=1) + parser.add_argument("--snapshot", type=str, default=None) + parser.add_argument("--result_file", type=str, default=None) + AppLauncher.add_app_launcher_args(parser) + parser.add_argument("--enable_cameras", action="store_true") + args = parser.parse_args() + enable_cameras = args.enable_cameras + app = AppLauncher(args).app + + import gymnasium as gym + import torch + + from autodata_examples.envs.isaac_lab_arena.franka_rope.environment import FRANKA_ROPE_ARENA_ENV_ID + from autodata_examples.envs.isaac_lab_arena.registration import register_environment_from_cli + + # Use the same backend-selection callback as the public Lab teleop/replay scripts. + register_environment_from_cli() + cfg = gym.spec(FRANKA_ROPE_ARENA_ENV_ID).kwargs["env_cfg_entry_point"] + cfg.seed = 42 + cfg.recorders.dataset_export_dir_path = tempfile.mkdtemp(prefix="rope_env_check_") + env = gym.make(FRANKA_ROPE_ARENA_ENV_ID, cfg=cfg).unwrapped + exit_code = 0 + try: + + def camera_observations(): + # Headless Kit renders on demand in the pinned Lab version. Request + # an actual frame before reading the RTX camera annotators. + if args.presets == "physx": + for visualizer in env.sim.visualizers: + if visualizer.pumps_app_update(): + visualizer.render_rgb_array() + for sensor in env.scene.sensors.values(): + if hasattr(sensor.cfg, "renderer_cfg"): + sensor.update(0.0, force_recompute=True) + return env.observation_manager.compute()["camera_obs"] + + env.reset() + rope = env.scene["object"] + robot = env.scene["robot"] + assert rope.data.nodal_pos_w.torch.shape == (args.num_envs, 549, 3) + expected_module = "isaaclab_physx" + assert type(rope).__module__.startswith(expected_module) + print(f"ROPE backend={args.presets} implementation={type(rope)}", flush=True) + initial_nodes = rope.data.nodal_pos_w.torch.clone() + actions = torch.zeros((args.num_envs, env.action_manager.total_action_dim), device=env.device) + actions[:, -1] = 1.0 + with torch.inference_mode(): + for step in range(100): + obs, _, terminated, _, _ = env.step(actions) + nodes = rope.data.nodal_pos_w.torch + assert torch.isfinite(nodes).all(), f"Non-finite rope nodes at step {step}" + assert not terminated.any(), ( + f"Unexpected task termination at step {step}: " + f"{env.termination_manager.get_active_iterable_terms(int(terminated.nonzero()[0]))}" + ) + assert torch.isfinite(obs["policy"]["object_nodal_pos"]).all() + if step % 20 == 0: + local = nodes - env.scene.env_origins[:, None, :] + print(f"ROPE step={step} min={local.amin((0, 1))} max={local.amax((0, 1))}", flush=True) + assert (rope.data.nodal_pos_w.torch - initial_nodes).abs().max() > 1.0e-4 + + # Gravity must settle the rope on the tabletop, without penetrating or exploding. + local = rope.data.nodal_pos_w.torch - env.scene.env_origins[:, None, :] + assert local[..., 2].min() > -0.01 + assert local[..., 2].max() < 0.15 + assert (local.amax(1) - local.amin(1)).amax() < 1.5 + + if enable_cameras: + images = camera_observations() + assert len(images) == 2 + for image in images.values(): + assert image.shape[:3] == (args.num_envs, 128, 128) + assert image.float().std() > 1.0 + if args.snapshot: + from PIL import Image + + snapshot = next(value for name, value in images.items() if "agentview" in name) + Image.fromarray(snapshot[0, ..., :3].cpu().numpy()).save(args.snapshot) + + # Relative IK and gripper actions must drive the PhysX articulation. + before_joints = robot.data.joint_pos.torch.clone() + before_images = {name: value.clone() for name, value in images.items()} if enable_cameras else {} + actions[:, 2] = 0.005 + actions[:, -1] = -1.0 + for _ in range(12): + obs, *_ = env.step(actions) + after_joints = robot.data.joint_pos.torch + assert torch.isfinite(after_joints).all() + assert (after_joints[:, :7] - before_joints[:, :7]).abs().max() > 0.005 + assert after_joints[:, -2:].mean() < before_joints[:, -2:].mean() - 0.005 + if enable_cameras: + images = camera_observations() + assert any( + (images[name].float() - before.float()).abs().mean() > 0.01 for name, before in before_images.items() + ), "Camera images did not change when the robot moved" + + # Reset one environment without changing the other environments' nodes. + before_reset = rope.data.nodal_pos_w.torch.clone() + env.reset(env_ids=torch.tensor([0], device=env.device)) + after_reset = rope.data.nodal_pos_w.torch + assert (after_reset[0] - before_reset[0]).abs().max() > 0.001 + torch.testing.assert_close(after_reset[1:], before_reset[1:], atol=1.0e-5, rtol=0) + + # Restoring a saved state must preserve nodal coordinates, including env offsets. + state = env.scene.get_state(is_relative=True) + saved_nodes = rope.data.nodal_pos_w.torch.clone() + env.reset() + env.reset_to(state, env_ids=None, is_relative=True) + torch.testing.assert_close(rope.data.nodal_pos_w.torch, saved_nodes, atol=1.0e-5, rtol=0) + actions.zero_() + actions[:, -1] = 1.0 + for _ in range(5): + env.step(actions) + assert torch.isfinite(rope.data.nodal_pos_w.torch).all() + print("ROPE smoke passed", flush=True) + if args.result_file: + from pathlib import Path + + Path(args.result_file).touch() + except BaseException: + import traceback + + traceback.print_exc() + exit_code = 1 + finally: + env.close() + # Kit terminates the process directly; preserve failures across shutdown. + app.close(exit_code=exit_code) + + +if __name__ == "__main__": + main() diff --git a/autodata_tests/e2e/test_softmimicgen_data_generation.py b/autodata_tests/e2e/test_softmimicgen_data_generation.py new file mode 100644 index 0000000..f3b06f8 --- /dev/null +++ b/autodata_tests/e2e/test_softmimicgen_data_generation.py @@ -0,0 +1,85 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import h5py +import numpy as np +import os +import tempfile + +import pytest + +from autodata_tests.utils.constants import TestPaths +from autodata_tests.utils.subprocess import run_subprocess +from autodata_tests.utils.utils import assert_valid_dataset + +HEADLESS = True +GENERATION_NUM_TRIALS = 1 + + +def _run_franka_rope_softmimicgen(num_envs: int, device: str, presets: str = "physx") -> None: + """Run SoftMimicGen data generation for the Franka rope task.""" + + with tempfile.TemporaryDirectory() as temp_dir: + output_file = os.path.join(temp_dir, "generated.hdf5") + + args = [ + TestPaths.python_path, + TestPaths.generate_dataset_script, + "--env_name", + "Isaac-Rope-Franka-Arena-IK-Rel-v0", + "--presets", + presets, + "--alg", + "softmimicgen", + "--task_descriptor", + os.path.join(TestPaths.tasks_dir, "franka_rope_softmimicgen.yaml"), + "--embodiment", + os.path.join(TestPaths.embodiments_dir, "franka_ik_rel.yaml"), + "--input_file", + os.path.join(TestPaths.test_data_dir, "annotated_dataset_franka_rope_softmimicgen.hdf5"), + "--output_file", + output_file, + "--generation_num_trials", + str(GENERATION_NUM_TRIALS), + "--num_envs", + str(num_envs), + "--device", + device, + "--viz", + "none" if HEADLESS else "kit", + ] + run_subprocess(args) + + assert_valid_dataset(output_file, min_num_demos=GENERATION_NUM_TRIALS) + with h5py.File(output_file, "r") as dataset: + for episode in dataset["data"].values(): + assert episode.attrs["success"] + initial_nodes = episode["initial_state/deformable_object/object/nodal_position"][:] + assert initial_nodes.shape == (1, 549, 3) + assert np.isfinite(initial_nodes).all() + assert np.abs(initial_nodes).max() < 2.0 + nodes = episode["states/deformable_object/object/nodal_position"][:] + assert nodes.shape == (len(episode["actions"]), 549, 3) + assert np.isfinite(nodes).all() + # Parallel environment origins must not leak into recorded positions. + assert np.abs(nodes).max() < 2.0 + + +@pytest.mark.with_subprocess +@pytest.mark.parametrize("presets", ["physx"]) +def test_franka_rope_softmimicgen_data_generation_single_env_cuda(presets): + """SoftMimicGen generation for the Franka rope task on a single env on GPU.""" + _run_franka_rope_softmimicgen(num_envs=1, device="cuda", presets=presets) + + +@pytest.mark.with_subprocess +@pytest.mark.parametrize("presets", ["physx"]) +def test_franka_rope_softmimicgen_data_generation_multi_env_cuda(presets): + """SoftMimicGen generation for the Franka rope task on multiple parallel envs on GPU.""" + _run_franka_rope_softmimicgen(num_envs=10, device="cuda", presets=presets) + + +if __name__ == "__main__": + for presets in ("physx",): + test_franka_rope_softmimicgen_data_generation_single_env_cuda(presets) + test_franka_rope_softmimicgen_data_generation_multi_env_cuda(presets) diff --git a/autodata_tests/interfaces/env/test_franka_rope_arena.py b/autodata_tests/interfaces/env/test_franka_rope_arena.py new file mode 100644 index 0000000..86696e9 --- /dev/null +++ b/autodata_tests/interfaces/env/test_franka_rope_arena.py @@ -0,0 +1,213 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Configuration tests for the AutoData-owned Arena Franka rope environment.""" + +import gymnasium as gym +import sys +from functools import partial +from types import SimpleNamespace + +import pytest +from isaaclab.sim import SimulationCfg +from isaaclab_arena.assets.registries import EnvironmentRegistry +from isaaclab_arena.tasks.task_termination_cfg import TaskTerminationCfg +from isaaclab_arena.utils.physics_backend import PhysicsBackend + +import autodata_examples.envs as env_registration +import autodata_examples.envs.isaac_lab_arena as arena_registration +from autodata_examples.envs.isaac_lab_arena import registration +from autodata_examples.envs.isaac_lab_arena.franka_rope import mdp +from autodata_examples.envs.isaac_lab_arena.franka_rope.embodiment import FrankaRopeIKEmbodiment +from autodata_examples.envs.isaac_lab_arena.franka_rope.environment import ( + FRANKA_ROPE_ARENA_ENV_ID, + FrankaRopeArenaEnvironment, + FrankaRopeArenaEnvironmentCfg, + _configure_franka_rope_env, +) +from autodata_examples.envs.isaac_lab_arena.franka_rope.rope_asset import FrankaRopeAsset +from autodata_examples.envs.isaac_lab_arena.franka_rope.task import FrankaRopeTask + + +def test_arena_environment_factory_has_typed_config(): + assert FrankaRopeArenaEnvironment.name == FRANKA_ROPE_ARENA_ENV_ID + assert FrankaRopeArenaEnvironment._legacy_argparse_cfg_type is FrankaRopeArenaEnvironmentCfg + + +def test_arena_environment_factory_is_registered(): + registry = EnvironmentRegistry() + assert registry.get_component_by_name(FRANKA_ROPE_ARENA_ENV_ID) is FrankaRopeArenaEnvironment + assert registry.get_environment_cfg_type(FrankaRopeArenaEnvironment) is FrankaRopeArenaEnvironmentCfg + + +@pytest.mark.parametrize("presets", [None, "physx"]) +def test_arena_registration_uses_run_configuration_and_preserves_make_kwargs(monkeypatch, presets): + captured = {} + variation_recorder = object() + + class FakeArenaEnvBuilder: + def __init__(self, arena_environment, cfg): + captured["arena_environment"] = arena_environment + captured["cfg"] = cfg + + def build_registered(self): + return FRANKA_ROPE_ARENA_ENV_ID, object(), {"variation_recorder": variation_recorder} + + monkeypatch.delitem(gym.registry, FRANKA_ROPE_ARENA_ENV_ID, raising=False) + monkeypatch.setitem( + sys.modules, + "isaaclab_arena.environments.arena_env_builder", + SimpleNamespace(ArenaEnvBuilder=FakeArenaEnvBuilder), + ) + monkeypatch.setattr( + FrankaRopeArenaEnvironment, + "build", + lambda _self, cfg: SimpleNamespace( + name=FRANKA_ROPE_ARENA_ENV_ID, + embodiment=SimpleNamespace(enable_cameras=cfg.enable_cameras), + ), + ) + + make_kwargs = env_registration.register_environment_for_run( + env_name=FRANKA_ROPE_ARENA_ENV_ID, + enable_cameras=True, + num_envs=8, + device="cuda:1", + seed=23, + presets=presets, + ) + + builder_cfg = captured["cfg"] + assert captured["arena_environment"].name == FRANKA_ROPE_ARENA_ENV_ID + assert captured["arena_environment"].embodiment.enable_cameras is True + assert builder_cfg.num_envs == 8 + assert builder_cfg.env_spacing == 2.5 + assert builder_cfg.seed == 23 + assert builder_cfg.solve_relations is False + assert builder_cfg.device == "cuda:1" + assert builder_cfg.presets == (None if presets is None else PhysicsBackend(presets)) + assert make_kwargs == {FRANKA_ROPE_ARENA_ENV_ID: {"variation_recorder": variation_recorder}} + + +def test_regular_isaac_lab_run_does_not_build_arena(monkeypatch): + arena_registration_called = False + + def register_arena_environment(**_kwargs): + nonlocal arena_registration_called + arena_registration_called = True + return {} + + monkeypatch.setattr(arena_registration, "build_and_register_arena_environment", register_arena_environment) + + make_kwargs = env_registration.register_environment_for_run( + env_name="Isaac-Regular-Lab-Env-v0", + enable_cameras=False, + num_envs=4, + device="cuda:0", + seed=17, + ) + + assert make_kwargs == {} + assert arena_registration_called is False + + +def test_external_callback_routes_arena_task(monkeypatch): + callback_called = False + + def register_arena_environment_from_cli(): + nonlocal callback_called + callback_called = True + return ["hydra.option=value"] + + monkeypatch.setattr(arena_registration, "register_environment_from_cli", register_arena_environment_from_cli) + monkeypatch.setattr(sys, "argv", ["replay_demos.py", "--task", FRANKA_ROPE_ARENA_ENV_ID]) + + remaining_args = env_registration.register_environments() + + assert callback_called is True + assert remaining_args == ["hydra.option=value"] + + +def test_external_callback_skips_arena_for_regular_lab_task(monkeypatch): + callback_called = False + + def register_arena_environment_from_cli(): + nonlocal callback_called + callback_called = True + return [] + + monkeypatch.setattr(arena_registration, "register_environment_from_cli", register_arena_environment_from_cli) + monkeypatch.setattr(sys, "argv", ["replay_demos.py", "--task", "Isaac-Regular-Lab-Env-v0"]) + + assert env_registration.register_environments() == [] + assert callback_called is False + + +def test_external_callback_binds_make_kwargs_without_gym_deepcopy(): + test_env_name = "AutoData-Test-Arena-Callback-v0" + variation_recorder = object() + gym.register(id=test_env_name, entry_point=lambda **_kwargs: None) + try: + registration._bind_gym_make_kwargs( # noqa: SLF001 + test_env_name, + {"variation_recorder": variation_recorder}, + ) + + env_spec = gym.spec(test_env_name) + assert isinstance(env_spec.entry_point, partial) + assert env_spec.entry_point.keywords["variation_recorder"] is variation_recorder + assert "variation_recorder" not in env_spec.kwargs + finally: + del gym.registry[test_env_name] + + +def test_rope_asset_uses_deformable_scene_key(): + rope = FrankaRopeAsset() + name, object_cfg = rope.get_object_cfg() + assert name == "object" + assert object_cfg.prim_path == "{ENV_REGEX_NS}/Object" + assert object_cfg.spawn.usd_path.endswith("/isaac_lab_arena/franka_rope/assets/Rope.usd") + + +def test_rope_mdp_terms_are_owned_by_arena_package(): + arena_module_prefix = "autodata_examples.envs.isaac_lab_arena.franka_rope.mdp" + assert mdp.object_nodal_pos.__module__.startswith(arena_module_prefix) + assert mdp.reset_rope_nodal_state.__module__.startswith(arena_module_prefix) + assert mdp.rope_ends_close_tracked.__module__.startswith(arena_module_prefix) + + +def test_rope_task_exposes_annotation_and_success_terms(): + task = FrankaRopeTask(episode_length_s=15.0) + assert task.get_observation_cfg().subtask_terms.grasp is not None + terminations = task.get_termination_cfg() + assert isinstance(terminations, TaskTerminationCfg) + assert terminations.timeout_s == 15.0 + assert len(terminations.success) == 1 + assert terminations.success[0].predicate_sequence == [mdp.rope_ends_close_tracked] + assert terminations.failures["object_dropping"].func is mdp.rope_below_minimum + assert task.get_events_cfg().reset_object_position is not None + + +def test_rope_environment_configures_lab_timing(): + env_cfg = SimpleNamespace(sim=SimulationCfg(), scene=SimpleNamespace(replicate_physics=True)) + + assert _configure_franka_rope_env(env_cfg) is env_cfg + assert env_cfg.decimation * env_cfg.sim.dt == 0.05 + assert env_cfg.sim.render_interval == 2 + assert env_cfg.num_rerenders_on_reset == 1 + assert env_cfg.scene.replicate_physics is False + assert set(env_cfg.teleop_devices.devices) == {"keyboard", "spacemouse"} + + +def test_rope_embodiment_preserves_arena_observation_names(): + embodiment = FrankaRopeIKEmbodiment(enable_cameras=False) + assert embodiment.name == "franka_rope_ik" + assert embodiment.observation_config.policy.eef_pos is not None + assert embodiment.observation_config.policy.eef_quat is not None + assert embodiment.observation_config.policy.object_nodal_pos is not None + assert embodiment.event_config.randomize_franka_joint_state.params["position_range"] == (-0.02, 0.02) + assert embodiment.reward_config is None + for camera_name in embodiment.camera_config.camera_names(): + camera = getattr(embodiment.camera_config, camera_name) + assert camera.renderer_cfg.renderer_type == "isaac_rtx" + assert camera.renderer_cfg.global_settings.antialiasing_mode == "DLSS" diff --git a/autodata_tests/test_data/annotated_dataset_franka_rope_softmimicgen.hdf5 b/autodata_tests/test_data/annotated_dataset_franka_rope_softmimicgen.hdf5 new file mode 100644 index 0000000..2c25c2e --- /dev/null +++ b/autodata_tests/test_data/annotated_dataset_franka_rope_softmimicgen.hdf5 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6fa2b5db4752242eb91cec522053ad1a040741764dcb71ebab6f3abe80e4992c +size 30514205 diff --git a/autodata_tests/utils/test_ci_cache.py b/autodata_tests/utils/test_ci_cache.py new file mode 100644 index 0000000..fe3a21f --- /dev/null +++ b/autodata_tests/utils/test_ci_cache.py @@ -0,0 +1,97 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CI must mount a writable home cache for Arena's composed robot assets.""" + +import json +import os +import pwd +import shutil +import subprocess +import sys +from pathlib import Path + +import pytest + +from autodata_tests.utils.constants import TestPaths + + +@pytest.fixture +def ci_launcher(tmp_path): + """Exercise the real shell entry points while replacing only Docker and X11.""" + root = tmp_path / "checkout with spaces" + for relative in ("scripts/ci/run_tests.sh", "docker/run_docker.sh"): + target = root / relative + target.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(Path(TestPaths.repo_root) / relative, target) + target.chmod(0o755) + + commands = tmp_path / "bin" + commands.mkdir() + docker = commands / "docker" + docker.write_text( + f"#!{sys.executable}\n" + "import json, os, sys\n" + "args = sys.argv[1:]\n" + "if args[:2] == ['images', '-q']:\n" + " print('test-image')\n" + "elif args[:2] == ['container', 'inspect']:\n" + " print('false')\n" + "elif args[:1] == ['run']:\n" + " with open(os.environ['TEST_DOCKER_LOG'], 'w') as output:\n" + " json.dump(args, output)\n" + "elif args[:2] != ['ps', '-a']:\n" + " raise SystemExit(f'Unexpected Docker command: {args}')\n" + ) + docker.chmod(0o755) + xhost = commands / "xhost" + xhost.write_text("#!/bin/sh\nexit 0\n") + xhost.chmod(0o755) + runner_temp = tmp_path / "runner temp" + runner_temp.mkdir() + log = tmp_path / "docker.json" + environment = os.environ.copy() + environment.pop("AUTODATA_DOCKER_CACHE_DIR", None) + environment.update( + PATH=f"{commands}:{environment['PATH']}", + RUNNER_TEMP=str(runner_temp), + TEST_DOCKER_LOG=str(log), + FORCE_REBUILD="false", + AUTODATA_RESULTS_DIR="", + ) + + def run(**variables): + result = subprocess.run( + ["bash", str(root / "scripts/ci/run_tests.sh")], + env=environment | variables, + capture_output=True, + text=True, + timeout=30, + ) + assert result.returncode == 0, result.stdout + result.stderr + args = json.loads(log.read_text()) + target = f"/home/{pwd.getpwuid(os.getuid()).pw_name}/.cache" + sources = [args[i + 1].split(":")[0] for i, arg in enumerate(args[:-1]) if arg == "-v"] + mounts = [source for source in sources if f"{source}:{target}" in args] + assert len(mounts) == 1 + source = Path(mounts[0]) + assert source.is_dir() + assert source.stat().st_uid == os.getuid() + (source / "arena-cache-write-check").write_text("writable") + assert "/isaac-sim/python.sh" in args and "pytest" in args + return source + + return run, runner_temp + + +def test_ci_cache_is_fresh_and_owned_by_runner(ci_launcher): + run, runner_temp = ci_launcher + first, second = run(), run() + assert first.parent == second.parent == runner_temp + assert first != second + + +def test_ci_cache_respects_explicit_mount(ci_launcher, tmp_path): + run, _ = ci_launcher + selected = tmp_path / "selected cache" + assert run(AUTODATA_DOCKER_CACHE_DIR=str(selected)) == selected diff --git a/docker/run_docker.sh b/docker/run_docker.sh index f623012..3946bf4 100755 --- a/docker/run_docker.sh +++ b/docker/run_docker.sh @@ -43,6 +43,8 @@ ISAACLAB_PATH="${WORKDIR}/submodules/IsaacLab-Arena/submodules/IsaacLab" # Optional host datasets directory mounted at /datasets (in addition to the repo's own datasets/). DATASETS_HOST_MOUNT_DIRECTORY="$HOME/datasets" +# CI can provide a writable cache without changing the shared host cache's ownership. +CACHE_HOST_MOUNT_DIRECTORY="${AUTODATA_DOCKER_CACHE_DIR:-$HOME/.cache}" FORCE_REBUILD=false NO_CACHE="" @@ -133,6 +135,9 @@ add_volume_if_it_exists() { [ -d "$1" ] && echo "-v $1:$2" } +# Create it as the host user before Docker can create a root-owned bind source. +mkdir -p "${CACHE_HOST_MOUNT_DIRECTORY}" + # Forward the host SSH agent if available (so git over SSH works in-container). SSH_DOCKER_ARGS=() if [ -n "${SSH_AUTH_SOCK:-}" ] && [ -S "$SSH_AUTH_SOCK" ]; then @@ -151,8 +156,8 @@ DOCKER_RUN_ARGS=( # Live-mount the repo: host edits are reflected in the container's editable installs. "-v" "${REPO_ROOT}:${WORKDIR}" $(add_volume_if_it_exists "$DATASETS_HOST_MOUNT_DIRECTORY" /datasets) - # Share the host Kit/pip cache to speed up shader warmup and reinstalls across runs. - "-v" "$HOME/.cache:/home/$(id -un)/.cache" + # Share the selected cache, including Arena's hard-coded ~/.cache location. + "-v" "${CACHE_HOST_MOUNT_DIRECTORY}:/home/$(id -un)/.cache" # X11 passthrough so "--viz kit" can open a window. "-v" "/tmp/.X11-unix:/tmp/.X11-unix:rw" "${SSH_DOCKER_ARGS[@]}" diff --git a/scripts/annotate_demos.py b/scripts/annotate_demos.py index ea522d6..83930ec 100644 --- a/scripts/annotate_demos.py +++ b/scripts/annotate_demos.py @@ -45,6 +45,9 @@ default=None, help="Environment name. Overrides the env name recorded in the source dataset.", ) +parser.add_argument( + "--presets", choices=("physx",), default=None, help="Physics backend preset for Arena environments." +) parser.add_argument( "--task_descriptor", type=str, @@ -102,7 +105,9 @@ import contextlib # noqa: E402 import gymnasium as gym # noqa: E402 import math # noqa: E402 +import numpy as np # noqa: E402 import os # noqa: E402 +import random # noqa: E402 import sys # noqa: E402 import torch # noqa: E402 import traceback # noqa: E402 @@ -116,9 +121,11 @@ from isaaclab.utils.datasets import EpisodeData, HDF5DatasetFileHandler # noqa: E402 from autodata_core.pool import DataGenInfoPool # noqa: E402 +from autodata_examples.envs import register_environment_for_run # noqa: E402 from autodata_interfaces.datastream import Datastream # noqa: E402 from autodata_interfaces.embodiments import embodiment_adapter_from_yaml # noqa: E402 from autodata_interfaces.env import get_env_name_from_dataset, setup_env_config, setup_output_paths # noqa: E402 +from autodata_interfaces.env.recorders import PreStepCameraObservationsRecorder # noqa: E402 from autodata_interfaces.env.registration import register_external_environment # noqa: E402 from autodata_interfaces.tasks.task_descriptor import TaskDescriptor # noqa: E402 @@ -236,6 +243,19 @@ def main() -> int: task_descriptor = TaskDescriptor.from_yaml(args_cli.task_descriptor) generation_policy = task_descriptor.get_generation_policy() + random.seed(generation_policy.seed) + np.random.seed(generation_policy.seed) + torch.manual_seed(generation_policy.seed) + + env_make_kwargs = register_environment_for_run( + env_name=env_name, + enable_cameras=args_cli.record_images, + num_envs=1, + device=args_cli.device, + seed=generation_policy.seed, + presets=args_cli.presets, + ) + # Start signals are required only by SkillGen. annotate_start_signals = generation_policy.use_skillgen @@ -268,10 +288,14 @@ def main() -> int: recorder_cfg=AnnotationRecorderManagerCfg(), enable_cameras=args_cli.record_images, ) + if args_cli.record_images: + env_cfg.recorders.record_pre_step_camera_observations = RecorderTermCfg( + class_type=PreStepCameraObservationsRecorder + ) # Only export episodes we explicitly mark successful (i.e. fully annotated). env_cfg.recorders.dataset_export_mode = DatasetExportMode.EXPORT_SUCCEEDED_ONLY - env = gym.make(env_name, cfg=env_cfg).unwrapped + env = gym.make(env_name, cfg=env_cfg, **env_make_kwargs.get(env_name, {})).unwrapped try: # Create the Datastream embodiment_adapter = embodiment_adapter_from_yaml(args_cli.embodiment) diff --git a/scripts/ci/run_tests.sh b/scripts/ci/run_tests.sh index f9701ae..2b91fdc 100755 --- a/scripts/ci/run_tests.sh +++ b/scripts/ci/run_tests.sh @@ -30,13 +30,14 @@ FORCE_REBUILD="${FORCE_REBUILD:-false}" # run_docker.sh gives the container a fresh environment, so this is forwarded # explicitly on the in-container command line rather than exported on the host. SUBPROCESS_TIMEOUT="${AUTODATA_SUBPROCESS_TIMEOUT:-1200}" -# Container-local cache directory. run_docker.sh bind-mounts the host's -# $HOME/.cache into the container; on the CI runner that path is root-owned (or -# auto-created as root), so the recreated non-root container user cannot write -# it -- warp fails to create ~/.cache/warp with a PermissionError. Pointing the -# cache env vars at a writable, container-local /tmp path sidesteps the mounted -# host cache entirely. The tools create these dirs themselves (makedirs). +# Writable container-local cache for Warp and tools that honor XDG_CACHE_HOME. +# Arena's hard-coded home cache is handled separately below. CONTAINER_CACHE_DIR="${CONTAINER_CACHE_DIR:-/tmp/autodata_ci_cache}" +# Arena's composed USD cache ignores XDG_CACHE_HOME and writes under ~/.cache. +# Mount a fresh directory owned by the runner instead of its shared host cache. +# RUNNER_TEMP is cleaned by Actions; local invocations fall back to TMPDIR or /tmp. +AUTODATA_DOCKER_CACHE_DIR="${AUTODATA_DOCKER_CACHE_DIR:-$(mktemp -d "${RUNNER_TEMP:-${TMPDIR:-/tmp}}/autodata-ci-cache.XXXXXX")}" +export AUTODATA_DOCKER_CACHE_DIR # Optional JUnit report dir (repo-relative). The repo is bind-mounted, so a report # written here lands on the host for CI to collect. RESULTS_DIR="${AUTODATA_RESULTS_DIR-}" @@ -56,8 +57,8 @@ fi # container (run_docker.sh passes trailing args through verbatim, and the values # are inherited by the data-generation subprocess pytest spawns). Call # /isaac-sim/python.sh -m pytest explicitly so it does not depend on the -# in-container `pytest` alias expanding. XDG_CACHE_HOME/WARP_CACHE_PATH steer all -# caches away from the mounted host $HOME/.cache (see CONTAINER_CACHE_DIR above). +# in-container `pytest` alias expanding. XDG_CACHE_HOME/WARP_CACHE_PATH keep those +# caches container-local; Arena uses the writable bind mount configured above. PYTEST_ARGS=( env "AUTODATA_SUBPROCESS_TIMEOUT=${SUBPROCESS_TIMEOUT}" diff --git a/scripts/generate_dataset.py b/scripts/generate_dataset.py index 2cbbbd2..812b418 100755 --- a/scripts/generate_dataset.py +++ b/scripts/generate_dataset.py @@ -87,6 +87,9 @@ ) parser.add_argument("--generation_num_trials", type=int, default=None, help="Number of demos to generate.") parser.add_argument("--num_envs", type=int, default=1, help="Number of parallel environments.") +parser.add_argument( + "--presets", choices=("physx",), default=None, help="Physics backend preset for Arena environments." +) parser.add_argument("--input_file", type=str, required=True, help="Source dataset HDF5 file.") parser.add_argument( "--output_file", @@ -146,6 +149,7 @@ from autodata_core import DataGenerator, get_algorithm # noqa: E402 from autodata_core.algorithms import REGISTERED_ALGORITHMS # noqa: E402 +from autodata_examples.envs import register_environment_for_run # noqa: E402 from autodata_interfaces.datastream import Datastream # noqa: E402 from autodata_interfaces.env import ( # noqa: E402 EnvironmentProfile, @@ -311,6 +315,19 @@ def main() -> None: if args_cli.generation_num_trials is not None: generation_policy_params.num_trials = args_cli.generation_num_trials + random.seed(generation_policy_params.seed) + np.random.seed(generation_policy_params.seed) + torch.manual_seed(generation_policy_params.seed) + + env_make_kwargs = register_environment_for_run( + env_name=env_name, + enable_cameras=args_cli.record_images, + num_envs=args_cli.num_envs, + device=args_cli.device, + seed=generation_policy_params.seed, + presets=args_cli.presets, + ) + # The algorithm's start-signal expectation is a class attribute, so we resolve it before # instantiating (the Datastream needs it, and SkillGen can only be instantiated once the # planners exist, which in turn need the Datastream). @@ -333,11 +350,7 @@ def main() -> None: enable_cameras=args_cli.record_images, ) - env = gym.make(env_name, cfg=env_cfg).unwrapped - - random.seed(generation_policy_params.seed) - np.random.seed(generation_policy_params.seed) - torch.manual_seed(generation_policy_params.seed) + env = gym.make(env_name, cfg=env_cfg, **env_make_kwargs.get(env_name, {})).unwrapped env.reset()