From 16c0b3dcd4b74015da965f81729132b4a7392a43 Mon Sep 17 00:00:00 2001 From: Kikutoji Date: Fri, 17 Jul 2026 19:40:57 +0900 Subject: [PATCH] Add mechanical search clutter box environment --- .gitignore | 2 + .../scenes/mechanical_search_clutter_box.usda | 3 + examples/episodes.py | 9 +- examples/run_gripper_toggle.py | 3 + examples/run_mechanical_search_scene.py | 144 ++++++++++++++++++ .../mechanical_search_clutter_box_task.py | 93 +++++++++++ 6 files changed, 253 insertions(+), 1 deletion(-) create mode 100644 assets/scenes/mechanical_search_clutter_box.usda create mode 100644 examples/run_mechanical_search_scene.py create mode 100644 robolab/tasks/benchmark/mechanical_search_clutter_box_task.py diff --git a/.gitignore b/.gitignore index f8cba3a6..5957684b 100644 --- a/.gitignore +++ b/.gitignore @@ -42,3 +42,5 @@ pyrightconfig.json .claude/scheduled_tasks.lock .claude/agents .claude/commands +# Mechanical Search scenes generated by examples/run_mechanical_search_scene.py +assets/scenes/mechanical_search_clutter_box_runtime_*.usda diff --git a/assets/scenes/mechanical_search_clutter_box.usda b/assets/scenes/mechanical_search_clutter_box.usda new file mode 100644 index 00000000..d46cc642 --- /dev/null +++ b/assets/scenes/mechanical_search_clutter_box.usda @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8f238807b30f8bf8d99868d39689fb033f895bc6d75fcac1dca50c37e6e33e85 +size 11783 diff --git a/examples/episodes.py b/examples/episodes.py index 0f866939..96f9bc68 100644 --- a/examples/episodes.py +++ b/examples/episodes.py @@ -28,7 +28,7 @@ def run_gripper_toggle_episode(env, env_cfg=None, *, save_videos=True, video_mode="all", - headless=False, num_steps=100, toggle_every=5): + headless=False, num_steps=100, toggle_every=5, settle_steps=0): """Toggle the gripper open/closed every `toggle_every` steps while holding the arm joints fixed. @@ -40,6 +40,13 @@ def run_gripper_toggle_episode(env, env_cfg=None, *, save_videos=True, video_mod robot = env.scene["robot"] obs, _ = env.reset() + if settle_steps > 0: + current_joint_pos = robot.data.joint_pos[0, :7] + gripper_action = torch.tensor([0.785398163], device=env.device) + actions = torch.cat([current_joint_pos, gripper_action]).unsqueeze(0) + for _ in tqdm(range(settle_steps), desc="Settling scene"): + obs, _, _, _, _ = env.step(actions) + instruction = getattr(env_cfg, "instruction", None) or "gripper_toggle" if isinstance(instruction, dict): instruction = instruction.get("default", "gripper_toggle") diff --git a/examples/run_gripper_toggle.py b/examples/run_gripper_toggle.py index 68a366dc..78a52801 100644 --- a/examples/run_gripper_toggle.py +++ b/examples/run_gripper_toggle.py @@ -40,6 +40,8 @@ parser.add_argument("--tag", nargs="+", default=None, help="List of tags of tasks to run on.") parser.add_argument("--num-steps", type=int, default=100, help="Number of steps per episode.") +parser.add_argument("--settle-steps", type=int, default=0, + help="Physics steps to run after reset before recording starts.") parser.add_argument("--toggle-every", type=int, default=15, help="Toggle gripper every N steps.") parser.add_argument("--video-mode", "--video_mode", type=str, default="all", choices=["all", "viewport", "sensor", "none"], @@ -91,6 +93,7 @@ def main(): video_mode=args_cli.video_mode, headless=args_cli.headless, num_steps=args_cli.num_steps, + settle_steps=args_cli.settle_steps, toggle_every=args_cli.toggle_every, ) end_episode(env) diff --git a/examples/run_mechanical_search_scene.py b/examples/run_mechanical_search_scene.py new file mode 100644 index 00000000..e413eb69 --- /dev/null +++ b/examples/run_mechanical_search_scene.py @@ -0,0 +1,144 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# isort: skip_file + +"""Generate and record a Mechanical Search clutter-box scene.""" + +import argparse +import json +import os +import re +import shutil +import sys +import traceback +from pathlib import Path + +import cv2 # noqa: F401 must be imported before isaaclab +from isaaclab.app import AppLauncher + +PACKAGE_DIR = Path(__file__).resolve().parents[1] + + +TARGET_POSES = { + "bottom": (0.545, -0.010, 0.070), + "middle": (0.546, -0.025, 0.116), + "top": (0.545, -0.080, 0.165), +} + + +def _write_runtime_scene(target: str, target_placement: str) -> Path: + if target != "banana": + raise ValueError("Only --target banana is supported by the current minimal scene.") + + source_scene = Path(PACKAGE_DIR) / "assets" / "scenes" / "mechanical_search_clutter_box.usda" + scene_dir = source_scene.parent + runtime_scene = scene_dir / f"mechanical_search_clutter_box_runtime_{target_placement}.usda" + text = source_scene.read_text() + x, y, z = TARGET_POSES[target_placement] + + pattern = r'(def "banana" \([\s\S]*?double3 xformOp:translate = \()[^)]+(\)[\s\S]*?uniform token\[\] xformOpOrder)' + replacement = rf"\g<1>{x:.3f}, {y:.3f}, {z:.3f}\g<2>" + updated, count = re.subn(pattern, replacement, text, count=1) + if count != 1: + raise RuntimeError(f"Could not update banana pose in {source_scene}") + + runtime_scene.write_text(updated) + return runtime_scene + + +parser = argparse.ArgumentParser(description="Record a Mechanical Search clutter-box scene.") +parser.add_argument("--num_envs", type=int, default=1, help="Number of environments to spawn.") +AppLauncher.add_app_launcher_args(parser) +parser.add_argument("--task", nargs="+", default=["MechanicalSearchClutterBoxTask"], + help="Task to run. Defaults to MechanicalSearchClutterBoxTask.") +parser.add_argument("--target", default="banana", choices=["banana"], help="Target object.") +parser.add_argument("--target-placement", default="middle", choices=sorted(TARGET_POSES), + help="Initial target layer in the clutter box.") +parser.add_argument("--num-steps", type=int, default=100, help="Recorded episode length.") +parser.add_argument("--settle-steps", type=int, default=80, + help="Physics steps to run after reset before recording starts.") +parser.add_argument("--toggle-every", type=int, default=15, help="Toggle gripper every N recorded steps.") +parser.add_argument("--video-mode", "--video_mode", type=str, default="all", + choices=["all", "viewport", "sensor", "none"], + help="Which videos to save.") + +args_cli, _ = parser.parse_known_args() +args_cli.enable_cameras = True +args_cli.save_videos = args_cli.video_mode != "none" + +runtime_scene = _write_runtime_scene(args_cli.target, args_cli.target_placement) +os.environ["ROBOLAB_MECHANICAL_SEARCH_SCENE"] = str(runtime_scene) +os.environ["ROBOLAB_MECHANICAL_SEARCH_TARGET"] = args_cli.target +os.environ["ROBOLAB_MECHANICAL_SEARCH_TARGET_PLACEMENT"] = args_cli.target_placement + +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +from episodes import run_gripper_toggle_episode # noqa: E402 +from robolab.constants import set_output_dir # noqa: E402 +from robolab.core.environments.factory import get_envs # noqa: E402 +from robolab.core.environments.runtime import create_env, end_episode # noqa: E402 +from robolab.registrations.droid.auto_env_registrations_jointpos import auto_register_droid_envs # noqa: E402 + +auto_register_droid_envs(task=args_cli.task) + + +def main(): + output_root = Path(PACKAGE_DIR) / "output" / "mechanical_search_visibility" + task_envs = get_envs(task=args_cli.task) + print(f"Running {len(task_envs)} environment(s): {task_envs}") + print(f"Runtime scene: {runtime_scene}") + + for task_env in task_envs: + scene_output_dir = output_root / task_env / args_cli.target_placement + scene_output_dir.mkdir(parents=True, exist_ok=True) + set_output_dir(str(scene_output_dir)) + shutil.copy2(runtime_scene, scene_output_dir / runtime_scene.name) + + env, env_cfg = create_env( + task_env, + device=args_cli.device, + num_envs=args_cli.num_envs, + use_fabric=True, + ) + try: + print(f"Running {task_env}: '{env_cfg.instruction}'") + run_gripper_toggle_episode( + env, + env_cfg, + save_videos=args_cli.save_videos, + video_mode=args_cli.video_mode, + headless=args_cli.headless, + num_steps=args_cli.num_steps, + settle_steps=args_cli.settle_steps, + toggle_every=args_cli.toggle_every, + ) + setup_path = scene_output_dir / "scene_setup.json" + setup_path.write_text(json.dumps({ + "task": "MechanicalSearchClutterBoxTask", + "task_env": task_env, + "target_object": args_cli.target, + "target_placement": args_cli.target_placement, + "runtime_scene": str(runtime_scene), + "recorded_steps": args_cli.num_steps, + "settle_steps": args_cli.settle_steps, + "box_asset": "procedural_usda_open_top_box", + "box_scale": [1.0, 1.0, 1.0], + "approx_box_size_m": [1.00, 0.50, 0.30], + "visibility_metrics": "not_computed_yet", + }, indent=2)) + end_episode(env) + finally: + env.close() + + simulation_app.close() + + +if __name__ == "__main__": + try: + main() + except Exception as e: + print(f"Terminated with error: {e}") + traceback.print_exc() + simulation_app.close() + sys.exit(1) diff --git a/robolab/tasks/benchmark/mechanical_search_clutter_box_task.py b/robolab/tasks/benchmark/mechanical_search_clutter_box_task.py new file mode 100644 index 00000000..dbb60a4e --- /dev/null +++ b/robolab/tasks/benchmark/mechanical_search_clutter_box_task.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os +from dataclasses import dataclass +from functools import partial + +import isaaclab.envs.mdp as mdp +from isaaclab.managers import TerminationTermCfg as DoneTerm +from isaaclab.utils import configclass + +from robolab.core.scenes.utils import import_scene +from robolab.core.task.conditionals import object_dropped, object_grabbed, object_outside_of +from robolab.core.task.subtask import Subtask +from robolab.core.task.task import Task + + +TARGET_OBJECT = os.environ.get("ROBOLAB_MECHANICAL_SEARCH_TARGET", "banana") +TARGET_PLACEMENT = os.environ.get("ROBOLAB_MECHANICAL_SEARCH_TARGET_PLACEMENT", "middle") +SCENE_FILE = os.environ.get("ROBOLAB_MECHANICAL_SEARCH_SCENE", "mechanical_search_clutter_box.usda") +CONTAINER_OBJECT = "search_box" + + +@configclass +class MechanicalSearchClutterBoxTerminations: + time_out = DoneTerm(func=mdp.time_out, time_out=True) + success = DoneTerm( + func=object_outside_of, + params={ + "object": TARGET_OBJECT, + "container": CONTAINER_OBJECT, + "tolerance": 0.0, + "require_gripper_detached": True, + }, + ) + + +@dataclass +class MechanicalSearchClutterBoxTask(Task): + """Task: Find and retrieve a target object from dense clutter in a box.""" + + contact_object_list = [ + "table", + CONTAINER_OBJECT, + TARGET_OBJECT, + "rubiks_cube", + "red_block", + "green_block", + "blue_block", + "yellow_block", + "tuna_can", + "jello", + "brick", + "dry_erase_marker", + "spam_can", + "tomato_soup_can", + "clamp", + "scissors", + "sugar_box", + "mustard", + ] + scene = import_scene(SCENE_FILE, contact_object_list) + terminations = MechanicalSearchClutterBoxTerminations + instruction = { + "default": f"Find the {TARGET_OBJECT} in the cluttered box", + "vague": "Find the target object in the cluttered box", + "specific": ( + f"Search through the densely packed objects in the box, grasp the {TARGET_PLACEMENT}-layer " + f"{TARGET_OBJECT}, and pick it from the cluttered box" + ), + } + episode_length_s: int = 180 + attributes = ["semantics", "spatial"] + subtasks = [ + Subtask( + name="target_out_of_clutter_box", + conditions={ + TARGET_OBJECT: [ + partial(object_grabbed, object=TARGET_OBJECT), + partial( + object_outside_of, + object=TARGET_OBJECT, + container=CONTAINER_OBJECT, + tolerance=0.0, + require_gripper_detached=True, + ), + partial(object_dropped, object=TARGET_OBJECT), + ] + }, + logical="all", + score=1.0, + ) + ]