|
| 1 | +# Copyright (c) 2025-2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md). |
| 2 | +# All rights reserved. |
| 3 | +# |
| 4 | +# SPDX-License-Identifier: Apache-2.0 |
| 5 | + |
| 6 | +import torch |
| 7 | +import traceback |
| 8 | + |
| 9 | +import pytest |
| 10 | +import warp as wp |
| 11 | + |
| 12 | +from isaaclab_arena.tests.utils.subprocess import run_simulation_app_function |
| 13 | + |
| 14 | +NUM_STEPS = 10 |
| 15 | +WARMUP_STEPS = 50 |
| 16 | +HEADLESS = True |
| 17 | +ENABLE_CAMERAS = True |
| 18 | + |
| 19 | + |
| 20 | +def get_test_environment(num_envs: int): |
| 21 | + """Build the G1 AGILE tabletop apple-to-plate environment for testing. |
| 22 | +
|
| 23 | + Uses a simplified scene layout (plain table, no spatial relations) to |
| 24 | + isolate task termination logic from the full production environment. |
| 25 | + """ |
| 26 | + |
| 27 | + from isaaclab_arena.assets.asset_registry import AssetRegistry |
| 28 | + from isaaclab_arena.cli.isaaclab_arena_cli import get_isaaclab_arena_cli_parser |
| 29 | + from isaaclab_arena.embodiments.g1.g1 import G1WBCAgileJointEmbodiment |
| 30 | + from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder |
| 31 | + from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment |
| 32 | + from isaaclab_arena.scene.scene import Scene |
| 33 | + from isaaclab_arena.tasks.pick_and_place_task import PickAndPlaceTask |
| 34 | + from isaaclab_arena.utils.pose import Pose |
| 35 | + |
| 36 | + asset_registry = AssetRegistry() |
| 37 | + background = asset_registry.get_asset_by_name("table")() |
| 38 | + apple = asset_registry.get_asset_by_name("apple_01_objaverse_robolab")() |
| 39 | + plate = asset_registry.get_asset_by_name("clay_plates_hot3d_robolab")() |
| 40 | + |
| 41 | + apple.set_initial_pose(Pose(position_xyz=(0.15, 0.15, 0.05), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) |
| 42 | + plate.set_initial_pose(Pose(position_xyz=(0.15, -0.15, 0.02), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) |
| 43 | + |
| 44 | + embodiment = G1WBCAgileJointEmbodiment(enable_cameras=ENABLE_CAMERAS) |
| 45 | + embodiment.set_initial_pose(Pose(position_xyz=(-0.4, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.0, 1.0))) |
| 46 | + |
| 47 | + scene = Scene(assets=[background, apple, plate]) |
| 48 | + task = PickAndPlaceTask( |
| 49 | + pick_up_object=apple, |
| 50 | + destination_location=plate, |
| 51 | + background_scene=background, |
| 52 | + episode_length_s=30.0, |
| 53 | + task_description="Pick up the apple from the table and place it onto the plate.", |
| 54 | + success_proximity_max_distance=0.15, |
| 55 | + ) |
| 56 | + |
| 57 | + isaaclab_arena_environment = IsaacLabArenaEnvironment( |
| 58 | + name="test_g1_agile_tabletop_apple_to_plate", |
| 59 | + embodiment=embodiment, |
| 60 | + scene=scene, |
| 61 | + task=task, |
| 62 | + ) |
| 63 | + |
| 64 | + args_cli = get_isaaclab_arena_cli_parser().parse_args(["--num_envs", str(num_envs)]) |
| 65 | + env_builder = ArenaEnvBuilder(isaaclab_arena_environment, args_cli) |
| 66 | + env = env_builder.make_registered() |
| 67 | + env.reset() |
| 68 | + |
| 69 | + return env, apple, plate |
| 70 | + |
| 71 | + |
| 72 | +def _step_with_standing_actions(env, num_steps: int) -> list[bool]: |
| 73 | + """Step the environment with standing idle actions and return termination flags.""" |
| 74 | + terminated_list = [] |
| 75 | + for _ in range(num_steps): |
| 76 | + with torch.inference_mode(): |
| 77 | + actions = torch.zeros(env.action_space.shape, device=env.unwrapped.device) |
| 78 | + # NOTE: Set base height to 0.75m to avoid robot squatting to match 0-height command. |
| 79 | + actions[:, -4] = 0.75 |
| 80 | + _, _, terminated, _, _ = env.step(actions) |
| 81 | + terminated_list.append(terminated.item()) |
| 82 | + return terminated_list |
| 83 | + |
| 84 | + |
| 85 | +def _replace_apple_at_initial_pose(env, apple) -> None: |
| 86 | + """Teleport the apple back to its initial position with zero velocity. |
| 87 | +
|
| 88 | + During warmup the G1 AGILE WBC policy can physically knock the apple |
| 89 | + off the table. Calling this after warmup restores the apple so the |
| 90 | + real test assertions start from a known-good state. |
| 91 | + """ |
| 92 | + from isaaclab.assets import RigidObject |
| 93 | + |
| 94 | + with torch.inference_mode(): |
| 95 | + apple_object: RigidObject = env.unwrapped.scene[apple.name] |
| 96 | + init_pos = torch.tensor([[0.15, 0.15, 0.05]], device=env.unwrapped.device) |
| 97 | + init_quat = torch.tensor([[1.0, 0.0, 0.0, 0.0]], device=env.unwrapped.device) |
| 98 | + apple_object.write_root_pose_to_sim(root_pose=torch.cat([init_pos, init_quat], dim=-1)) |
| 99 | + apple_object.write_root_velocity_to_sim(root_velocity=torch.zeros((1, 6), device=env.unwrapped.device)) |
| 100 | + |
| 101 | + |
| 102 | +def _test_initial_state_not_terminated(simulation_app) -> bool: |
| 103 | + """Apple starts away from the plate -- task must not be terminated.""" |
| 104 | + |
| 105 | + env, apple, plate = get_test_environment(num_envs=1) |
| 106 | + |
| 107 | + try: |
| 108 | + # Warmup: let the G1 AGILE WBC policy stabilise the robot before |
| 109 | + # checking termination. During the first few dozen sim steps the |
| 110 | + # robot's lower-body controller settles, which can cause brief |
| 111 | + # physics transients (vibrations, contacts) that may nudge the |
| 112 | + # apple and trigger the object-dropped termination spuriously. |
| 113 | + _step_with_standing_actions(env, WARMUP_STEPS) |
| 114 | + |
| 115 | + # Re-place the apple at its initial pose after warmup. The robot's |
| 116 | + # stabilisation during warmup can knock the apple off the table, |
| 117 | + # triggering the object_dropped termination before we even start |
| 118 | + # the real assertion steps. Re-placing ensures the test validates |
| 119 | + # the steady-state behaviour, not the warmup transient. |
| 120 | + _replace_apple_at_initial_pose(env, apple) |
| 121 | + |
| 122 | + # After warmup the robot should be stable. Assert that the task |
| 123 | + # does not terminate over the next NUM_STEPS steps. |
| 124 | + terminated_list = _step_with_standing_actions(env, NUM_STEPS) |
| 125 | + for step, terminated in enumerate(terminated_list): |
| 126 | + assert not terminated, f"Task terminated unexpectedly at post-warmup step {step}/{NUM_STEPS}" |
| 127 | + except Exception as e: |
| 128 | + print(f"Error: {e}") |
| 129 | + traceback.print_exc() |
| 130 | + return False |
| 131 | + finally: |
| 132 | + env.close() |
| 133 | + |
| 134 | + return True |
| 135 | + |
| 136 | + |
| 137 | +def _test_apple_on_plate_succeeds(simulation_app) -> bool: |
| 138 | + """Teleporting the apple onto the plate should trigger success termination.""" |
| 139 | + |
| 140 | + from isaaclab.assets import RigidObject |
| 141 | + |
| 142 | + env, apple, plate = get_test_environment(num_envs=1) |
| 143 | + |
| 144 | + try: |
| 145 | + # Warmup: stabilise the robot before teleporting the apple. |
| 146 | + _step_with_standing_actions(env, WARMUP_STEPS) |
| 147 | + |
| 148 | + with torch.inference_mode(): |
| 149 | + plate_object: RigidObject = env.unwrapped.scene[plate.name] |
| 150 | + apple_object: RigidObject = env.unwrapped.scene[apple.name] |
| 151 | + |
| 152 | + plate_pos = wp.to_torch(plate_object.data.root_pos_w)[0] |
| 153 | + target_quat = torch.tensor([[1.0, 0.0, 0.0, 0.0]], device=env.unwrapped.device) |
| 154 | + |
| 155 | + # Place the apple slightly above the plate so it falls onto it |
| 156 | + apple_target_pos = plate_pos.clone().unsqueeze(0) |
| 157 | + apple_target_pos[0, 2] += 0.05 |
| 158 | + |
| 159 | + apple_object.write_root_pose_to_sim(root_pose=torch.cat([apple_target_pos, target_quat], dim=-1)) |
| 160 | + apple_object.write_root_velocity_to_sim(root_velocity=torch.zeros((1, 6), device=env.unwrapped.device)) |
| 161 | + |
| 162 | + terminated_ever = False |
| 163 | + for _ in range(NUM_STEPS * 10): |
| 164 | + actions = torch.zeros(env.action_space.shape, device=env.unwrapped.device) |
| 165 | + # Set base height command to 0.75 to keep robot standing |
| 166 | + actions[:, -4] = 0.75 |
| 167 | + _, _, terminated, _, _ = env.step(actions) |
| 168 | + if terminated.item(): |
| 169 | + terminated_ever = True |
| 170 | + break |
| 171 | + |
| 172 | + assert terminated_ever, "Task should terminate after apple is placed on plate" |
| 173 | + print("Success: apple-on-plate termination detected") |
| 174 | + |
| 175 | + except Exception as e: |
| 176 | + print(f"Error: {e}") |
| 177 | + traceback.print_exc() |
| 178 | + return False |
| 179 | + |
| 180 | + finally: |
| 181 | + env.close() |
| 182 | + |
| 183 | + return True |
| 184 | + |
| 185 | + |
| 186 | +@pytest.mark.with_cameras |
| 187 | +def test_initial_state_not_terminated(): |
| 188 | + result = run_simulation_app_function( |
| 189 | + _test_initial_state_not_terminated, |
| 190 | + headless=HEADLESS, |
| 191 | + enable_cameras=ENABLE_CAMERAS, |
| 192 | + ) |
| 193 | + assert result, f"Test {_test_initial_state_not_terminated.__name__} failed" |
| 194 | + |
| 195 | + |
| 196 | +@pytest.mark.with_cameras |
| 197 | +def test_apple_on_plate_succeeds(): |
| 198 | + result = run_simulation_app_function( |
| 199 | + _test_apple_on_plate_succeeds, |
| 200 | + headless=HEADLESS, |
| 201 | + enable_cameras=ENABLE_CAMERAS, |
| 202 | + ) |
| 203 | + assert result, f"Test {_test_apple_on_plate_succeeds.__name__} failed" |
| 204 | + |
| 205 | + |
| 206 | +if __name__ == "__main__": |
| 207 | + test_initial_state_not_terminated() |
| 208 | + test_apple_on_plate_succeeds() |
0 commit comments