diff --git a/scripts/benchmark/atomic_action/move_end_effector_benchmark.py b/scripts/benchmark/atomic_action/move_end_effector_benchmark.py index 8cd4a75c7..b1d1dea97 100644 --- a/scripts/benchmark/atomic_action/move_end_effector_benchmark.py +++ b/scripts/benchmark/atomic_action/move_end_effector_benchmark.py @@ -60,6 +60,7 @@ class PoseCase: } DEFAULT_POSE_CASES = tuple(POSE_CASES.keys()) MOVE_SAMPLE_INTERVAL = 80 +# Endpoint accuracy of the planned trajectory against the commanded pose. SUCCESS_TOLERANCE_M = 0.01 @@ -142,7 +143,8 @@ def _run_case( binding=binding, motion_policy=MotionPolicy(sample_count=MOVE_SAMPLE_INTERVAL), ), - ) + ), + atomic_engine.initial_context(control_dt=sim.sim_config.physics_dt), ) ) is_success = bool(result.plan_success.all().item()) @@ -321,7 +323,9 @@ def main() -> None: if __name__ == "__main__": - main() + from scripts.tutorials.atomic_action.tutorial_utils import run_tutorial + + run_tutorial(main) __all__ = ["add_benchmark_args", "run_all_benchmarks"] diff --git a/scripts/benchmark/atomic_action/move_held_object_benchmark.py b/scripts/benchmark/atomic_action/move_held_object_benchmark.py index 8f577ed52..c71a7ee93 100644 --- a/scripts/benchmark/atomic_action/move_held_object_benchmark.py +++ b/scripts/benchmark/atomic_action/move_held_object_benchmark.py @@ -191,7 +191,7 @@ def _prepare_held_state( make_pre_pick_eef_pose, ) - hand_open, hand_close = get_hand_open_close_qpos(robot, sim.device) + hand_open, hand_close = get_hand_open_close_qpos(robot) pickup_args = _make_pickup_args(args, object_preset, profile) atomic_engine = AtomicActionEngine( motion_generator=motion_gen, @@ -790,7 +790,9 @@ def main() -> None: if __name__ == "__main__": - main() + from scripts.tutorials.atomic_action.tutorial_utils import run_tutorial + + run_tutorial(main) __all__ = ["add_benchmark_args", "run_all_benchmarks"] diff --git a/scripts/benchmark/atomic_action/move_joints_benchmark.py b/scripts/benchmark/atomic_action/move_joints_benchmark.py index 9d53c20c5..6b0774429 100644 --- a/scripts/benchmark/atomic_action/move_joints_benchmark.py +++ b/scripts/benchmark/atomic_action/move_joints_benchmark.py @@ -151,7 +151,10 @@ def _run_case( reset_robot(robot, initial_qpos) steps = _targets_for_sequence(atomic_engine, case, sim.device) elapsed, mem_delta, peak_gpu, result = timed_call( - lambda: atomic_engine.compile(steps) + lambda: atomic_engine.compile( + steps, + atomic_engine.initial_context(control_dt=sim.sim_config.physics_dt), + ) ) is_success = bool(result.plan_success.all().item()) traj = result.trajectory.positions @@ -333,7 +336,9 @@ def main() -> None: if __name__ == "__main__": - main() + from scripts.tutorials.atomic_action.tutorial_utils import run_tutorial + + run_tutorial(main) __all__ = ["add_benchmark_args", "run_all_benchmarks"] diff --git a/scripts/benchmark/atomic_action/pickup_benchmark.py b/scripts/benchmark/atomic_action/pickup_benchmark.py index 2bd6a1acb..eaecbe45b 100644 --- a/scripts/benchmark/atomic_action/pickup_benchmark.py +++ b/scripts/benchmark/atomic_action/pickup_benchmark.py @@ -150,7 +150,7 @@ def _run_case( settle_steps=2, ) initial_obj_position = object_position_tuple(obj) - hand_open, hand_close = get_hand_open_close_qpos(robot, sim.device) + hand_open, hand_close = get_hand_open_close_qpos(robot) initialize_pre_pick_robot_pose(robot, obj, hand_open) case_args = _make_pickup_args(args, approach, object_preset, profile) approach_direction = resolve_pickup_approach_direction( @@ -595,7 +595,9 @@ def main() -> None: if __name__ == "__main__": - main() + from scripts.tutorials.atomic_action.tutorial_utils import run_tutorial + + run_tutorial(main) __all__ = ["add_benchmark_args", "run_all_benchmarks"] diff --git a/scripts/benchmark/atomic_action/place_benchmark.py b/scripts/benchmark/atomic_action/place_benchmark.py index 4ad32f1d6..37419eaa9 100644 --- a/scripts/benchmark/atomic_action/place_benchmark.py +++ b/scripts/benchmark/atomic_action/place_benchmark.py @@ -189,7 +189,7 @@ def _prepare_held_state( initialize_pre_pick_robot_pose, ) - hand_open, hand_close = get_hand_open_close_qpos(robot, sim.device) + hand_open, hand_close = get_hand_open_close_qpos(robot) initialize_pre_pick_robot_pose(robot, obj, hand_open) pickup_args = _make_pickup_args(args, object_preset, profile) atomic_engine = AtomicActionEngine( @@ -769,7 +769,9 @@ def main() -> None: if __name__ == "__main__": - main() + from scripts.tutorials.atomic_action.tutorial_utils import run_tutorial + + run_tutorial(main) __all__ = ["add_benchmark_args", "run_all_benchmarks"] diff --git a/scripts/tutorials/atomic_action/move_end_effector.py b/scripts/tutorials/atomic_action/move_end_effector.py index 6aa4bfe1c..65114dc37 100644 --- a/scripts/tutorials/atomic_action/move_end_effector.py +++ b/scripts/tutorials/atomic_action/move_end_effector.py @@ -46,6 +46,9 @@ replay_trajectory, run_tutorial, ) +from scripts.tutorials.atomic_action.tutorial_utils import ( + initialize_benchmark_simulation, +) MOVE_SAMPLE_INTERVAL = 80 POST_TRAJECTORY_STEPS = 120 @@ -128,3 +131,16 @@ def main() -> None: if __name__ == "__main__": run_tutorial(main) + + +def initialize_simulation(args) -> "SimulationManager": + """Create the tutorial simulation for interactive or benchmark runs.""" + return initialize_benchmark_simulation(args) + + +def create_robot(sim: "SimulationManager") -> "Robot": + """Add the default MoveEndEffector tutorial robot.""" + robot = add_tutorial_robot(sim, "ur5") + # DexSim binds body_data (and therefore get_qpos) only after prepare(). + sim.prepare() + return robot diff --git a/scripts/tutorials/atomic_action/move_held_object.py b/scripts/tutorials/atomic_action/move_held_object.py index 4341110e0..8b98d1a19 100644 --- a/scripts/tutorials/atomic_action/move_held_object.py +++ b/scripts/tutorials/atomic_action/move_held_object.py @@ -60,6 +60,11 @@ replay_trajectory, run_tutorial, ) +from scripts.tutorials.atomic_action.tutorial_utils import ( + compute_pick_close_end_step, + initialize_benchmark_simulation, + make_eef_pose_at, +) OBJECT_MESH_PATH = "PaperCup/paper_cup.ply" OBJECT_XY = (-0.42, -0.08) @@ -226,3 +231,21 @@ def main() -> None: if __name__ == "__main__": run_tutorial(main) + + +def initialize_simulation(args) -> "SimulationManager": + """Create the tutorial simulation for interactive or benchmark runs.""" + return initialize_benchmark_simulation(args) + + +def create_robot(sim: "SimulationManager") -> "Robot": + """Add the default MoveHeldObject tutorial robot.""" + robot = add_tutorial_robot(sim, "ur5") + # DexSim binds body_data (and therefore get_qpos) only after prepare(). + sim.prepare() + return robot + + +def make_pre_pick_eef_pose(robot: "Robot", position) -> "torch.Tensor": + """Build the canonical top-down EEF pose at ``position``.""" + return make_eef_pose_at(robot, position) diff --git a/scripts/tutorials/atomic_action/move_joints.py b/scripts/tutorials/atomic_action/move_joints.py index c1a315dc9..b0bd44336 100644 --- a/scripts/tutorials/atomic_action/move_joints.py +++ b/scripts/tutorials/atomic_action/move_joints.py @@ -45,6 +45,9 @@ replay_trajectory, run_tutorial, ) +from scripts.tutorials.atomic_action.tutorial_utils import ( + initialize_benchmark_simulation, +) MOVE_JOINTS_SAMPLE_INTERVAL = 80 POST_TRAJECTORY_STEPS = 120 @@ -144,3 +147,16 @@ def offset_from_home(offsets: tuple[float, ...]) -> torch.Tensor: if __name__ == "__main__": run_tutorial(main) + + +def initialize_simulation(args) -> "SimulationManager": + """Create the tutorial simulation for interactive or benchmark runs.""" + return initialize_benchmark_simulation(args) + + +def create_robot(sim: "SimulationManager") -> "Robot": + """Add the default MoveJoints tutorial robot.""" + robot = add_tutorial_robot(sim, "ur5") + # DexSim binds body_data (and therefore get_qpos) only after prepare(). + sim.prepare() + return robot diff --git a/scripts/tutorials/atomic_action/pickup.py b/scripts/tutorials/atomic_action/pickup.py index e35bc902a..c85bf7760 100644 --- a/scripts/tutorials/atomic_action/pickup.py +++ b/scripts/tutorials/atomic_action/pickup.py @@ -56,6 +56,10 @@ replay_trajectory, run_tutorial, ) +from scripts.tutorials.atomic_action.tutorial_utils import ( + compute_pick_close_end_step, + initialize_benchmark_simulation, +) OBJECT_SIZE = (0.05, 0.05, 0.05) OBJECT_XY = (-0.42, -0.08) @@ -205,3 +209,16 @@ def main() -> None: if __name__ == "__main__": run_tutorial(main) + + +def initialize_simulation(args) -> "SimulationManager": + """Create the tutorial simulation for interactive or benchmark runs.""" + return initialize_benchmark_simulation(args) + + +def create_robot(sim: "SimulationManager") -> "Robot": + """Add the default PickUp tutorial robot.""" + robot = add_tutorial_robot(sim, "ur5", tcp_z=0.15) + # DexSim binds body_data (and therefore get_qpos) only after prepare(). + sim.prepare() + return robot diff --git a/scripts/tutorials/atomic_action/place.py b/scripts/tutorials/atomic_action/place.py index f943bd2e1..d67f2adb7 100644 --- a/scripts/tutorials/atomic_action/place.py +++ b/scripts/tutorials/atomic_action/place.py @@ -60,6 +60,10 @@ replay_trajectory, run_tutorial, ) +from scripts.tutorials.atomic_action.tutorial_utils import ( + compute_pick_close_end_step, + initialize_benchmark_simulation, +) OBJECT_SIZE = (0.05, 0.05, 0.05) OBJECT_XY = (-0.42, -0.08) @@ -226,3 +230,16 @@ def main() -> None: if __name__ == "__main__": run_tutorial(main) + + +def initialize_simulation(args) -> "SimulationManager": + """Create the tutorial simulation for interactive or benchmark runs.""" + return initialize_benchmark_simulation(args) + + +def create_robot(sim: "SimulationManager") -> "Robot": + """Add the default Place tutorial robot.""" + robot = add_tutorial_robot(sim, "ur5", tcp_z=0.15) + # DexSim binds body_data (and therefore get_qpos) only after prepare(). + sim.prepare() + return robot diff --git a/scripts/tutorials/atomic_action/tutorial_utils.py b/scripts/tutorials/atomic_action/tutorial_utils.py index 1167510b1..5f7307d76 100644 --- a/scripts/tutorials/atomic_action/tutorial_utils.py +++ b/scripts/tutorials/atomic_action/tutorial_utils.py @@ -1600,3 +1600,43 @@ def create_tutorial_robot_cfg( "stop_auto_play_recording", "draw_axis_marker", ] + + +def initialize_benchmark_simulation(args) -> "SimulationManager": + """Create the tutorial simulation from a benchmark-style namespace. + + Benchmark argument namespaces carry only ``device``/``renderer``; fill the + remaining launcher fields with tutorial defaults so + :func:`create_tutorial_simulation` accepts them unchanged. + + Args: + args: Namespace with optional ``num_envs``/``device``/``renderer``. + + Returns: + The shared tutorial simulation. + """ + namespace = argparse.Namespace( + num_envs=getattr(args, "num_envs", 1), + device=getattr(args, "device", "cpu"), + renderer=getattr(args, "renderer", "auto"), + headless=True, + ) + return create_tutorial_simulation(namespace) + + +def compute_pick_close_end_step(compiled=None, invocation_index: int = 0) -> int: + """Trajectory step where PickUp's hand-close segment ends (lift start). + + Args: + compiled: Optional compiled engine result; when given, the exact + ``lift`` segment start of the selected invocation is returned. + invocation_index: Invocation to inspect within ``compiled``. + + Returns: + Step index separating the grasp phase from the lift phase. Without a + compiled result this uses the tutorial defaults + (``sample_count=120`` + ``hand_interp_steps=12`` + ``settle=0``). + """ + if compiled is not None: + return int(compiled.segment(invocation_index, "lift").start) + return 120 + 12