diff --git a/autodata_examples/envs/isaac_lab_arena/franka_rope/assets/RopeNewton.usda b/autodata_examples/envs/isaac_lab_arena/franka_rope/assets/RopeNewton.usda new file mode 100644 index 0000000..87de490 --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/assets/RopeNewton.usda @@ -0,0 +1,19 @@ +#usda 1.0 +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +( + defaultPrim = "World" +) + +# Keep the shared simulation topology, round visual mesh, and material. +# Newton uses the simulation tetrahedra directly; the PhysX collision mesh is redundant. +def Xform "World" ( + prepend references = @Rope.usd@ +) +{ + over "collision_mesh" ( + active = false + ) + { + } +} diff --git a/autodata_examples/envs/isaac_lab_arena/franka_rope/embodiment.py b/autodata_examples/envs/isaac_lab_arena/franka_rope/embodiment.py index e75c49a..b831c80 100644 --- a/autodata_examples/envs/isaac_lab_arena/franka_rope/embodiment.py +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/embodiment.py @@ -19,6 +19,7 @@ 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.physics_backend import PhysicsBackend from isaaclab_arena.utils.pose import Pose from isaaclab_physx.renderers import IsaacRtxRendererCfg, IsaacRtxRendererGlobalSettingsCfg @@ -159,6 +160,35 @@ def __init__( self.event_config = FrankaRopeEventsCfg() self.reward_config = None + def _configure_physics_backend(self, backend: PhysicsBackend) -> None: + """Apply Newton gravity compensation and mesh contact settings to the Franka.""" + if backend is PhysicsBackend.NEWTON: + from isaaclab_newton.renderers import NewtonWarpRendererCfg + from isaaclab_newton.sim.schemas import ( + MujocoRigidBodyCfg, + NewtonMaterialPropertiesCfg, + NewtonSDFCollisionPropertiesCfg, + ) + + from .newton_cfg import spawn_rope_franka_newton + + spawn = self.scene_config.robot.spawn + # Match the shared high-PD Franka's PhysX-only disable_gravity setting. + spawn.rigid_props = MujocoRigidBodyCfg(gravcomp=1.0) + spawn.make_uninstanceable = True + spawn.func = spawn_rope_franka_newton + # Full-surface soft contacts require SDFs for the mesh colliders. + spawn.collision_props = NewtonSDFCollisionPropertiesCfg(sdf_max_resolution=64) + spawn.physics_material = NewtonMaterialPropertiesCfg( + static_friction=3.0, dynamic_friction=3.0, restitution=0.0 + ) + # Maintain pose tracking under the deformable solver's contact feedback. + self.scene_config.robot.actuators["panda_shoulder"].stiffness = 800.0 + self.scene_config.robot.actuators["panda_forearm"].stiffness = 800.0 + # Read rigid poses and deformable triangles directly from Newton state. + for camera_name in self.camera_config.camera_names(): + getattr(self.camera_config, camera_name).renderer_cfg = NewtonWarpRendererCfg() + @register_retargeter class FrankaRopeKeyboardRetargeter(RetargetterBase): diff --git a/autodata_examples/envs/isaac_lab_arena/franka_rope/environment.py b/autodata_examples/envs/isaac_lab_arena/franka_rope/environment.py index 1781e57..6f6f633 100644 --- a/autodata_examples/envs/isaac_lab_arena/franka_rope/environment.py +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/environment.py @@ -12,7 +12,7 @@ from isaaclab_arena.environments.arena_environment_factory import ArenaEnvironmentCfg, ArenaEnvironmentFactory from isaaclab_arena.utils.pose import Pose -from .rope_asset import FrankaRopeAsset +from .rope_asset import FrankaRopeAsset, rope_newton_spawn_cfg from .task import FrankaRopeTask if TYPE_CHECKING: @@ -31,13 +31,35 @@ class FrankaRopeArenaEnvironmentCfg(ArenaEnvironmentCfg): def _configure_franka_rope_env(env_cfg: Any) -> Any: - """Apply rope timing and asset settings for PhysX.""" + """Apply rope timing and asset/solver settings for Arena's selected backend.""" + + from isaaclab_newton.physics import NewtonCfg 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 + if isinstance(env_cfg.sim.physics, NewtonCfg): + from isaaclab_newton.sim.schemas import MujocoCollisionCfg + + from .newton_cfg import rope_newton_physics_cfg, spawn_rope_table_newton + + env_cfg.sim.physics = rope_newton_physics_cfg() + # Refresh Kit once per control step; the solver still advances at 1 ms. + env_cfg.sim.render_interval = env_cfg.decimation + env_cfg.scene.object.spawn = rope_newton_spawn_cfg() + # A discontinuous half-mesh reset can invert the connecting tetrahedra. + # Blend the same sampled poses through the middle half of the rope. + env_cfg.events.reset_object_position.params["transition_length"] = 0.4 + env_cfg.scene.replicate_physics = True + env_cfg.scene.table.spawn.make_uninstanceable = True + env_cfg.scene.table.spawn.func = spawn_rope_table_newton + # A firm, critically damped table contact in MJWarp. Priority keeps the + # table response from being softened by the gripper's default material. + # These MuJoCo parameters do not change VBD's rope contact properties. + env_cfg.scene.table.spawn.collision_props = MujocoCollisionCfg(priority=1, solref=(0.004, 1.0)) + else: + env_cfg.scene.replicate_physics = False env_cfg.teleop_devices = DevicesCfg( devices={ "keyboard": Se3KeyboardCfg( 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 index 2297c31..b2cf0f1 100644 --- a/autodata_examples/envs/isaac_lab_arena/franka_rope/mdp/events.py +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/mdp/events.py @@ -78,6 +78,7 @@ def reset_rope_nodal_state( pose_range: dict[str, tuple[float, float]], partial_pose_range: dict[str, tuple[float, float]], asset_cfg: SceneEntityCfg = SceneEntityCfg("object"), + transition_length: float | None = None, ) -> None: """Reset and randomize the rope, including a transform of one rope half. @@ -87,6 +88,8 @@ def reset_rope_nodal_state( 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. + transition_length: Length [m] over which to blend the half-rope transform. + None preserves the original piecewise reset used by PhysX. """ rope: DeformableObject = env.scene[asset_cfg.name] @@ -101,6 +104,24 @@ def reset_rope_nodal_state( ) 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) + if transition_length is None: + nodal_state[..., node_ids, :3] = rope.transform_nodal_pos(nodal_state[..., node_ids, :3], position, quaternion) + else: + assert transition_length > 0.0 + nodes = nodal_state[..., :3] + center = nodes.mean(dim=1, keepdim=True) + half_center = nodes[:, node_ids].mean(dim=1, keepdim=True) + axis = torch.nn.functional.normalize(half_center - center, dim=-1) + distance = ((nodes - center) * axis).sum(dim=-1, keepdim=True) + fraction = (distance / transition_length + 0.5).clamp(0.0, 1.0) + # A C2-continuous transition avoids tearing/inverting the tetrahedra at + # the boundary of the two halves. End sections retain their rigid poses. + weight = fraction**3 * (10.0 - 15.0 * fraction + 6.0 * fraction**2) + transformed = ( + math_utils.quat_apply(quaternion[:, None, :].expand(-1, nodes.shape[1], -1), nodes - half_center) + + half_center + + position[:, None, :] + ) + nodal_state[..., :3] = nodes + weight * (transformed - nodes) rope.write_nodal_state_to_sim_index(nodal_state, env_ids=env_ids) diff --git a/autodata_examples/envs/isaac_lab_arena/franka_rope/newton_cfg.py b/autodata_examples/envs/isaac_lab_arena/franka_rope/newton_cfg.py new file mode 100644 index 0000000..6722bec --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/newton_cfg.py @@ -0,0 +1,141 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Newton rigid--deformable physics for the shared Franka rope task.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from isaaclab.sim.spawners.from_files import UsdFileCfg, spawn_from_usd +from isaaclab.sim.utils import clone +from isaaclab_contrib.coupling import CouplerEntryCfg, CouplerProxyCfg, CouplerProxyMappingCfg +from isaaclab_newton.physics import ( + MJWarpSolverCfg, + NewtonCfg, + NewtonCollisionPipelineCfg, + NewtonSoftContactCfg, + VBDSolverCfg, +) + +if TYPE_CHECKING: + from pxr import Usd + + +@clone +def spawn_rope_franka_newton( + prim_path: str, + cfg: UsdFileCfg, + translation: tuple[float, float, float] | None = None, + orientation: tuple[float, float, float, float] | None = None, + **kwargs, +) -> Usd.Prim: + """Use analytic finger colliders for stable volumetric rope contacts. + + The boxes occupy the original collision meshes' bounds. Visual meshes and + articulation joints remain those of the shared Franka. + + Args: + prim_path: Prim path or pattern at which to spawn the robot. + cfg: Robot USD spawner configuration. + translation: Translation relative to the parent [m]. + orientation: Rotation quaternion in XYZW order. + **kwargs: Additional arguments forwarded to the USD spawner. + + Returns: + The spawned robot prim. + """ + from pxr import Gf, UsdGeom + + prim = spawn_from_usd(prim_path, cfg, translation, orientation, **kwargs) + for finger in ("panda_leftfinger", "panda_rightfinger"): + collider = prim.GetStage().GetPrimAtPath(f"{prim.GetPath()}/{finger}/collisions/collisions") + assert collider.IsValid() and collider.GetTypeName() == "Mesh", "Expected the legacy Franka finger collider" + points = UsdGeom.Mesh(collider).GetPointsAttr().Get() + lower = Gf.Vec3d(*(min(point[axis] for point in points) for axis in range(3))) + upper = Gf.Vec3d(*(max(point[axis] for point in points) for axis in range(3))) + collider.SetTypeName("Cube") + cube = UsdGeom.Cube(collider) + cube.CreateSizeAttr(1.0) + cube.AddTranslateOp(opSuffix="rope_contact_center").Set((lower + upper) * 0.5) + cube.AddScaleOp(opSuffix="rope_contact_size").Set(upper - lower) + return prim + + +@clone +def spawn_rope_table_newton( + prim_path: str, + cfg: UsdFileCfg, + translation: tuple[float, float, float] | None = None, + orientation: tuple[float, float, float, float] | None = None, + **kwargs, +) -> Usd.Prim: + """Use an analytic box for the Seattle table's existing unit-cube collider. + + Args: + prim_path: Prim path or pattern at which to spawn the table. + cfg: Table USD spawner configuration. + translation: Translation relative to the parent [m]. + orientation: Rotation quaternion in XYZW order. + **kwargs: Additional arguments forwarded to the USD spawner. + + Returns: + The spawned table prim. + """ + from pxr import UsdGeom + + prim = spawn_from_usd(prim_path, cfg, translation, orientation, **kwargs) + collider = prim.GetStage().GetPrimAtPath(f"{prim.GetPath()}/Collisions/Cube") + assert collider.IsValid() and collider.GetTypeName() == "Mesh", "Expected the Seattle table's cube mesh collider" + collider.SetTypeName("Cube") + UsdGeom.Cube(collider).CreateSizeAttr(1.0) + return prim + + +def rope_newton_physics_cfg() -> NewtonCfg: + """Couple the Franka's rigid dynamics to the rope's volumetric VBD solve.""" + return NewtonCfg( + solver_cfg=CouplerProxyCfg( + class_type="autodata_examples.envs.isaac_lab_arena.franka_rope.newton_visuals:NewtonRopeManager", + entries=[ + CouplerEntryCfg( + name="rigid", + solver_cfg=MJWarpSolverCfg(cone="elliptic", ls_iterations=20, integrator="implicitfast"), + bodies=[r"/World/envs/env_[^/]+/Robot"], + include_static_shapes=True, + ), + CouplerEntryCfg( + name="soft", + solver_cfg=VBDSolverCfg( + iterations=10, + rigid_body_particle_contact_buffer_size=4096, + particle_enable_self_contact=True, + particle_self_contact_radius=0.002, + particle_self_contact_margin=0.003, + ), + all_particles=True, + ), + ], + proxies=[ + CouplerProxyMappingCfg( + source="rigid", + destination="soft", + bodies=[ + r"/World/envs/env_[^/]+/Robot/.*panda_hand", + r"/World/envs/env_[^/]+/Robot/.*panda_(left|right)finger", + ], + mode="lagged", + collide_interval=1, + collision_pipeline=NewtonCollisionPipelineCfg(enable_rigid_soft_full_surface_contact=True), + ), + ], + iterations=1, + ), + soft_contact_cfg=NewtonSoftContactCfg( + # Resolve contacts around the narrow rope, with enough friction to carry it. + soft_contact_ke=1.0e5, + soft_contact_kd=10.0, + soft_contact_mu=10.0, + ), + num_substeps=10, + ) diff --git a/autodata_examples/envs/isaac_lab_arena/franka_rope/newton_visuals.py b/autodata_examples/envs/isaac_lab_arena/franka_rope/newton_visuals.py new file mode 100644 index 0000000..c03af54 --- /dev/null +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/newton_visuals.py @@ -0,0 +1,117 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Deform the rope's original round visual mesh for Newton's Kit viewport.""" + +import numpy as np + +import warp as wp +from isaaclab.scene_data.deformable_vis_remap import VolumeVisRemap, build_volume_vis_barycentric_remap +from isaaclab.sim.utils import find_matching_prim_paths +from isaaclab.sim.utils.stage import get_current_stage +from isaaclab_contrib.coupling.coupler import NewtonCouplerManager + + +@wp.kernel(enable_backward=False) +def _skin_rope_points( + points: wp.fabricarrayarray(dtype=wp.vec3f), + matrices: wp.fabricarray(dtype=wp.mat44d), + offsets: wp.fabricarray(dtype=wp.uint32), + particles: wp.array(dtype=wp.vec3f), + indices: wp.array2d(dtype=wp.int32), + weights: wp.array2d(dtype=wp.float32), +): + instance, vertex = wp.tid() + offset = int(offsets[instance]) + world = wp.vec3f(0.0) + for corner in range(4): + world += weights[vertex, corner] * particles[offset + indices[vertex, corner]] + world_to_local = wp.inverse(wp.transpose(wp.mat44f(matrices[instance]))) + points[instance][vertex] = wp.transform_point(world_to_local, world) + + +class NewtonRopeManager(NewtonCouplerManager): + """Add visual skinning to the standard rigid/deformable Newton coupler. + + The pinned Lab backend copies simulation nodes into visual vertices by index. + This rope has different simulation and visual topologies. Embed its original + round mesh in the simulation tetrahedra and update every visual vertex before + Kit renders. The solver, collision geometry, and particle state are unchanged. + """ + + _rope_vis_remap: VolumeVisRemap | None = None + _rope_offset_attr = "autodata:rope:particleOffset" + + @classmethod + def _initialize_rope_skin(cls) -> None: + import usdrt + from pxr import Gf, UsdGeom + + assert len(cls._deformable_registry) == 1, "The rope task expects one deformable asset" + entry = cls._deformable_registry[0] + stage = get_current_stage() + paths = find_matching_prim_paths(entry.vis_mesh_prim_path) + assert paths, "Rope visual mesh is missing" + visual = UsdGeom.Mesh(stage.GetPrimAtPath(paths[0])) + simulation = UsdGeom.TetMesh( + stage.GetPrimAtPath(visual.GetPath().GetParentPath().AppendChild("simulation_mesh")) + ) + xforms = UsdGeom.XformCache() + visual_to_sim = ( + xforms.GetLocalToWorldTransform(visual.GetPrim()) + * xforms.GetLocalToWorldTransform(simulation.GetPrim()).GetInverse() + ) + visual_points = np.asarray( + [visual_to_sim.Transform(Gf.Vec3d(point)) for point in visual.GetPointsAttr().Get()], dtype=np.float32 + ) + cls._rope_vis_remap = build_volume_vis_barycentric_remap( + np.asarray(simulation.GetPointsAttr().Get(), dtype=np.float32), + np.asarray(simulation.GetTetVertexIndicesAttr().Get(), dtype=np.int32).reshape(-1), + visual_points, + device=str(cls._device), + ) + assert cls._rope_vis_remap is not None, "Cannot embed the rope visual mesh in its simulation tetrahedra" + for path in paths: + env_index = int(path.split("/envs/env_", 1)[1].split("/", 1)[0]) + prim = cls._usdrt_stage.GetPrimAtPath(path) + usdrt.Rt.Xformable(prim).SetWorldXformFromUsd() + prim.CreateAttribute(cls._rope_offset_attr, usdrt.Sdf.ValueTypeNames.UInt, True).Set( + entry.particle_offsets[env_index] + ) + + @classmethod + def pre_render(cls) -> None: + """Refresh the round skin before Kit renders, including after state restores.""" + super().pre_render() + if cls._usdrt_stage is None or cls._clone_physics_only or cls._state_0 is None: + return + import usdrt + + if cls._rope_vis_remap is None: + cls._initialize_rope_skin() + selection = cls._usdrt_stage.SelectPrims( + require_attrs=[ + (usdrt.Sdf.ValueTypeNames.Point3fArray, "points", usdrt.Usd.Access.ReadWrite), + (usdrt.Sdf.ValueTypeNames.UInt, cls._rope_offset_attr, usdrt.Usd.Access.Read), + (usdrt.Sdf.ValueTypeNames.Matrix4d, "omni:fabric:worldMatrix", usdrt.Usd.Access.Read), + ], + device=str(cls._device), + ) + wp.launch( + _skin_rope_points, + dim=(selection.GetCount(), cls._rope_vis_remap.tet_vertex_indices.shape[0]), + inputs=[ + wp.fabricarrayarray(data=selection, attrib="points", dtype=wp.vec3f), + wp.fabricarray(data=selection, attrib="omni:fabric:worldMatrix"), + wp.fabricarray(data=selection, attrib=cls._rope_offset_attr), + cls._state_0.particle_q, + cls._rope_vis_remap.tet_vertex_indices, + cls._rope_vis_remap.bary_weights, + ], + device=cls._device, + ) + + @classmethod + def _solver_specific_clear(cls) -> None: + cls._rope_vis_remap = None + super()._solver_specific_clear() 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 index 1eb21cd..33ff0c3 100644 --- a/autodata_examples/envs/isaac_lab_arena/franka_rope/rope_asset.py +++ b/autodata_examples/envs/isaac_lab_arena/franka_rope/rope_asset.py @@ -13,6 +13,28 @@ _ROPE_USD_PATH = Path(__file__).resolve().parent / "assets" / "Rope.usd" +def rope_newton_spawn_cfg() -> UsdFileCfg: + """Load the shared rope topology with Newton material properties.""" + from isaaclab_newton.sim.schemas import NewtonDeformableBodyPropertiesCfg + from isaaclab_newton.sim.spawners.materials import NewtonDeformableBodyMaterialCfg + + # VBD needs a compliant volume material for the gripper to bend this slender + # mesh. Keep the shared topology and density, with backend-specific elasticity. + youngs_modulus, poissons_ratio = 2.0e5, 0.45 + return UsdFileCfg( + usd_path=str(_ROPE_USD_PATH.with_name("RopeNewton.usda")), + deformable_props=NewtonDeformableBodyPropertiesCfg(), + physics_material_path="NewtonPhysicsMaterial", + physics_material=NewtonDeformableBodyMaterialCfg( + density=1000.0, + particle_radius=0.002, + k_mu=youngs_modulus / (2.0 * (1.0 + poissons_ratio)), + k_lambda=youngs_modulus * poissons_ratio / ((1.0 + poissons_ratio) * (1.0 - 2.0 * poissons_ratio)), + k_damp=1000.0, + ), + ) + + class FrankaRopeAsset(Asset): """Deformable rope asset that can participate in an Arena :class:`Scene`.""" diff --git a/autodata_tests/e2e/test_newton_rope_workflow.py b/autodata_tests/e2e/test_newton_rope_workflow.py new file mode 100644 index 0000000..a37b044 --- /dev/null +++ b/autodata_tests/e2e/test_newton_rope_workflow.py @@ -0,0 +1,184 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Record, annotate, and generate Newton demonstrations through the public tools.""" + +import h5py +import json +import numpy as np +import os +from pathlib import Path + +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 + + +def configure_newton_recording() -> list[str]: + """Feed reproducible SpaceMouse commands while keeping real physics and success checks.""" + import torch + + import __main__ as recorder + from isaaclab.devices import Se3SpaceMouse + from isaaclab.envs import ManagerBasedRLEnv + from isaaclab.utils.datasets import HDF5DatasetFileHandler + + from autodata_examples.envs import register_environments + + assert recorder.args_cli.teleop_device == "spacemouse" + assert not recorder.args_cli.reset_sim_buffer_each_episode + remaining_args = register_environments() + reader = HDF5DatasetFileHandler() + reader.open(str(Path(TestPaths.test_data_dir) / "annotated_dataset_franka_rope_softmimicgen_newton.hdf5")) + # Only the initial state and robot actions serve as test input. No annotations, + # subsequent states, or success labels are copied into the new recording. + episode = reader.load_episode("demo_0", "cpu") + reader.close() + actions = episode.data["actions"] + original_reset = ManagerBasedRLEnv.reset + + def reset_from_source(env, *args, **kwargs): + original_reset(env, *args, **kwargs) + assert type(env.scene["object"]).__module__.startswith("isaaclab_contrib.deformable") + return env.reset_to(episode.get_initial_state(), env_ids=None, is_relative=True) + + def reset_device(device): + device._test_action_index = 0 + + def advance(device): + index = device._test_action_index + device._test_action_index += 1 + assert index < len(actions) + 30, "Scripted recording did not reach the real task success condition" + if index < len(actions): + return actions[index].clone() + hold = torch.zeros_like(actions[-1]) + hold[-1] = actions[-1, -1] + return hold + + ManagerBasedRLEnv.reset = reset_from_source + Se3SpaceMouse._find_device = lambda self: None + Se3SpaceMouse._run_device = lambda self: None + Se3SpaceMouse.reset = reset_device + Se3SpaceMouse.advance = advance + return remaining_args + + +@pytest.mark.with_subprocess +def test_newton_rope_record_annotate_generate(tmp_path): + """Require newly recorded raw data to become usable SoftMimicGen input.""" + raw_file = tmp_path / "recorded.hdf5" + annotated_file = tmp_path / "annotated.hdf5" + generated_file = tmp_path / "generated.hdf5" + lab_record = ( + Path(TestPaths.repo_root) / "submodules/IsaacLab-Arena/submodules/IsaacLab/scripts/tools/record_demos.py" + ) + process_env = {**os.environ, "HEADLESS": "1"} + env_name = "Isaac-Rope-Franka-Arena-IK-Rel-v0" + camera_args = ["--viz", "kit", "--enable_cameras"] + task_args = [ + "--task_descriptor", + str(Path(TestPaths.tasks_dir) / "franka_rope_softmimicgen.yaml"), + "--embodiment", + str(Path(TestPaths.embodiments_dir) / "franka_ik_rel.yaml"), + ] + + run_subprocess( + [ + TestPaths.python_path, + str(lab_record), + "--external_callback", + "autodata_tests.e2e.test_newton_rope_workflow.configure_newton_recording", + "--task", + env_name, + "--presets", + "newton", + "--teleop_device", + "spacemouse", + "--no-reset_sim_buffer_each_episode", + "--dataset_file", + str(raw_file), + "--num_demos", + "2", + "--step_hz", + "20", + *camera_args, + ], + env=process_env, + ) + assert_valid_dataset(str(raw_file), min_num_demos=2) + with h5py.File(raw_file, "r") as raw: + assert len(raw["data"]) == 2 + for episode in raw["data"].values(): + assert episode.attrs["success"] + assert "datagen_info" not in episode["obs"] + + run_subprocess( + [ + TestPaths.python_path, + str(Path(TestPaths.scripts_dir) / "annotate_demos.py"), + "--auto", + "--presets", + "newton", + "--device", + "cuda", + "--input_file", + str(raw_file), + "--output_file", + str(annotated_file), + *task_args, + *camera_args, + ], + env=process_env, + ) + assert_valid_dataset(str(annotated_file), min_num_demos=2) + with h5py.File(raw_file, "r") as raw, h5py.File(annotated_file, "r") as annotated: + assert len(annotated["data"]) == 2 + assert json.loads(annotated["data"].attrs["env_args"])["env_name"] == env_name + for name, episode in annotated["data"].items(): + np.testing.assert_array_equal(episode["actions"][:], raw[f"data/{name}/actions"][:]) + info = episode["obs/datagen_info"] + signal = info["subtask_term_signals/grasp"][:] + assert len(signal) == len(episode["actions"]) + assert not signal[0] and signal[-1] + assert np.all(np.diff(signal.astype(int)) >= 0) + nodes = info["object_nodal_position/object"][:] + assert nodes.shape == (len(signal), 549, 3) + assert np.isfinite(nodes).all() + assert set(info["eef_pose"]) == {"franka"} + for camera in ("agentview_image_rgb", "robot0_eye_in_hand_image_rgb"): + assert episode[f"obs/{camera}"].shape[:3] == (len(signal), 128, 128) + + run_subprocess( + [ + TestPaths.python_path, + TestPaths.generate_dataset_script, + "--env_name", + env_name, + "--presets", + "newton", + "--alg", + "softmimicgen", + "--input_file", + str(annotated_file), + "--output_file", + str(generated_file), + "--generation_num_trials", + "1", + "--num_envs", + "3", + "--device", + "cuda", + *task_args, + *camera_args, + ], + env=process_env, + ) + assert_valid_dataset(str(generated_file), min_num_demos=1) + with h5py.File(generated_file, "r") as generated: + for episode in generated["data"].values(): + assert episode.attrs["success"] + nodes = episode["initial_state/deformable_object/object/nodal_position"][:] + assert np.isfinite(nodes).all() and np.abs(nodes).max() < 2.0 + assert not (tmp_path / "generated_failed.hdf5").exists() diff --git a/autodata_tests/e2e/test_rope_environment.py b/autodata_tests/e2e/test_rope_environment.py index 63abc7b..3cf2598 100644 --- a/autodata_tests/e2e/test_rope_environment.py +++ b/autodata_tests/e2e/test_rope_environment.py @@ -1,7 +1,7 @@ # 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.""" +"""Simulate the shared rope-folding task with each supported physics backend.""" import os @@ -12,8 +12,68 @@ from autodata_tests.utils.subprocess import run_subprocess +@wp.kernel +def _read_rope_visual_points( + points: wp.fabricarrayarray(dtype=wp.vec3f), + matrices: wp.fabricarray(dtype=wp.mat44d), + offsets: wp.fabricarray(dtype=wp.uint32), + result: wp.array2d(dtype=wp.vec3f), + sim_count: int, +): + i, node = wp.tid() + env_index = int(offsets[i]) // sim_count + result[env_index, node] = wp.transform_point(wp.transpose(wp.mat44f(matrices[i])), points[i][node]) + + +def _assert_kit_rope_visual(env): + """Compare Kit's actual GPU mesh points with the simulated rope in each world.""" + import torch + + import usdrt + from isaaclab.sim.utils.stage import get_current_stage + from pxr import UsdGeom + + env.sim.render() + nodes = env.scene["object"].data.nodal_pos_w.torch + stage = get_current_stage() + for env_index in range(env.num_envs): + root = f"/World/envs/env_{env_index}/Object" + visual = UsdGeom.Mesh(stage.GetPrimAtPath(f"{root}/Xform")) + assert visual.GetPrim().IsActive() + assert len(visual.GetPointsAttr().Get()) == 1056 + assert len(visual.GetFaceVertexCountsAttr().Get()) == 1026 + # CPU attribute reads are not valid once Fabric owns the point data on the GPU. + selection = get_current_stage(fabric=True).SelectPrims( + require_attrs=[ + (usdrt.Sdf.ValueTypeNames.Point3fArray, "points", usdrt.Usd.Access.Read), + (usdrt.Sdf.ValueTypeNames.UInt, "autodata:rope:particleOffset", usdrt.Usd.Access.Read), + (usdrt.Sdf.ValueTypeNames.Matrix4d, "omni:fabric:worldMatrix", usdrt.Usd.Access.Read), + ], + device=env.device, + ) + assert selection.GetCount() == env.num_envs, f"Only {selection.GetCount()}/{env.num_envs} rope skins found" + visual_points = wp.empty((env.num_envs, 1056), dtype=wp.vec3f, device=env.device) + wp.launch( + _read_rope_visual_points, + dim=visual_points.shape, + inputs=[ + wp.fabricarrayarray(data=selection, attrib="points", dtype=wp.vec3f), + wp.fabricarray(data=selection, attrib="omni:fabric:worldMatrix"), + wp.fabricarray(data=selection, attrib="autodata:rope:particleOffset"), + visual_points, + nodes.shape[1], + ], + device=env.device, + ) + remap = env.sim.physics_manager._rope_vis_remap + indices = wp.to_torch(remap.tet_vertex_indices).long() + weights = wp.to_torch(remap.bary_weights) + expected = (nodes[:, indices] * weights[None, :, :, None]).sum(2) + torch.testing.assert_close(wp.to_torch(visual_points), expected, atol=1.0e-5, rtol=0) + + @pytest.mark.with_subprocess -@pytest.mark.parametrize("presets", ["physx"]) +@pytest.mark.parametrize("presets", ["physx", "newton"]) @pytest.mark.parametrize("num_envs", [1, 3]) def test_rope_environment(presets, num_envs, tmp_path): result_file = tmp_path / "passed" @@ -37,6 +97,102 @@ def test_rope_environment(presets, num_envs, tmp_path): assert result_file.exists(), "Simulation exited before completing the checks" +@pytest.mark.with_subprocess +def test_rope_gripper_table_contacts_newton(tmp_path): + """Require physical table support under downward commands with open and closed fingers.""" + result_file = tmp_path / "table_contacts_passed" + run_subprocess( + [ + TestPaths.python_path, + "-m", + __name__, + "--presets", + "newton", + "--num_envs", + "1", + "--viz", + "kit", + "--enable_cameras", + "--table_contacts", + "--result_file", + str(result_file), + ], + env={**os.environ, "HEADLESS": "1"}, + ) + assert result_file.exists(), "Simulation exited before completing the table-contact checks" + + +def _assert_newton_gripper_table_contacts(env): + import itertools + import math + import torch + + from isaaclab.utils.math import compute_pose_error, quat_apply, quat_apply_inverse + from newton import ShapeFlags + + assert env.num_envs == 1 + model = env.sim.physics_manager.get_model() + shape_names = list(map(str, model.shape_label)) + finger_ids = [i for i, name in enumerate(shape_names) if "finger/collisions" in name] + table_ids = [i for i, name in enumerate(shape_names) if "/table/Collisions/Cube" in name] + assert len(finger_ids) == 2 and len(table_ids) == 1 + for entry in ("rigid", "soft"): + view = env.sim.physics_manager._solver.view(entry) + table_flags = [ + int(view.shape_flags.numpy()[i]) + for i, name in enumerate(view.shape_label) + if "/table/Collisions/Cube" in str(name) + ] + assert table_flags and all(flags & int(ShapeFlags.COLLIDE_SHAPES) for flags in table_flags) + + shape_poses = wp.to_torch(model.shape_transform) + shape_scales = wp.to_torch(model.shape_scale) + shape_bodies = wp.to_torch(model.shape_body).long() + signs = torch.tensor(list(itertools.product((-1.0, 1.0), repeat=3)), device=env.device) + + def box_corners(shape_ids): + poses = shape_poses[shape_ids] + points = quat_apply( + poses[:, None, 3:].expand(-1, 8, -1).reshape(-1, 4), + (shape_scales[shape_ids, None] * signs).reshape(-1, 3), + ).reshape(-1, 8, 3) + points += poses[:, None, :3] + bodies = wp.to_torch(env.sim.physics_manager.get_state_0().body_q) + for index, body in enumerate(shape_bodies[shape_ids].tolist()): + if body >= 0: + points[index] = quat_apply(bodies[body, 3:].expand(8, -1), points[index]) + bodies[body, :3] + return points + + table_top = box_corners(table_ids)[..., 2].max().item() + robot = env.scene["robot"] + action = torch.zeros((1, 7), device=env.device) + orientation = torch.tensor([[math.sqrt(0.5), -math.sqrt(0.5), 0.0, 0.0]], device=env.device) + with torch.inference_mode(): + action[:, -1] = 1.0 + env.step(action) # Refresh the frame sensor after reset. + for gripper in (1.0, -1.0): + for height, count in ((0.13, 60), (-0.04, 75)): + target = torch.tensor([[0.45, -0.35, height]], device=env.device) + max_penetration = 0.0 + for _ in range(count): + frame = env.scene["ee_frame"].data + pos, quat = frame.target_pos_w.torch[:, 0], frame.target_quat_w.torch[:, 0] + _, rotation = compute_pose_error(pos, quat, target, orientation, rot_error_type="axis_angle") + root_quat = robot.data.root_quat_w.torch + action[:, :3] = (quat_apply_inverse(root_quat, target - pos) / 0.5).clamp(-0.04, 0.04) + action[:, 3:6] = (quat_apply_inverse(root_quat, rotation) / 0.5).clamp(-0.2, 0.2) + action[:, -1] = gripper + env.step(action) + max_penetration = max(max_penetration, table_top - box_corners(finger_ids)[..., 2].min().item()) + assert torch.isfinite(env.scene["object"].data.nodal_pos_w.torch).all() + if height < table_top: + pos = env.scene["ee_frame"].data.target_pos_w.torch[0, 0] + assert torch.linalg.vector_norm(pos[:2] - target[0, :2]) < 0.025, "Missed the test contact area" + assert pos[2] < 0.04, "Gripper never reached the tabletop" + assert max_penetration < 0.001, f"Fingers penetrated the table by {max_penetration:.6f} m" + print(f"ROPE table contact: gripper={gripper}, max_penetration={max_penetration:.6f} m", flush=True) + + def main() -> None: """Check spawning, stepping, contacts, observations and indexed resets.""" import argparse @@ -46,10 +202,11 @@ def main() -> None: from isaaclab.app import AppLauncher parser = argparse.ArgumentParser() - parser.add_argument("--presets", choices=("physx",), required=True) + parser.add_argument("--presets", choices=("physx", "newton"), 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) + parser.add_argument("--table_contacts", action="store_true") AppLauncher.add_app_launcher_args(parser) parser.add_argument("--enable_cameras", action="store_true") args = parser.parse_args() @@ -67,9 +224,21 @@ def main() -> None: 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_") + if args.table_contacts: + assert args.presets == "newton" + cfg.recorders = {} + cfg.terminations = {} env = gym.make(FRANKA_ROPE_ARENA_ENV_ID, cfg=cfg).unwrapped exit_code = 0 try: + if args.table_contacts: + env.reset() + _assert_newton_gripper_table_contacts(env) + if args.result_file: + from pathlib import Path + + Path(args.result_file).touch() + return def camera_observations(): # Headless Kit renders on demand in the pinned Lab version. Request @@ -84,10 +253,12 @@ def camera_observations(): return env.observation_manager.compute()["camera_obs"] env.reset() + if args.presets == "newton": + _assert_kit_rope_visual(env) 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" + expected_module = "isaaclab_physx" if args.presets == "physx" else "isaaclab_contrib.deformable" 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() @@ -106,6 +277,8 @@ def camera_observations(): 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) + if args.presets == "newton": + _assert_kit_rope_visual(env) 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. @@ -126,7 +299,7 @@ def camera_observations(): 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. + # Relative IK and gripper actions must drive the Newton/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 @@ -149,6 +322,8 @@ def camera_observations(): 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) + if args.presets == "newton": + _assert_kit_rope_visual(env) # Restoring a saved state must preserve nodal coordinates, including env offsets. state = env.scene.get_state(is_relative=True) @@ -156,6 +331,8 @@ def camera_observations(): 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) + if args.presets == "newton": + _assert_kit_rope_visual(env) actions.zero_() actions[:, -1] = 1.0 for _ in range(5): diff --git a/autodata_tests/e2e/test_softmimicgen_data_generation.py b/autodata_tests/e2e/test_softmimicgen_data_generation.py index f3b06f8..eef0e0f 100644 --- a/autodata_tests/e2e/test_softmimicgen_data_generation.py +++ b/autodata_tests/e2e/test_softmimicgen_data_generation.py @@ -21,6 +21,7 @@ def _run_franka_rope_softmimicgen(num_envs: int, device: str, presets: str = "ph with tempfile.TemporaryDirectory() as temp_dir: output_file = os.path.join(temp_dir, "generated.hdf5") + dataset_suffix = "_newton" if presets == "newton" else "" args = [ TestPaths.python_path, @@ -36,7 +37,7 @@ def _run_franka_rope_softmimicgen(num_envs: int, device: str, presets: str = "ph "--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"), + os.path.join(TestPaths.test_data_dir, f"annotated_dataset_franka_rope_softmimicgen{dataset_suffix}.hdf5"), "--output_file", output_file, "--generation_num_trials", @@ -66,20 +67,20 @@ def _run_franka_rope_softmimicgen(num_envs: int, device: str, presets: str = "ph @pytest.mark.with_subprocess -@pytest.mark.parametrize("presets", ["physx"]) +@pytest.mark.parametrize("presets", ["physx", "newton"]) 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"]) +@pytest.mark.parametrize("presets", ["physx", "newton"]) 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",): + for presets in ("physx", "newton"): 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 index 86696e9..bd2b56b 100644 --- a/autodata_tests/interfaces/env/test_franka_rope_arena.py +++ b/autodata_tests/interfaces/env/test_franka_rope_arena.py @@ -25,7 +25,7 @@ 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.rope_asset import FrankaRopeAsset, rope_newton_spawn_cfg from autodata_examples.envs.isaac_lab_arena.franka_rope.task import FrankaRopeTask @@ -40,7 +40,7 @@ def test_arena_environment_factory_is_registered(): assert registry.get_environment_cfg_type(FrankaRopeArenaEnvironment) is FrankaRopeArenaEnvironmentCfg -@pytest.mark.parametrize("presets", [None, "physx"]) +@pytest.mark.parametrize("presets", [None, "physx", "newton"]) def test_arena_registration_uses_run_configuration_and_preserves_make_kwargs(monkeypatch, presets): captured = {} variation_recorder = object() @@ -169,6 +169,28 @@ def test_rope_asset_uses_deformable_scene_key(): assert object_cfg.spawn.usd_path.endswith("/isaac_lab_arena/franka_rope/assets/Rope.usd") +def test_newton_rope_preserves_simulation_topology(): + from pxr import Usd + + physx_stage = Usd.Stage.Open(FrankaRopeAsset().object_cfg.spawn.usd_path) + newton_stage = Usd.Stage.Open(rope_newton_spawn_cfg().usd_path) + for attribute in ("points", "tetVertexIndices", "omniphysics:restShapePoints", "omniphysics:restTetVtxIndices"): + original = physx_stage.GetPrimAtPath("/World/simulation_mesh").GetAttribute(attribute).Get() + adapted = newton_stage.GetPrimAtPath("/World/simulation_mesh").GetAttribute(attribute).Get() + assert original is not None + assert original == adapted + assert len(newton_stage.GetPrimAtPath("/World/simulation_mesh").GetAttribute("points").Get()) == 549 + assert physx_stage.GetPrimAtPath("/World/collision_mesh").IsActive() + assert not newton_stage.GetPrimAtPath("/World/collision_mesh").IsActive() + assert physx_stage.GetPrimAtPath("/World/Xform").IsActive() + assert newton_stage.GetPrimAtPath("/World/Xform").IsActive() + for attribute in ("points", "faceVertexCounts", "faceVertexIndices", "primvars:st"): + original = physx_stage.GetPrimAtPath("/World/Xform").GetAttribute(attribute).Get() + adapted = newton_stage.GetPrimAtPath("/World/Xform").GetAttribute(attribute).Get() + assert adapted == original + assert sum(prim.GetTypeName() == "TetMesh" for prim in newton_stage.Traverse()) == 1 + + 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) @@ -211,3 +233,116 @@ def test_rope_embodiment_preserves_arena_observation_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" + + +def test_newton_embodiment_compensates_gravity_without_changing_physx(): + physx = FrankaRopeIKEmbodiment() + newton = FrankaRopeIKEmbodiment() + physx.configure_physics_backend(PhysicsBackend.PHYSX) + newton.configure_physics_backend(PhysicsBackend.NEWTON) + assert physx.scene_config.robot.spawn.rigid_props.disable_gravity is True + assert physx.scene_config.robot.spawn.collision_props is None + assert newton.scene_config.robot.spawn.rigid_props.gravcomp == 1.0 + assert newton.scene_config.robot.spawn.make_uninstanceable is True + assert newton.scene_config.robot.spawn.collision_props.sdf_max_resolution == 64 + assert physx.scene_config.robot.actuators["panda_shoulder"].stiffness == 400.0 + assert newton.scene_config.robot.actuators["panda_shoulder"].stiffness == 800.0 + assert physx.action_config.gripper_action.close_command_expr == {"panda_finger_.*": 0.0} + assert newton.action_config.gripper_action.close_command_expr == {"panda_finger_.*": 0.0} + for name in newton.camera_config.camera_names(): + assert getattr(newton.camera_config, name).renderer_cfg.renderer_type == "newton_warp" + assert getattr(physx.camera_config, name).renderer_cfg.renderer_type == "isaac_rtx" + + +def test_newton_rope_reset_preserves_tetrahedra(monkeypatch): + """Reset poses must not invert or excessively stretch the volume mesh.""" + import itertools + import numpy as np + import torch + + from isaaclab.assets.deformable_object.base_deformable_object import BaseDeformableObject + from isaaclab.utils import math as math_utils + from pxr import Gf, Usd, UsdGeom + + from autodata_examples.envs.isaac_lab_arena.franka_rope.mdp import events + + stage = Usd.Stage.Open(FrankaRopeAsset().object_cfg.spawn.usd_path) + mesh = UsdGeom.TetMesh(stage.GetPrimAtPath("/World/simulation_mesh")) + transform = UsdGeom.XformCache().GetLocalToWorldTransform(mesh.GetPrim()) + nodes = torch.tensor( + np.asarray([transform.Transform(Gf.Vec3d(point)) for point in mesh.GetPointsAttr().Get()]), + dtype=torch.float32, + ) + # Match the task's initial orientation: the rope's long axis is world Y. + nodes = math_utils.quat_apply(torch.tensor([0.0, 0.0, 2**-0.5, 2**-0.5]).expand(len(nodes), -1), nodes) + corners = torch.tensor(list(itertools.product((0.0, 0.1), (0.0, 0.1), (-torch.pi / 6, torch.pi / 6)))) + count = len(corners) + state = torch.zeros(count, len(nodes), 6) + state[..., :3] = nodes + written = {} + rope = SimpleNamespace(device="cpu", data=SimpleNamespace(default_nodal_state_w=SimpleNamespace(torch=state))) + rope.transform_nodal_pos = lambda *args: BaseDeformableObject.transform_nodal_pos(rope, *args) + rope.write_nodal_state_to_sim_index = lambda values, env_ids: written.update(state=values, env_ids=env_ids) + translations = torch.zeros(count, 3) + translations[:, :2] = corners[:, :2] + rotations = math_utils.quat_from_euler_xyz(torch.zeros(count), torch.zeros(count), corners[:, 2]) + samples = iter( + ((torch.zeros(count, 3), torch.tensor([0.0, 0.0, 0.0, 1.0]).repeat(count, 1)), (translations, rotations)) + ) + monkeypatch.setattr(events, "_sample_pose_delta", lambda *_: next(samples)) + env_ids = torch.arange(count) + events.reset_rope_nodal_state(SimpleNamespace(scene={"object": rope}), env_ids, {}, {}, transition_length=0.4) + indices = torch.tensor(np.asarray(mesh.GetTetVertexIndicesAttr().Get()), dtype=torch.long) + rest = nodes[indices[:, 1:]] - nodes[indices[:, :1]] + result = written["state"][..., :3] + deformed = result[:, indices[:, 1:]] - result[:, indices[:, :1]] + gradient = torch.linalg.solve(rest.transpose(-1, -2), deformed.transpose(-1, -2), left=False) + assert torch.linalg.det(gradient).min() > 0.25 + assert torch.linalg.svdvals(gradient).max() < 2.0 + torch.testing.assert_close(written["state"][..., 3:], torch.zeros_like(state[..., 3:])) + torch.testing.assert_close(written["env_ids"], env_ids) + torch.testing.assert_close(state[..., :3], nodes.expand(count, -1, -1)) + + +def test_rope_reset_smoothing_is_newton_only(): + from isaaclab_newton.physics import NewtonCfg + + for physics in (None, NewtonCfg()): + task = FrankaRopeTask() + cfg = SimpleNamespace( + sim=SimulationCfg() if physics is None else SimulationCfg(physics=physics), + scene=SimpleNamespace( + object=FrankaRopeAsset().object_cfg, + table=SimpleNamespace(spawn=SimpleNamespace()), + ), + events=task.get_events_cfg(), + ) + _configure_franka_rope_env(cfg) + params = cfg.events.reset_object_position.params + if physics is None: + assert "transition_length" not in params + else: + assert params["transition_length"] == 0.4 + assert "transition_length" not in FrankaRopeTask().get_events_cfg().reset_object_position.params + + +def test_newton_rope_contact_buffer_fits_full_surface(): + from collections import Counter + from itertools import combinations + + from pxr import Usd, UsdGeom + + from autodata_examples.envs.isaac_lab_arena.franka_rope.newton_cfg import rope_newton_physics_cfg + + stage = Usd.Stage.Open(rope_newton_spawn_cfg().usd_path) + mesh = UsdGeom.TetMesh(stage.GetPrimAtPath("/World/simulation_mesh")) + face_counts = Counter( + tuple(sorted(face)) for tet in mesh.GetTetVertexIndicesAttr().Get() for face in combinations(tet, 3) + ) + surface_faces = [face for face, count in face_counts.items() if count == 1] + surface_edges = {edge for face in surface_faces for edge in combinations(face, 2)} + contact_count = len(mesh.GetPointsAttr().Get()) + len(surface_faces) + len(surface_edges) + soft_solver = next( + entry.solver_cfg for entry in rope_newton_physics_cfg().solver_cfg.entries if entry.name == "soft" + ) + assert soft_solver.rigid_body_particle_contact_buffer_size >= contact_count diff --git a/autodata_tests/test_data/annotated_dataset_franka_rope_softmimicgen_newton.hdf5 b/autodata_tests/test_data/annotated_dataset_franka_rope_softmimicgen_newton.hdf5 new file mode 100644 index 0000000..54b4f86 --- /dev/null +++ b/autodata_tests/test_data/annotated_dataset_franka_rope_softmimicgen_newton.hdf5 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c56d7d7856b07a7b87255b26ba4a0334dde8e2ce548b88e1215d61798732907a +size 13597900 diff --git a/scripts/annotate_demos.py b/scripts/annotate_demos.py index 83930ec..45ee7f1 100644 --- a/scripts/annotate_demos.py +++ b/scripts/annotate_demos.py @@ -46,7 +46,13 @@ 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." + "--presets", choices=("physx", "newton"), default=None, help="Physics backend preset for Arena environments." +) +parser.add_argument( + "--reset_sim_buffer_each_episode", + action=argparse.BooleanOptionalAction, + default=None, + help="Flush simulation buffers before replay. Defaults to enabled for PhysX and disabled for Newton.", ) parser.add_argument( "--task_descriptor", @@ -288,6 +294,12 @@ def main() -> int: recorder_cfg=AnnotationRecorderManagerCfg(), enable_cameras=args_cli.record_images, ) + if args_cli.reset_sim_buffer_each_episode is None: + from isaaclab_newton.physics import NewtonCfg + + # Preserve the established PhysX replay behavior. Newton keeps solver buffers + # across episode restores, as in Lab's public replay tool. + args_cli.reset_sim_buffer_each_episode = not isinstance(env_cfg.sim.physics, NewtonCfg) if args_cli.record_images: env_cfg.recorders.record_pre_step_camera_observations = RecorderTermCfg( class_type=PreStepCameraObservationsRecorder @@ -397,8 +409,9 @@ def replay_episode( initial_state = episode.data["initial_state"] actions = episode.data["actions"] - # Reset the simulation before restoring state (use of env.sim.reset() for deterministic replay) - env.sim.reset() + # A full solver reset is distinct from restoring an episode's scene state. + if args_cli.reset_sim_buffer_each_episode: + env.sim.reset() env.recorder_manager.reset() env.reset_to(initial_state, torch.tensor([0], device=env.device), is_relative=True) diff --git a/scripts/generate_dataset.py b/scripts/generate_dataset.py index 812b418..b5e9f64 100755 --- a/scripts/generate_dataset.py +++ b/scripts/generate_dataset.py @@ -88,7 +88,7 @@ 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." + "--presets", choices=("physx", "newton"), 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(