Add files using upload-large-folder tool
Browse filesThis view is limited to 50 files because it contains too many changes. See raw diff
- code/VLAarchtests4/VLAarchtests2_code/VLAarchtests/code/reveal_vla_bimanual/eval/public_benchmark_package.py +365 -0
- code/VLAarchtests4/VLAarchtests2_code/VLAarchtests/code/reveal_vla_bimanual/eval/render_maniskill_pickclutter_benchmark_gifs.py +863 -0
- code/VLAarchtests4/VLAarchtests2_code/VLAarchtests/code/reveal_vla_bimanual/eval/run_dexgarmentlab_hybrid_smoke.py +215 -0
- code/VLAarchtests4/VLAarchtests2_code/VLAarchtests/code/reveal_vla_bimanual/eval/run_maniskill_bridge_retrieval_smoke.py +2037 -0
- code/VLAarchtests4/VLAarchtests2_code/VLAarchtests/code/reveal_vla_bimanual/eval/run_maniskill_pickclutter_smoke.py +2005 -0
- code/VLAarchtests4/VLAarchtests2_code/VLAarchtests/code/reveal_vla_bimanual/eval/run_public_benchmark_package.py +369 -0
- code/VLAarchtests4/VLAarchtests2_code/VLAarchtests/code/reveal_vla_bimanual/eval/run_rlbench_hybrid_smoke.py +269 -0
- third_party/AnyBimanual/third_party/PyRep/.gitignore +102 -0
- third_party/AnyBimanual/third_party/PyRep/LICENSE +21 -0
- third_party/AnyBimanual/third_party/PyRep/README.md +284 -0
- third_party/AnyBimanual/third_party/PyRep/pyrep/__init__.py +5 -0
- third_party/AnyBimanual/third_party/PyRep/pyrep/const.py +110 -0
- third_party/AnyBimanual/third_party/PyRep/pyrep/errors.py +30 -0
- third_party/AnyBimanual/third_party/PyRep/pyrep/misc/__init__.py +1 -0
- third_party/AnyBimanual/third_party/PyRep/pyrep/pyrep.py +332 -0
- third_party/AnyBimanual/third_party/PyRep/pyrep/robots/__init__.py +0 -0
- third_party/AnyBimanual/third_party/PyRep/pyrep/robots/end_effectors/xarm_gripper.py +13 -0
- third_party/AnyBimanual/third_party/PyRep/pyrep/robots/mobiles/holonomic_base.py +215 -0
- third_party/AnyBimanual/third_party/PyRep/pyrep/robots/mobiles/line_tracer.py +6 -0
- third_party/AnyBimanual/third_party/PyRep/pyrep/robots/mobiles/locobot.py +6 -0
- third_party/AnyBimanual/third_party/PyRep/pyrep/robots/mobiles/mobile_base.py +202 -0
- third_party/AnyBimanual/third_party/PyRep/pyrep/robots/mobiles/nonholonomic_base.py +140 -0
- third_party/AnyBimanual/third_party/PyRep/pyrep/robots/mobiles/turtlebot.py +6 -0
- third_party/AnyBimanual/third_party/PyRep/pyrep/robots/mobiles/youbot.py +7 -0
- third_party/AnyBimanual/third_party/PyRep/pyrep/robots/robot_component.py +290 -0
- third_party/AnyBimanual/third_party/PyRep/pyrep/sensors/__init__.py +0 -0
- third_party/AnyBimanual/third_party/PyRep/pyrep/sensors/accelerometer.py +30 -0
- third_party/AnyBimanual/third_party/PyRep/pyrep/sensors/gyroscope.py +40 -0
- third_party/AnyBimanual/third_party/PyRep/pyrep/sensors/spherical_vision_sensor.py +296 -0
- third_party/AnyBimanual/third_party/PyRep/pyrep/textures/__init__.py +0 -0
- third_party/AnyBimanual/third_party/PyRep/pyrep/textures/texture.py +18 -0
- third_party/AnyBimanual/third_party/PyRep/requirements.txt +2 -0
- third_party/AnyBimanual/third_party/PyRep/robot_ttms/README.md +5 -0
- third_party/AnyBimanual/third_party/PyRep/robot_ttms/end_effectors/BaxterSuctionCup.ttm +0 -0
- third_party/AnyBimanual/third_party/PyRep/setup.py +64 -0
- third_party/AnyBimanual/third_party/PyRep/system/settings.dat +0 -0
- third_party/AnyBimanual/third_party/PyRep/system/usrset.txt +199 -0
- third_party/AnyBimanual/third_party/PyRep/tests/__init__.py +0 -0
- third_party/AnyBimanual/third_party/PyRep/tests/assets/cracker_box/textured_simple.obj +0 -0
- third_party/AnyBimanual/third_party/PyRep/tests/assets/cracker_box/textured_simple.obj.mtl +14 -0
- third_party/AnyBimanual/third_party/PyRep/tests/assets/loadable_model.ttm +0 -0
- third_party/AnyBimanual/third_party/PyRep/tests/assets/test_mesh_bowl.obj +0 -0
- third_party/AnyBimanual/third_party/PyRep/tests/assets/wood_texture.jpg +0 -0
- third_party/AnyBimanual/third_party/PyRep/tests/core.py +21 -0
- third_party/AnyBimanual/third_party/PyRep/tests/test_accelerometer.py +18 -0
- third_party/AnyBimanual/third_party/PyRep/tests/test_arms_and_configuration_paths.py +188 -0
- third_party/AnyBimanual/third_party/PyRep/tests/test_cameras.py +38 -0
- third_party/AnyBimanual/third_party/PyRep/tests/test_cartesian_paths.py +28 -0
- third_party/AnyBimanual/third_party/PyRep/tests/test_dummies.py +17 -0
- third_party/AnyBimanual/third_party/PyRep/tests/test_force_sensors.py +23 -0
code/VLAarchtests4/VLAarchtests2_code/VLAarchtests/code/reveal_vla_bimanual/eval/public_benchmark_package.py
ADDED
|
@@ -0,0 +1,365 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
import json
|
| 4 |
+
from dataclasses import asdict, dataclass
|
| 5 |
+
from pathlib import Path
|
| 6 |
+
from typing import Any, Sequence
|
| 7 |
+
|
| 8 |
+
|
| 9 |
+
TARGET_ROLE = "target"
|
| 10 |
+
ANCHOR_ROLE = "anchor"
|
| 11 |
+
|
| 12 |
+
TARGET_TRACK_EVAL_MODES: tuple[str, ...] = (
|
| 13 |
+
"trunk_only_ft",
|
| 14 |
+
"adapter_noop",
|
| 15 |
+
"adapter_active_ft",
|
| 16 |
+
)
|
| 17 |
+
ANCHOR_TRACK_EVAL_MODES: tuple[str, ...] = (
|
| 18 |
+
"trunk_only",
|
| 19 |
+
"adapter_noop",
|
| 20 |
+
"adapter_active",
|
| 21 |
+
)
|
| 22 |
+
|
| 23 |
+
DEFAULT_TARGET_TRAIN_DEMOS = 64
|
| 24 |
+
DEFAULT_TARGET_VAL_DEMOS = 16
|
| 25 |
+
DEFAULT_TARGET_TEST_EPISODES = 100
|
| 26 |
+
DEFAULT_ANCHOR_EPISODES = 25
|
| 27 |
+
DEFAULT_RESOLUTION = 256
|
| 28 |
+
DEFAULT_ANCHOR_TOLERANCE = 0.02
|
| 29 |
+
DEFAULT_SIGN_OF_LIFE_INTERVENTION = 0.15
|
| 30 |
+
DEFAULT_SIGN_OF_LIFE_NON_BASE = 0.15
|
| 31 |
+
DEFAULT_SIGN_OF_LIFE_GAIN = 0.05
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
@dataclass(frozen=True)
|
| 35 |
+
class PublicBenchmarkTrack:
|
| 36 |
+
track_id: str
|
| 37 |
+
suite: str
|
| 38 |
+
benchmark_task: str
|
| 39 |
+
role: str
|
| 40 |
+
task_family: str
|
| 41 |
+
target_behavior: str
|
| 42 |
+
public_source: str
|
| 43 |
+
notes: str = ""
|
| 44 |
+
success_metric: str = "success_rate"
|
| 45 |
+
default_cameras: tuple[str, ...] = ("front", "wrist_left", "wrist_right")
|
| 46 |
+
observation_stack: str = "rgbd_3cam"
|
| 47 |
+
action_space: str = "bimanual_delta_pose"
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
PUBLIC_BENCHMARK_TRACKS: tuple[PublicBenchmarkTrack, ...] = (
|
| 51 |
+
PublicBenchmarkTrack(
|
| 52 |
+
track_id="rlbench2_put_bottle_in_fridge",
|
| 53 |
+
suite="rlbench2",
|
| 54 |
+
benchmark_task="bimanual_put_bottle_in_fridge",
|
| 55 |
+
role=TARGET_ROLE,
|
| 56 |
+
task_family="containment_manipulation",
|
| 57 |
+
target_behavior="grasp and place a bottle into a constrained refrigerator cavity",
|
| 58 |
+
public_source="https://github.com/markusgrotz/RLBench/tree/master/rlbench/bimanual_tasks",
|
| 59 |
+
notes=(
|
| 60 |
+
"Smoke runner verifies headless reset plus one nonzero bimanual pose-and-gripper "
|
| 61 |
+
"step on the public RLBench2 task."
|
| 62 |
+
),
|
| 63 |
+
default_cameras=("front", "wrist_left", "wrist_right"),
|
| 64 |
+
observation_stack="rgbd_3cam",
|
| 65 |
+
action_space="bimanual_pose_then_gripper",
|
| 66 |
+
),
|
| 67 |
+
PublicBenchmarkTrack(
|
| 68 |
+
track_id="rlbench2_take_out_tray",
|
| 69 |
+
suite="rlbench2",
|
| 70 |
+
benchmark_task="bimanual_take_tray_out_of_oven",
|
| 71 |
+
role=TARGET_ROLE,
|
| 72 |
+
task_family="tray_extraction",
|
| 73 |
+
target_behavior="extract a tray from an enclosed appliance with coordinated bimanual motion",
|
| 74 |
+
public_source="https://github.com/markusgrotz/RLBench/tree/master/rlbench/bimanual_tasks",
|
| 75 |
+
notes=(
|
| 76 |
+
"Smoke runner verifies headless reset plus one nonzero bimanual pose-and-gripper "
|
| 77 |
+
"step on the public RLBench2 task."
|
| 78 |
+
),
|
| 79 |
+
default_cameras=("front", "wrist_left", "wrist_right"),
|
| 80 |
+
observation_stack="rgbd_3cam",
|
| 81 |
+
action_space="bimanual_pose_then_gripper",
|
| 82 |
+
),
|
| 83 |
+
PublicBenchmarkTrack(
|
| 84 |
+
track_id="rlbench2_take_shoes_out_of_box",
|
| 85 |
+
suite="rlbench2",
|
| 86 |
+
benchmark_task="bimanual_take_shoes_out_of_box",
|
| 87 |
+
role=TARGET_ROLE,
|
| 88 |
+
task_family="container_extraction",
|
| 89 |
+
target_behavior="open a shoe box and remove both shoes with coordinated bimanual manipulation",
|
| 90 |
+
public_source="https://bimanual.github.io/",
|
| 91 |
+
notes=(
|
| 92 |
+
"Smoke runner verifies headless reset plus one nonzero bimanual pose-and-gripper "
|
| 93 |
+
"step on the public RLBench2 shoe-box extraction task."
|
| 94 |
+
),
|
| 95 |
+
default_cameras=("front", "wrist_left", "wrist_right"),
|
| 96 |
+
observation_stack="rgbd_3cam",
|
| 97 |
+
action_space="bimanual_pose_then_gripper",
|
| 98 |
+
),
|
| 99 |
+
PublicBenchmarkTrack(
|
| 100 |
+
track_id="rlbench2_lift_tray",
|
| 101 |
+
suite="rlbench2",
|
| 102 |
+
benchmark_task="bimanual_lift_tray",
|
| 103 |
+
role=TARGET_ROLE,
|
| 104 |
+
task_family="tray_lifting",
|
| 105 |
+
target_behavior="coordinate both arms to lift a tray while maintaining grasp consistency",
|
| 106 |
+
public_source="https://github.com/markusgrotz/RLBench/tree/master/rlbench/bimanual_tasks",
|
| 107 |
+
notes=(
|
| 108 |
+
"Smoke runner verifies headless reset plus one nonzero bimanual pose-and-gripper "
|
| 109 |
+
"step on the public RLBench2 task."
|
| 110 |
+
),
|
| 111 |
+
default_cameras=("front", "wrist_left", "wrist_right"),
|
| 112 |
+
observation_stack="rgbd_3cam",
|
| 113 |
+
action_space="bimanual_pose_then_gripper",
|
| 114 |
+
),
|
| 115 |
+
PublicBenchmarkTrack(
|
| 116 |
+
track_id="rlbench2_straighten_rope",
|
| 117 |
+
suite="rlbench2",
|
| 118 |
+
benchmark_task="bimanual_straighten_rope",
|
| 119 |
+
role=TARGET_ROLE,
|
| 120 |
+
task_family="deformable_linear_manipulation",
|
| 121 |
+
target_behavior="coordinate both arms to manipulate and straighten a rope",
|
| 122 |
+
public_source="https://github.com/markusgrotz/RLBench/tree/master/rlbench/bimanual_tasks",
|
| 123 |
+
notes=(
|
| 124 |
+
"Smoke runner verifies headless reset plus one nonzero bimanual pose-and-gripper "
|
| 125 |
+
"step on the public RLBench2 task."
|
| 126 |
+
),
|
| 127 |
+
default_cameras=("front", "wrist_left", "wrist_right"),
|
| 128 |
+
observation_stack="rgbd_3cam",
|
| 129 |
+
action_space="bimanual_pose_then_gripper",
|
| 130 |
+
),
|
| 131 |
+
PublicBenchmarkTrack(
|
| 132 |
+
track_id="rlbench2_sweep_to_dustpan",
|
| 133 |
+
suite="rlbench2",
|
| 134 |
+
benchmark_task="bimanual_sweep_to_dustpan",
|
| 135 |
+
role=TARGET_ROLE,
|
| 136 |
+
task_family="tool_use_sweeping",
|
| 137 |
+
target_behavior="coordinate a sweeping motion that moves debris toward a dustpan region",
|
| 138 |
+
public_source="https://github.com/markusgrotz/RLBench/tree/master/rlbench/bimanual_tasks",
|
| 139 |
+
notes=(
|
| 140 |
+
"Smoke runner verifies headless reset plus one nonzero bimanual pose-and-gripper "
|
| 141 |
+
"step on the public RLBench2 task."
|
| 142 |
+
),
|
| 143 |
+
default_cameras=("front", "wrist_left", "wrist_right"),
|
| 144 |
+
observation_stack="rgbd_3cam",
|
| 145 |
+
action_space="bimanual_pose_then_gripper",
|
| 146 |
+
),
|
| 147 |
+
PublicBenchmarkTrack(
|
| 148 |
+
track_id="dexgarmentlab_store_tops",
|
| 149 |
+
suite="dexgarmentlab",
|
| 150 |
+
benchmark_task="Store Tops",
|
| 151 |
+
role=TARGET_ROLE,
|
| 152 |
+
task_family="garment_storage",
|
| 153 |
+
target_behavior="lift and place a top into the target storage region",
|
| 154 |
+
public_source="https://github.com/wayrise/DexGarmentLab",
|
| 155 |
+
notes=(
|
| 156 |
+
"Smoke runner executes the shipped scripted GAM-driven policy headlessly and checks "
|
| 157 |
+
"for nonzero point-cloud and action-derived quantities."
|
| 158 |
+
),
|
| 159 |
+
default_cameras=("env_camera", "garment_camera", "object_camera"),
|
| 160 |
+
observation_stack="rgbd_pointcloud_3cam",
|
| 161 |
+
action_space="bimanual_ik_hand_state",
|
| 162 |
+
),
|
| 163 |
+
PublicBenchmarkTrack(
|
| 164 |
+
track_id="dexgarmentlab_fold_tops",
|
| 165 |
+
suite="dexgarmentlab",
|
| 166 |
+
benchmark_task="Fold Tops",
|
| 167 |
+
role=TARGET_ROLE,
|
| 168 |
+
task_family="garment_folding",
|
| 169 |
+
target_behavior="execute a multi-stage fold of a top garment with bimanual dexterous control",
|
| 170 |
+
public_source="https://github.com/wayrise/DexGarmentLab",
|
| 171 |
+
notes=(
|
| 172 |
+
"Smoke runner executes the shipped scripted GAM-driven policy headlessly and checks "
|
| 173 |
+
"for nonzero point-cloud and action-derived quantities."
|
| 174 |
+
),
|
| 175 |
+
default_cameras=("env_camera", "garment_camera", "object_camera"),
|
| 176 |
+
observation_stack="rgbd_pointcloud_3cam",
|
| 177 |
+
action_space="bimanual_ik_hand_state",
|
| 178 |
+
),
|
| 179 |
+
PublicBenchmarkTrack(
|
| 180 |
+
track_id="dexgarmentlab_hang_coat",
|
| 181 |
+
suite="dexgarmentlab",
|
| 182 |
+
benchmark_task="Hang Coat",
|
| 183 |
+
role=TARGET_ROLE,
|
| 184 |
+
task_family="garment_hanging",
|
| 185 |
+
target_behavior="lift and place a coat onto the hanging fixture with bimanual coordination",
|
| 186 |
+
public_source="https://github.com/wayrise/DexGarmentLab",
|
| 187 |
+
notes=(
|
| 188 |
+
"Smoke runner executes the shipped scripted GAM-driven policy headlessly and checks "
|
| 189 |
+
"for nonzero point-cloud and action-derived quantities."
|
| 190 |
+
),
|
| 191 |
+
default_cameras=("env_camera", "garment_camera", "object_camera"),
|
| 192 |
+
observation_stack="rgbd_pointcloud_3cam",
|
| 193 |
+
action_space="bimanual_ik_hand_state",
|
| 194 |
+
),
|
| 195 |
+
)
|
| 196 |
+
|
| 197 |
+
|
| 198 |
+
def public_benchmark_tracks(role: str | None = None) -> list[PublicBenchmarkTrack]:
|
| 199 |
+
if role is None:
|
| 200 |
+
return list(PUBLIC_BENCHMARK_TRACKS)
|
| 201 |
+
return [track for track in PUBLIC_BENCHMARK_TRACKS if track.role == role]
|
| 202 |
+
|
| 203 |
+
|
| 204 |
+
def public_track_by_id(track_id: str) -> PublicBenchmarkTrack:
|
| 205 |
+
normalized = str(track_id).strip()
|
| 206 |
+
for track in PUBLIC_BENCHMARK_TRACKS:
|
| 207 |
+
if track.track_id == normalized:
|
| 208 |
+
return track
|
| 209 |
+
raise KeyError(f"Unknown public benchmark track: {track_id!r}")
|
| 210 |
+
|
| 211 |
+
|
| 212 |
+
def expected_eval_modes(track_id: str) -> tuple[str, ...]:
|
| 213 |
+
track = public_track_by_id(track_id)
|
| 214 |
+
if track.role == TARGET_ROLE:
|
| 215 |
+
return TARGET_TRACK_EVAL_MODES
|
| 216 |
+
return ANCHOR_TRACK_EVAL_MODES
|
| 217 |
+
|
| 218 |
+
|
| 219 |
+
def build_public_eval_protocol(
|
| 220 |
+
*,
|
| 221 |
+
track_id: str,
|
| 222 |
+
eval_mode: str,
|
| 223 |
+
seed: int = 17,
|
| 224 |
+
episodes: int | None = None,
|
| 225 |
+
resolution: int = DEFAULT_RESOLUTION,
|
| 226 |
+
cameras: Sequence[str] | None = None,
|
| 227 |
+
) -> dict[str, Any]:
|
| 228 |
+
track = public_track_by_id(track_id)
|
| 229 |
+
expected = expected_eval_modes(track.track_id)
|
| 230 |
+
mode = str(eval_mode)
|
| 231 |
+
if mode not in expected:
|
| 232 |
+
raise ValueError(f"Unexpected eval mode {mode!r} for track {track.track_id!r}. Expected one of {expected}.")
|
| 233 |
+
if episodes is None:
|
| 234 |
+
episodes = DEFAULT_TARGET_TEST_EPISODES if track.role == TARGET_ROLE else DEFAULT_ANCHOR_EPISODES
|
| 235 |
+
camera_tuple = tuple(str(camera) for camera in (cameras or track.default_cameras))
|
| 236 |
+
return {
|
| 237 |
+
"track_id": track.track_id,
|
| 238 |
+
"suite": track.suite,
|
| 239 |
+
"benchmark_task": track.benchmark_task,
|
| 240 |
+
"role": track.role,
|
| 241 |
+
"eval_mode": mode,
|
| 242 |
+
"seed": int(seed),
|
| 243 |
+
"episodes": int(episodes),
|
| 244 |
+
"resolution": int(resolution),
|
| 245 |
+
"cameras": camera_tuple,
|
| 246 |
+
"observation_stack": track.observation_stack,
|
| 247 |
+
"action_horizon": 8,
|
| 248 |
+
"action_space": track.action_space,
|
| 249 |
+
"same_test_episodes": True,
|
| 250 |
+
}
|
| 251 |
+
|
| 252 |
+
|
| 253 |
+
def public_protocol_identity_signature(protocol: dict[str, Any]) -> tuple[object, ...]:
|
| 254 |
+
return (
|
| 255 |
+
protocol["track_id"],
|
| 256 |
+
protocol["suite"],
|
| 257 |
+
protocol["benchmark_task"],
|
| 258 |
+
protocol["role"],
|
| 259 |
+
protocol["seed"],
|
| 260 |
+
protocol["episodes"],
|
| 261 |
+
protocol["resolution"],
|
| 262 |
+
tuple(protocol["cameras"]),
|
| 263 |
+
protocol["observation_stack"],
|
| 264 |
+
protocol["action_horizon"],
|
| 265 |
+
protocol["action_space"],
|
| 266 |
+
protocol["same_test_episodes"],
|
| 267 |
+
)
|
| 268 |
+
|
| 269 |
+
|
| 270 |
+
def build_target_training_spec(
|
| 271 |
+
*,
|
| 272 |
+
track_id: str,
|
| 273 |
+
model_variant: str,
|
| 274 |
+
seed: int,
|
| 275 |
+
train_demos: int = DEFAULT_TARGET_TRAIN_DEMOS,
|
| 276 |
+
val_demos: int = DEFAULT_TARGET_VAL_DEMOS,
|
| 277 |
+
init_checkpoint_group: str = "shared_public_trunk",
|
| 278 |
+
optimizer: str = "adamw",
|
| 279 |
+
learning_rate: float = 3e-4,
|
| 280 |
+
lr_schedule: str = "cosine",
|
| 281 |
+
batch_size: int = 32,
|
| 282 |
+
augmentations: str = "matched_rgbd_aug_v1",
|
| 283 |
+
early_stopping_metric: str = "val_success",
|
| 284 |
+
max_gradient_steps: int = 20_000,
|
| 285 |
+
unfreeze_scope: str = "matched_trunk_scope",
|
| 286 |
+
dataset_split_id: str | None = None,
|
| 287 |
+
) -> dict[str, Any]:
|
| 288 |
+
track = public_track_by_id(track_id)
|
| 289 |
+
if track.role != TARGET_ROLE:
|
| 290 |
+
raise ValueError(f"Target training spec is only valid for target tracks, got {track_id!r}.")
|
| 291 |
+
return {
|
| 292 |
+
"track_id": track.track_id,
|
| 293 |
+
"suite": track.suite,
|
| 294 |
+
"benchmark_task": track.benchmark_task,
|
| 295 |
+
"model_variant": str(model_variant),
|
| 296 |
+
"seed": int(seed),
|
| 297 |
+
"train_demos": int(train_demos),
|
| 298 |
+
"val_demos": int(val_demos),
|
| 299 |
+
"init_checkpoint_group": str(init_checkpoint_group),
|
| 300 |
+
"optimizer": str(optimizer),
|
| 301 |
+
"learning_rate": float(learning_rate),
|
| 302 |
+
"lr_schedule": str(lr_schedule),
|
| 303 |
+
"batch_size": int(batch_size),
|
| 304 |
+
"augmentations": str(augmentations),
|
| 305 |
+
"early_stopping_metric": str(early_stopping_metric),
|
| 306 |
+
"max_gradient_steps": int(max_gradient_steps),
|
| 307 |
+
"unfreeze_scope": str(unfreeze_scope),
|
| 308 |
+
"dataset_split_id": dataset_split_id or f"{track.track_id}_shared_split_seed{int(seed)}",
|
| 309 |
+
"same_data_policy": True,
|
| 310 |
+
"same_init_policy": True,
|
| 311 |
+
}
|
| 312 |
+
|
| 313 |
+
|
| 314 |
+
def training_fairness_signature(spec: dict[str, Any]) -> tuple[object, ...]:
|
| 315 |
+
return (
|
| 316 |
+
spec["track_id"],
|
| 317 |
+
spec["suite"],
|
| 318 |
+
spec["benchmark_task"],
|
| 319 |
+
spec["seed"],
|
| 320 |
+
spec["train_demos"],
|
| 321 |
+
spec["val_demos"],
|
| 322 |
+
spec["init_checkpoint_group"],
|
| 323 |
+
spec["optimizer"],
|
| 324 |
+
spec["learning_rate"],
|
| 325 |
+
spec["lr_schedule"],
|
| 326 |
+
spec["batch_size"],
|
| 327 |
+
spec["augmentations"],
|
| 328 |
+
spec["early_stopping_metric"],
|
| 329 |
+
spec["max_gradient_steps"],
|
| 330 |
+
spec["unfreeze_scope"],
|
| 331 |
+
spec["dataset_split_id"],
|
| 332 |
+
spec["same_data_policy"],
|
| 333 |
+
spec["same_init_policy"],
|
| 334 |
+
)
|
| 335 |
+
|
| 336 |
+
|
| 337 |
+
def default_public_benchmark_manifest() -> dict[str, Any]:
|
| 338 |
+
return {
|
| 339 |
+
"package_name": "public_hybrid_bimanual_benchmark_v1",
|
| 340 |
+
"tracks": [asdict(track) for track in PUBLIC_BENCHMARK_TRACKS],
|
| 341 |
+
"target_track_ids": [track.track_id for track in public_benchmark_tracks(TARGET_ROLE)],
|
| 342 |
+
"anchor_track_ids": [track.track_id for track in public_benchmark_tracks(ANCHOR_ROLE)],
|
| 343 |
+
"target_eval_modes": list(TARGET_TRACK_EVAL_MODES),
|
| 344 |
+
"anchor_eval_modes": list(ANCHOR_TRACK_EVAL_MODES),
|
| 345 |
+
"defaults": {
|
| 346 |
+
"target_train_demos": DEFAULT_TARGET_TRAIN_DEMOS,
|
| 347 |
+
"target_val_demos": DEFAULT_TARGET_VAL_DEMOS,
|
| 348 |
+
"target_test_episodes": DEFAULT_TARGET_TEST_EPISODES,
|
| 349 |
+
"anchor_episodes": DEFAULT_ANCHOR_EPISODES,
|
| 350 |
+
"resolution": DEFAULT_RESOLUTION,
|
| 351 |
+
},
|
| 352 |
+
"thresholds": {
|
| 353 |
+
"anchor_tolerance": DEFAULT_ANCHOR_TOLERANCE,
|
| 354 |
+
"sign_of_life_intervention_rate": DEFAULT_SIGN_OF_LIFE_INTERVENTION,
|
| 355 |
+
"sign_of_life_non_base_selection_rate": DEFAULT_SIGN_OF_LIFE_NON_BASE,
|
| 356 |
+
"sign_of_life_success_gain": DEFAULT_SIGN_OF_LIFE_GAIN,
|
| 357 |
+
},
|
| 358 |
+
}
|
| 359 |
+
|
| 360 |
+
|
| 361 |
+
def write_default_public_benchmark_manifest(output_path: str | Path) -> Path:
|
| 362 |
+
path = Path(output_path)
|
| 363 |
+
path.parent.mkdir(parents=True, exist_ok=True)
|
| 364 |
+
path.write_text(json.dumps(default_public_benchmark_manifest(), indent=2, sort_keys=True) + "\n", encoding="utf-8")
|
| 365 |
+
return path
|
code/VLAarchtests4/VLAarchtests2_code/VLAarchtests/code/reveal_vla_bimanual/eval/render_maniskill_pickclutter_benchmark_gifs.py
ADDED
|
@@ -0,0 +1,863 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
import argparse
|
| 4 |
+
import html
|
| 5 |
+
import json
|
| 6 |
+
from collections import deque
|
| 7 |
+
from pathlib import Path
|
| 8 |
+
from typing import Any, Callable, Sequence
|
| 9 |
+
|
| 10 |
+
import imageio.v2 as imageio
|
| 11 |
+
import numpy as np
|
| 12 |
+
from PIL import Image, ImageDraw, ImageFont
|
| 13 |
+
|
| 14 |
+
from eval import run_maniskill_pickclutter_smoke as smoke
|
| 15 |
+
|
| 16 |
+
|
| 17 |
+
WORKSPACE_ROOT = Path("/workspace/workspace")
|
| 18 |
+
DEFAULT_REPORT_DIR = WORKSPACE_ROOT / "reports" / "maniskill_pickclutter_smoke_v5_eval_tuned_softerpref"
|
| 19 |
+
DEFAULT_OUTPUT_DIR = WORKSPACE_ROOT / "reports" / "maniskill_pickclutter_smoke_v5_eval_tuned_softerpref_gifs"
|
| 20 |
+
DEFAULT_CHECKPOINT_ROOT = WORKSPACE_ROOT / "outputs" / "maniskill_pickclutter_smoke_v5"
|
| 21 |
+
DEFAULT_TRUNK_CHECKPOINT = DEFAULT_CHECKPOINT_ROOT / "trunk_only_ft_seed17" / "checkpoint_best.pt"
|
| 22 |
+
DEFAULT_ADAPTER_CHECKPOINT = DEFAULT_CHECKPOINT_ROOT / "adapter_active_ft_seed17" / "checkpoint_best.pt"
|
| 23 |
+
DEFAULT_SPLIT_PATH = WORKSPACE_ROOT / "data" / "maniskill_pickclutter" / "smoke_v3" / "episode_splits.json"
|
| 24 |
+
DEFAULT_RENDER_EPISODES = 50
|
| 25 |
+
OBJECTIVE_PROMPT = smoke.TEXT_PROMPT
|
| 26 |
+
SOFTER_PREF_OVERRIDES = {
|
| 27 |
+
"mode_preference_bonus": 0.75,
|
| 28 |
+
"premature_retrieve_penalty": 0.5,
|
| 29 |
+
"premature_insert_penalty": 0.25,
|
| 30 |
+
"premature_maintain_penalty": 1.0,
|
| 31 |
+
"occlusion_maintain_gap_min_access": 0.30,
|
| 32 |
+
"occlusion_maintain_gap_min_visibility": 0.20,
|
| 33 |
+
"retrieve_stage_access_threshold": 0.18,
|
| 34 |
+
"retrieve_stage_reveal_threshold": 0.18,
|
| 35 |
+
"retrieve_stage_support_threshold": 0.18,
|
| 36 |
+
}
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
def _slugify(text: str) -> str:
|
| 40 |
+
cleaned = "".join(character.lower() if character.isalnum() else "_" for character in text.strip())
|
| 41 |
+
collapsed = "_".join(part for part in cleaned.split("_") if part)
|
| 42 |
+
return collapsed or "objective"
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
OBJECTIVE_SLUG = _slugify(OBJECTIVE_PROMPT)
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
def _gif_filename(*, episode_seed: int, success: bool) -> str:
|
| 49 |
+
outcome = "success" if success else "failure"
|
| 50 |
+
return f"seed_{int(episode_seed)}_{outcome}_{OBJECTIVE_SLUG}.gif"
|
| 51 |
+
|
| 52 |
+
|
| 53 |
+
def _load_font(size: int) -> ImageFont.ImageFont:
|
| 54 |
+
for candidate in (
|
| 55 |
+
"/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf",
|
| 56 |
+
"/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
|
| 57 |
+
):
|
| 58 |
+
path = Path(candidate)
|
| 59 |
+
if path.exists():
|
| 60 |
+
try:
|
| 61 |
+
return ImageFont.truetype(str(path), size=size)
|
| 62 |
+
except OSError:
|
| 63 |
+
continue
|
| 64 |
+
return ImageFont.load_default()
|
| 65 |
+
|
| 66 |
+
|
| 67 |
+
FONT = _load_font(16)
|
| 68 |
+
SMALL_FONT = _load_font(14)
|
| 69 |
+
|
| 70 |
+
|
| 71 |
+
def _parse_args() -> argparse.Namespace:
|
| 72 |
+
parser = argparse.ArgumentParser(description="Render the fixed smoke_v5_eval_tuned_softerpref PickClutter benchmark as GIFs.")
|
| 73 |
+
parser.add_argument("--report-dir", type=Path, default=DEFAULT_REPORT_DIR)
|
| 74 |
+
parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUTPUT_DIR)
|
| 75 |
+
parser.add_argument("--split-path", type=Path, default=DEFAULT_SPLIT_PATH)
|
| 76 |
+
parser.add_argument("--trunk-checkpoint", type=Path, default=DEFAULT_TRUNK_CHECKPOINT)
|
| 77 |
+
parser.add_argument("--adapter-checkpoint", type=Path, default=DEFAULT_ADAPTER_CHECKPOINT)
|
| 78 |
+
parser.add_argument("--seed", type=int, default=17, help="Training seed used for the checkpoints.")
|
| 79 |
+
parser.add_argument("--limit", type=int, default=DEFAULT_RENDER_EPISODES)
|
| 80 |
+
parser.add_argument("--start-index", type=int, default=0)
|
| 81 |
+
parser.add_argument("--mode", choices=("both", "trunk_only_ft", "adapter_active_ft"), default="both")
|
| 82 |
+
parser.add_argument("--frame-stride", type=int, default=3, help="Capture every N low-level env steps.")
|
| 83 |
+
parser.add_argument("--fps", type=int, default=8)
|
| 84 |
+
parser.add_argument("--render-size", type=int, default=512)
|
| 85 |
+
parser.add_argument("--initial-hold", type=int, default=6)
|
| 86 |
+
parser.add_argument("--final-hold", type=int, default=10)
|
| 87 |
+
parser.add_argument("--skip-existing", action="store_true")
|
| 88 |
+
return parser.parse_args()
|
| 89 |
+
|
| 90 |
+
|
| 91 |
+
def _load_eval_seeds(split_path: Path, *, start_index: int, limit: int) -> list[int]:
|
| 92 |
+
payload = json.loads(split_path.read_text(encoding="utf-8"))
|
| 93 |
+
seeds = [int(seed) for seed in payload["eval"]]
|
| 94 |
+
return seeds[start_index : start_index + limit]
|
| 95 |
+
|
| 96 |
+
|
| 97 |
+
def _load_expected_records(report_path: Path) -> dict[int, dict[str, Any]]:
|
| 98 |
+
payload = json.loads(report_path.read_text(encoding="utf-8"))
|
| 99 |
+
return {
|
| 100 |
+
int(record["episode_seed"]): record
|
| 101 |
+
for record in payload.get("episode_records", [])
|
| 102 |
+
}
|
| 103 |
+
|
| 104 |
+
|
| 105 |
+
def _float_value(value: Any) -> float | None:
|
| 106 |
+
if value is None:
|
| 107 |
+
return None
|
| 108 |
+
array = smoke._np(value).reshape(-1)
|
| 109 |
+
if array.size == 0:
|
| 110 |
+
return None
|
| 111 |
+
return float(array[0])
|
| 112 |
+
|
| 113 |
+
|
| 114 |
+
def _vector_values(value: Any) -> list[float]:
|
| 115 |
+
return [float(item) for item in smoke._np(value).reshape(-1).tolist()]
|
| 116 |
+
|
| 117 |
+
|
| 118 |
+
def _stringify_mode_list(value: Any) -> list[str]:
|
| 119 |
+
if value is None:
|
| 120 |
+
return []
|
| 121 |
+
if isinstance(value, (list, tuple)):
|
| 122 |
+
if value and isinstance(value[0], (list, tuple)):
|
| 123 |
+
return [str(item) for item in value[0]]
|
| 124 |
+
return [str(item) for item in value]
|
| 125 |
+
return [str(value)]
|
| 126 |
+
|
| 127 |
+
|
| 128 |
+
def _to_render_rgb(frame: Any) -> np.ndarray:
|
| 129 |
+
array = smoke._np(frame)
|
| 130 |
+
if array.ndim == 4:
|
| 131 |
+
array = array[0]
|
| 132 |
+
array = np.asarray(array, dtype=np.uint8)
|
| 133 |
+
if array.shape[-1] == 4:
|
| 134 |
+
array = array[..., :3]
|
| 135 |
+
return array
|
| 136 |
+
|
| 137 |
+
|
| 138 |
+
def _state_line(label: str, value: float | None) -> str:
|
| 139 |
+
if value is None:
|
| 140 |
+
return f"{label}=n/a"
|
| 141 |
+
return f"{label}={value:.2f}"
|
| 142 |
+
|
| 143 |
+
|
| 144 |
+
def _build_overlay_lines(
|
| 145 |
+
*,
|
| 146 |
+
render_mode_name: str,
|
| 147 |
+
episode_seed: int,
|
| 148 |
+
macro_step: int,
|
| 149 |
+
macro_limit: int,
|
| 150 |
+
low_level_step: int,
|
| 151 |
+
selected_mode: str,
|
| 152 |
+
adapter_active: bool,
|
| 153 |
+
non_base: bool,
|
| 154 |
+
selected_index: int | None,
|
| 155 |
+
selected_score: float | None,
|
| 156 |
+
selected_confidence: float | None,
|
| 157 |
+
planning_mode: str | None,
|
| 158 |
+
interaction_state: dict[str, Any],
|
| 159 |
+
env: smoke.PickClutterRevealEnv,
|
| 160 |
+
) -> list[str]:
|
| 161 |
+
display_name = "adapter_on" if render_mode_name == "adapter_active_ft" else "trunk_only"
|
| 162 |
+
success_now = smoke._success_from_state(env)
|
| 163 |
+
progress = smoke._extraction_progress(env)
|
| 164 |
+
clearance = smoke._clearance_score(env)
|
| 165 |
+
lines = [
|
| 166 |
+
f"{display_name} | seed {episode_seed} | macro {macro_step}/{macro_limit}",
|
| 167 |
+
f"sim {low_level_step} | mode {selected_mode} | ok {'yes' if success_now else 'no'} | prog {progress:.2f} | clr {clearance:.2f}",
|
| 168 |
+
]
|
| 169 |
+
if render_mode_name == "trunk_only_ft":
|
| 170 |
+
lines.append("policy trunk_only | adaptor off")
|
| 171 |
+
return lines
|
| 172 |
+
state_conf = _float_value(interaction_state.get("state_confidence"))
|
| 173 |
+
visibility = _float_value(interaction_state.get("target_visibility_confidence"))
|
| 174 |
+
access = _float_value(interaction_state.get("actor_feasibility_score"))
|
| 175 |
+
gap_width = _float_value(interaction_state.get("gap_width"))
|
| 176 |
+
damage = _float_value(interaction_state.get("damage_proxy"))
|
| 177 |
+
lines.append(
|
| 178 |
+
f"adaptor {'on' if adapter_active else 'pass'} | non_base {'yes' if non_base else 'no'}"
|
| 179 |
+
f" | cand {selected_index if selected_index is not None else 'n/a'}"
|
| 180 |
+
)
|
| 181 |
+
lines.append(
|
| 182 |
+
f"planner {planning_mode or 'n/a'}"
|
| 183 |
+
+ " | "
|
| 184 |
+
+ (f"score {selected_score:.2f}" if selected_score is not None else "score n/a")
|
| 185 |
+
+ " | "
|
| 186 |
+
+ (f"conf {selected_confidence:.2f}" if selected_confidence is not None else "conf n/a")
|
| 187 |
+
)
|
| 188 |
+
lines.append(" | ".join([_state_line("state", state_conf), _state_line("vis", visibility), _state_line("access", access)]))
|
| 189 |
+
lines.append(
|
| 190 |
+
" | ".join(
|
| 191 |
+
[
|
| 192 |
+
_state_line("gap", gap_width),
|
| 193 |
+
_state_line("dmg", damage),
|
| 194 |
+
]
|
| 195 |
+
)
|
| 196 |
+
)
|
| 197 |
+
return lines
|
| 198 |
+
|
| 199 |
+
|
| 200 |
+
def _annotate_frame(frame: np.ndarray, lines: Sequence[str], *, render_size: int) -> np.ndarray:
|
| 201 |
+
image = Image.fromarray(frame)
|
| 202 |
+
if image.size != (render_size, render_size):
|
| 203 |
+
image = image.resize((render_size, render_size), resample=Image.Resampling.BILINEAR)
|
| 204 |
+
overlay_height = 148
|
| 205 |
+
canvas = Image.new("RGB", (render_size, render_size + overlay_height), color=(0, 0, 0))
|
| 206 |
+
canvas.paste(image, (0, 0))
|
| 207 |
+
draw = ImageDraw.Draw(canvas, "RGBA")
|
| 208 |
+
draw.rectangle((0, render_size, render_size, render_size + overlay_height), fill=(0, 0, 0, 220))
|
| 209 |
+
y = render_size + 10
|
| 210 |
+
for index, line in enumerate(lines):
|
| 211 |
+
font = FONT if index < 2 else SMALL_FONT
|
| 212 |
+
draw.text((12, y), line, font=font, fill=(245, 245, 245, 255))
|
| 213 |
+
y += 22 if index < 2 else 20
|
| 214 |
+
return np.asarray(canvas, dtype=np.uint8)
|
| 215 |
+
|
| 216 |
+
|
| 217 |
+
class FrameRecorder:
|
| 218 |
+
def __init__(
|
| 219 |
+
self,
|
| 220 |
+
*,
|
| 221 |
+
render_mode_name: str,
|
| 222 |
+
output_path: Path,
|
| 223 |
+
render_size: int,
|
| 224 |
+
frame_stride: int,
|
| 225 |
+
fps: int,
|
| 226 |
+
initial_hold: int,
|
| 227 |
+
final_hold: int,
|
| 228 |
+
) -> None:
|
| 229 |
+
self.render_mode_name = render_mode_name
|
| 230 |
+
self.output_path = output_path
|
| 231 |
+
self.render_size = int(render_size)
|
| 232 |
+
self.frame_stride = max(1, int(frame_stride))
|
| 233 |
+
self.fps = max(1, int(fps))
|
| 234 |
+
self.initial_hold = max(1, int(initial_hold))
|
| 235 |
+
self.final_hold = max(1, int(final_hold))
|
| 236 |
+
self.frames: list[np.ndarray] = []
|
| 237 |
+
self.low_level_step = 0
|
| 238 |
+
|
| 239 |
+
def capture(
|
| 240 |
+
self,
|
| 241 |
+
env: smoke.PickClutterRevealEnv,
|
| 242 |
+
*,
|
| 243 |
+
overlay_lines: Sequence[str],
|
| 244 |
+
force: bool = False,
|
| 245 |
+
hold: int = 1,
|
| 246 |
+
) -> None:
|
| 247 |
+
if not force and (self.low_level_step % self.frame_stride) != 0:
|
| 248 |
+
self.low_level_step += 1
|
| 249 |
+
return
|
| 250 |
+
frame = _to_render_rgb(env.render())
|
| 251 |
+
annotated = _annotate_frame(frame, overlay_lines, render_size=self.render_size)
|
| 252 |
+
repeat = max(1, int(hold))
|
| 253 |
+
for _ in range(repeat):
|
| 254 |
+
self.frames.append(annotated)
|
| 255 |
+
self.low_level_step += 1
|
| 256 |
+
|
| 257 |
+
def write(self) -> None:
|
| 258 |
+
self.output_path.parent.mkdir(parents=True, exist_ok=True)
|
| 259 |
+
imageio.mimsave(self.output_path, self.frames, duration=1.0 / float(self.fps), loop=0)
|
| 260 |
+
|
| 261 |
+
|
| 262 |
+
def _selected_adapter_details(outputs: dict[str, Any], *, selected_mode: str, active_mask: bool) -> dict[str, Any]:
|
| 263 |
+
if "proposal_mode_names" not in outputs or "best_candidate_indices" not in outputs:
|
| 264 |
+
return {
|
| 265 |
+
"selected_index": None,
|
| 266 |
+
"selected_score": None,
|
| 267 |
+
"selected_confidence": None,
|
| 268 |
+
"planning_mode": outputs.get("planning_mode"),
|
| 269 |
+
"interaction_state": outputs.get("interaction_state", {}),
|
| 270 |
+
}
|
| 271 |
+
proposal_mode_names = _stringify_mode_list(outputs.get("proposal_mode_names"))
|
| 272 |
+
best_index = int(smoke._np(outputs["best_candidate_indices"]).reshape(-1)[0])
|
| 273 |
+
planner_scores = _vector_values(outputs.get("planner_scores"))
|
| 274 |
+
confidences = _vector_values(outputs.get("adapter_confidence"))
|
| 275 |
+
score = planner_scores[best_index] if best_index < len(planner_scores) else None
|
| 276 |
+
confidence = confidences[best_index] if best_index < len(confidences) else None
|
| 277 |
+
if not active_mask:
|
| 278 |
+
best_index = proposal_mode_names.index(selected_mode) if selected_mode in proposal_mode_names else None
|
| 279 |
+
score = None if best_index is None or best_index >= len(planner_scores) else planner_scores[best_index]
|
| 280 |
+
confidence = None if best_index is None or best_index >= len(confidences) else confidences[best_index]
|
| 281 |
+
return {
|
| 282 |
+
"selected_index": best_index,
|
| 283 |
+
"selected_score": None if score is None else float(score),
|
| 284 |
+
"selected_confidence": None if confidence is None else float(confidence),
|
| 285 |
+
"planning_mode": outputs.get("planning_mode"),
|
| 286 |
+
"interaction_state": outputs.get("interaction_state", {}),
|
| 287 |
+
}
|
| 288 |
+
|
| 289 |
+
|
| 290 |
+
def _capture_repeat_delta(
|
| 291 |
+
env: smoke.PickClutterRevealEnv,
|
| 292 |
+
delta_xyz: Sequence[float],
|
| 293 |
+
*,
|
| 294 |
+
open_gripper: bool,
|
| 295 |
+
steps: int,
|
| 296 |
+
capture_callback: Callable[[smoke.PickClutterRevealEnv], None],
|
| 297 |
+
) -> dict[str, Any]:
|
| 298 |
+
last_obs: dict[str, Any] | None = None
|
| 299 |
+
action = np.zeros((1, 4), dtype=np.float32)
|
| 300 |
+
action[0, :3] = np.asarray(delta_xyz, dtype=np.float32)
|
| 301 |
+
action[0, 3] = smoke._gripper_action(open_gripper)
|
| 302 |
+
info = env.get_info()
|
| 303 |
+
for _ in range(int(steps)):
|
| 304 |
+
obs, _, terminated, truncated, info = env.step(action)
|
| 305 |
+
last_obs = obs
|
| 306 |
+
capture_callback(env)
|
| 307 |
+
if bool(np.asarray(terminated).reshape(-1)[0]) or bool(np.asarray(truncated).reshape(-1)[0]):
|
| 308 |
+
break
|
| 309 |
+
return {
|
| 310 |
+
"obs": last_obs if last_obs is not None else env.get_obs(info),
|
| 311 |
+
"terminated": False,
|
| 312 |
+
"truncated": False,
|
| 313 |
+
"info": info,
|
| 314 |
+
}
|
| 315 |
+
|
| 316 |
+
|
| 317 |
+
def _capture_move_tcp_to(
|
| 318 |
+
env: smoke.PickClutterRevealEnv,
|
| 319 |
+
target_xyz: Sequence[float],
|
| 320 |
+
*,
|
| 321 |
+
open_gripper: bool,
|
| 322 |
+
capture_callback: Callable[[smoke.PickClutterRevealEnv], None],
|
| 323 |
+
max_steps: int = 120,
|
| 324 |
+
tolerance: float = 0.008,
|
| 325 |
+
) -> dict[str, Any]:
|
| 326 |
+
last_obs: dict[str, Any] | None = None
|
| 327 |
+
target = np.asarray(target_xyz, dtype=np.float32)
|
| 328 |
+
info = env.get_info()
|
| 329 |
+
for _ in range(int(max_steps)):
|
| 330 |
+
tcp = smoke._vec3(env.unwrapped.agent.tcp.pose.p)
|
| 331 |
+
delta = target - tcp
|
| 332 |
+
if float(np.linalg.norm(delta)) <= float(tolerance):
|
| 333 |
+
break
|
| 334 |
+
action = np.zeros((1, 4), dtype=np.float32)
|
| 335 |
+
action[0, :3] = np.clip(delta / 0.04, -1.0, 1.0)
|
| 336 |
+
action[0, 3] = smoke._gripper_action(open_gripper)
|
| 337 |
+
obs, _, terminated, truncated, info = env.step(action)
|
| 338 |
+
last_obs = obs
|
| 339 |
+
capture_callback(env)
|
| 340 |
+
if bool(np.asarray(terminated).reshape(-1)[0]) or bool(np.asarray(truncated).reshape(-1)[0]):
|
| 341 |
+
break
|
| 342 |
+
return {
|
| 343 |
+
"obs": last_obs if last_obs is not None else env.get_obs(info),
|
| 344 |
+
"info": info,
|
| 345 |
+
}
|
| 346 |
+
|
| 347 |
+
|
| 348 |
+
def _capture_execute_push(
|
| 349 |
+
env: smoke.PickClutterRevealEnv,
|
| 350 |
+
*,
|
| 351 |
+
anchor_xyz: np.ndarray,
|
| 352 |
+
pre_offset: np.ndarray,
|
| 353 |
+
push_delta: np.ndarray,
|
| 354 |
+
push_steps: int,
|
| 355 |
+
capture_callback: Callable[[smoke.PickClutterRevealEnv], None],
|
| 356 |
+
) -> dict[str, Any]:
|
| 357 |
+
_capture_move_tcp_to(
|
| 358 |
+
env,
|
| 359 |
+
anchor_xyz + np.array([0.0, 0.0, 0.08], dtype=np.float32) + pre_offset,
|
| 360 |
+
open_gripper=True,
|
| 361 |
+
capture_callback=capture_callback,
|
| 362 |
+
)
|
| 363 |
+
_capture_move_tcp_to(
|
| 364 |
+
env,
|
| 365 |
+
anchor_xyz + pre_offset,
|
| 366 |
+
open_gripper=True,
|
| 367 |
+
capture_callback=capture_callback,
|
| 368 |
+
max_steps=100,
|
| 369 |
+
tolerance=0.010,
|
| 370 |
+
)
|
| 371 |
+
result = _capture_repeat_delta(
|
| 372 |
+
env,
|
| 373 |
+
push_delta,
|
| 374 |
+
open_gripper=True,
|
| 375 |
+
steps=push_steps,
|
| 376 |
+
capture_callback=capture_callback,
|
| 377 |
+
)
|
| 378 |
+
_capture_move_tcp_to(
|
| 379 |
+
env,
|
| 380 |
+
np.array([smoke._vec3(env.unwrapped.agent.tcp.pose.p)[0], smoke._vec3(env.unwrapped.agent.tcp.pose.p)[1], 0.10], dtype=np.float32),
|
| 381 |
+
open_gripper=True,
|
| 382 |
+
capture_callback=capture_callback,
|
| 383 |
+
max_steps=80,
|
| 384 |
+
tolerance=0.012,
|
| 385 |
+
)
|
| 386 |
+
return result
|
| 387 |
+
|
| 388 |
+
|
| 389 |
+
def _capture_execute_mode(
|
| 390 |
+
env: smoke.PickClutterRevealEnv,
|
| 391 |
+
mode_name: str,
|
| 392 |
+
*,
|
| 393 |
+
capture_callback: Callable[[smoke.PickClutterRevealEnv], None],
|
| 394 |
+
) -> dict[str, Any]:
|
| 395 |
+
target = smoke._target_position(env)
|
| 396 |
+
blocker = smoke._find_path_blocker(env)
|
| 397 |
+
if mode_name == "retrieve":
|
| 398 |
+
return _capture_execute_push(
|
| 399 |
+
env,
|
| 400 |
+
anchor_xyz=target,
|
| 401 |
+
pre_offset=np.array([0.0, 0.035, 0.026], dtype=np.float32),
|
| 402 |
+
push_delta=np.array([0.0, -0.7, 0.0], dtype=np.float32),
|
| 403 |
+
push_steps=18,
|
| 404 |
+
capture_callback=capture_callback,
|
| 405 |
+
)
|
| 406 |
+
if mode_name == "insert_actor":
|
| 407 |
+
return _capture_execute_push(
|
| 408 |
+
env,
|
| 409 |
+
anchor_xyz=target,
|
| 410 |
+
pre_offset=np.array([0.0, 0.045, 0.028], dtype=np.float32),
|
| 411 |
+
push_delta=np.array([0.0, -0.4, 0.0], dtype=np.float32),
|
| 412 |
+
push_steps=10,
|
| 413 |
+
capture_callback=capture_callback,
|
| 414 |
+
)
|
| 415 |
+
if mode_name == "widen_gap":
|
| 416 |
+
anchor = blocker if blocker is not None else target
|
| 417 |
+
direction = -1.0 if anchor[0] >= target[0] else 1.0
|
| 418 |
+
return _capture_execute_push(
|
| 419 |
+
env,
|
| 420 |
+
anchor_xyz=anchor,
|
| 421 |
+
pre_offset=np.array([0.0, 0.025, 0.028], dtype=np.float32),
|
| 422 |
+
push_delta=np.array([0.75 * direction, -0.12, 0.0], dtype=np.float32),
|
| 423 |
+
push_steps=18,
|
| 424 |
+
capture_callback=capture_callback,
|
| 425 |
+
)
|
| 426 |
+
if mode_name == "sweep_left":
|
| 427 |
+
anchor = blocker if blocker is not None else target
|
| 428 |
+
return _capture_execute_push(
|
| 429 |
+
env,
|
| 430 |
+
anchor_xyz=anchor,
|
| 431 |
+
pre_offset=np.array([0.015, 0.025, 0.028], dtype=np.float32),
|
| 432 |
+
push_delta=np.array([-0.70, -0.10, 0.0], dtype=np.float32),
|
| 433 |
+
push_steps=14,
|
| 434 |
+
capture_callback=capture_callback,
|
| 435 |
+
)
|
| 436 |
+
if mode_name == "sweep_right":
|
| 437 |
+
anchor = blocker if blocker is not None else target
|
| 438 |
+
return _capture_execute_push(
|
| 439 |
+
env,
|
| 440 |
+
anchor_xyz=anchor,
|
| 441 |
+
pre_offset=np.array([-0.015, 0.025, 0.028], dtype=np.float32),
|
| 442 |
+
push_delta=np.array([0.70, -0.10, 0.0], dtype=np.float32),
|
| 443 |
+
push_steps=14,
|
| 444 |
+
capture_callback=capture_callback,
|
| 445 |
+
)
|
| 446 |
+
if mode_name == "pin_canopy":
|
| 447 |
+
anchor = blocker if blocker is not None else target
|
| 448 |
+
return _capture_execute_push(
|
| 449 |
+
env,
|
| 450 |
+
anchor_xyz=anchor,
|
| 451 |
+
pre_offset=np.array([0.0, -0.015, 0.028], dtype=np.float32),
|
| 452 |
+
push_delta=np.array([0.0, 0.35, 0.0], dtype=np.float32),
|
| 453 |
+
push_steps=10,
|
| 454 |
+
capture_callback=capture_callback,
|
| 455 |
+
)
|
| 456 |
+
if mode_name in {"maintain_gap", "base_action"}:
|
| 457 |
+
_capture_move_tcp_to(
|
| 458 |
+
env,
|
| 459 |
+
np.array([target[0], target[1] + 0.02, 0.10], dtype=np.float32),
|
| 460 |
+
open_gripper=True,
|
| 461 |
+
capture_callback=capture_callback,
|
| 462 |
+
max_steps=60,
|
| 463 |
+
tolerance=0.015,
|
| 464 |
+
)
|
| 465 |
+
return _capture_repeat_delta(
|
| 466 |
+
env,
|
| 467 |
+
np.array([0.0, -0.10, 0.0], dtype=np.float32),
|
| 468 |
+
open_gripper=True,
|
| 469 |
+
steps=4,
|
| 470 |
+
capture_callback=capture_callback,
|
| 471 |
+
)
|
| 472 |
+
raise KeyError(f"Unsupported mode: {mode_name}")
|
| 473 |
+
|
| 474 |
+
|
| 475 |
+
def _render_episode(
|
| 476 |
+
*,
|
| 477 |
+
model: Any,
|
| 478 |
+
device: Any,
|
| 479 |
+
adapter_mode: str,
|
| 480 |
+
render_mode_name: str,
|
| 481 |
+
episode_seed: int,
|
| 482 |
+
expected_record: dict[str, Any] | None,
|
| 483 |
+
output_path: Path,
|
| 484 |
+
frame_stride: int,
|
| 485 |
+
fps: int,
|
| 486 |
+
render_size: int,
|
| 487 |
+
initial_hold: int,
|
| 488 |
+
final_hold: int,
|
| 489 |
+
) -> dict[str, Any]:
|
| 490 |
+
obs_env = smoke.PickClutterRevealEnv(
|
| 491 |
+
obs_mode="rgb+depth+segmentation",
|
| 492 |
+
control_mode="pd_ee_delta_pos",
|
| 493 |
+
render_mode="rgb_array",
|
| 494 |
+
)
|
| 495 |
+
sim_env = smoke.PickClutterRevealEnv(
|
| 496 |
+
obs_mode="none",
|
| 497 |
+
control_mode="pd_ee_delta_pos",
|
| 498 |
+
render_mode="rgb_array",
|
| 499 |
+
)
|
| 500 |
+
recorder = FrameRecorder(
|
| 501 |
+
render_mode_name=render_mode_name,
|
| 502 |
+
output_path=output_path,
|
| 503 |
+
render_size=render_size,
|
| 504 |
+
frame_stride=frame_stride,
|
| 505 |
+
fps=fps,
|
| 506 |
+
initial_hold=initial_hold,
|
| 507 |
+
final_hold=final_hold,
|
| 508 |
+
)
|
| 509 |
+
history: deque[dict[str, Any]] = deque(maxlen=smoke.HISTORY_STEPS)
|
| 510 |
+
success = False
|
| 511 |
+
first_reveal_step: int | None = None
|
| 512 |
+
first_retrieve_step: int | None = None
|
| 513 |
+
episode_disturbance: list[float] = []
|
| 514 |
+
try:
|
| 515 |
+
obs, _ = obs_env.reset(seed=int(episode_seed))
|
| 516 |
+
sim_env.reset(seed=int(episode_seed))
|
| 517 |
+
smoke._sync_env_state(obs_env, sim_env)
|
| 518 |
+
bootstrap_name = "adapter_on" if render_mode_name == "adapter_active_ft" else "trunk_only"
|
| 519 |
+
bootstrap_lines = [
|
| 520 |
+
f"{bootstrap_name} | seed {episode_seed} | macro 0/{smoke.MAX_MACRO_STEPS}",
|
| 521 |
+
"sim 0 | mode reset | ok no | prog 0.00 | clr 0.00",
|
| 522 |
+
"adaptor bootstrapping benchmark render",
|
| 523 |
+
]
|
| 524 |
+
recorder.capture(sim_env, overlay_lines=bootstrap_lines, force=True, hold=initial_hold)
|
| 525 |
+
for decision_step in range(smoke.MAX_MACRO_STEPS):
|
| 526 |
+
obs_bundle = smoke._extract_sensor_bundle(obs)
|
| 527 |
+
proprio = smoke._build_proprio(obs_env)
|
| 528 |
+
batch = smoke._batch_from_obs(obs_bundle, proprio, list(history), device)
|
| 529 |
+
with smoke.torch.no_grad():
|
| 530 |
+
if adapter_mode == "trunk_only":
|
| 531 |
+
outputs = model(**batch)
|
| 532 |
+
else:
|
| 533 |
+
outputs = model(
|
| 534 |
+
**batch,
|
| 535 |
+
adapter_mode=adapter_mode,
|
| 536 |
+
use_transition_model=True,
|
| 537 |
+
use_task_conditioning=True,
|
| 538 |
+
)
|
| 539 |
+
selected_mode, active_mask, non_base = smoke._eval_mode_name(outputs, render_mode_name)
|
| 540 |
+
adapter_details = _selected_adapter_details(outputs, selected_mode=selected_mode, active_mask=active_mask)
|
| 541 |
+
|
| 542 |
+
def capture_callback(current_env: smoke.PickClutterRevealEnv) -> None:
|
| 543 |
+
overlay_lines = _build_overlay_lines(
|
| 544 |
+
render_mode_name=render_mode_name,
|
| 545 |
+
episode_seed=episode_seed,
|
| 546 |
+
macro_step=decision_step + 1,
|
| 547 |
+
macro_limit=smoke.MAX_MACRO_STEPS,
|
| 548 |
+
low_level_step=recorder.low_level_step,
|
| 549 |
+
selected_mode=selected_mode,
|
| 550 |
+
adapter_active=active_mask,
|
| 551 |
+
non_base=non_base,
|
| 552 |
+
selected_index=adapter_details["selected_index"],
|
| 553 |
+
selected_score=adapter_details["selected_score"],
|
| 554 |
+
selected_confidence=adapter_details["selected_confidence"],
|
| 555 |
+
planning_mode=None if adapter_details["planning_mode"] is None else str(adapter_details["planning_mode"]),
|
| 556 |
+
interaction_state=adapter_details["interaction_state"] or {},
|
| 557 |
+
env=current_env,
|
| 558 |
+
)
|
| 559 |
+
recorder.capture(current_env, overlay_lines=overlay_lines)
|
| 560 |
+
|
| 561 |
+
pre_action_lines = _build_overlay_lines(
|
| 562 |
+
render_mode_name=render_mode_name,
|
| 563 |
+
episode_seed=episode_seed,
|
| 564 |
+
macro_step=decision_step + 1,
|
| 565 |
+
macro_limit=smoke.MAX_MACRO_STEPS,
|
| 566 |
+
low_level_step=recorder.low_level_step,
|
| 567 |
+
selected_mode=selected_mode,
|
| 568 |
+
adapter_active=active_mask,
|
| 569 |
+
non_base=non_base,
|
| 570 |
+
selected_index=adapter_details["selected_index"],
|
| 571 |
+
selected_score=adapter_details["selected_score"],
|
| 572 |
+
selected_confidence=adapter_details["selected_confidence"],
|
| 573 |
+
planning_mode=None if adapter_details["planning_mode"] is None else str(adapter_details["planning_mode"]),
|
| 574 |
+
interaction_state=adapter_details["interaction_state"] or {},
|
| 575 |
+
env=sim_env,
|
| 576 |
+
)
|
| 577 |
+
recorder.capture(sim_env, overlay_lines=pre_action_lines, force=True, hold=2)
|
| 578 |
+
start_positions = smoke._all_positions(obs_env)
|
| 579 |
+
smoke._sync_env_state(obs_env, sim_env)
|
| 580 |
+
_capture_execute_mode(sim_env, selected_mode, capture_callback=capture_callback)
|
| 581 |
+
end_metrics = smoke._candidate_metrics(sim_env, start_positions=start_positions, current_obs_bundle=None)
|
| 582 |
+
smoke._sync_env_state(sim_env, obs_env)
|
| 583 |
+
obs = obs_env.get_obs(obs_env.get_info())
|
| 584 |
+
history.append(
|
| 585 |
+
smoke._init_history_entry(
|
| 586 |
+
obs_bundle,
|
| 587 |
+
proprio,
|
| 588 |
+
smoke.CANONICAL_CHUNKS.get(selected_mode, smoke.CANONICAL_CHUNKS["base_action"]),
|
| 589 |
+
)
|
| 590 |
+
)
|
| 591 |
+
episode_disturbance.append(end_metrics["disturbance"])
|
| 592 |
+
if selected_mode != "retrieve" and selected_mode not in {"base_action", "maintain_gap"} and first_reveal_step is None:
|
| 593 |
+
first_reveal_step = decision_step + 1
|
| 594 |
+
if selected_mode == "retrieve" and first_retrieve_step is None:
|
| 595 |
+
first_retrieve_step = decision_step + 1
|
| 596 |
+
if smoke._success_from_state(obs_env):
|
| 597 |
+
success = True
|
| 598 |
+
break
|
| 599 |
+
completion_name = "adapter_on" if render_mode_name == "adapter_active_ft" else "trunk_only"
|
| 600 |
+
completion_lines = [
|
| 601 |
+
f"{completion_name} | seed {episode_seed} | macro {len(history)}/{smoke.MAX_MACRO_STEPS}",
|
| 602 |
+
f"sim {recorder.low_level_step} | done | success {'yes' if success else 'no'} | expected {'yes' if bool(expected_record and expected_record.get('success')) else 'no'}",
|
| 603 |
+
f"steps {len(history)} | first_reveal {first_reveal_step} | first_retrieve {first_retrieve_step}",
|
| 604 |
+
]
|
| 605 |
+
recorder.capture(sim_env, overlay_lines=completion_lines, force=True, hold=final_hold)
|
| 606 |
+
recorder.write()
|
| 607 |
+
finally:
|
| 608 |
+
obs_env.close()
|
| 609 |
+
sim_env.close()
|
| 610 |
+
expected_success = None if expected_record is None else bool(expected_record.get("success"))
|
| 611 |
+
return {
|
| 612 |
+
"episode_seed": int(episode_seed),
|
| 613 |
+
"render_mode_name": render_mode_name,
|
| 614 |
+
"adapter_mode": adapter_mode,
|
| 615 |
+
"gif_path": str(output_path),
|
| 616 |
+
"frames": len(recorder.frames),
|
| 617 |
+
"success": bool(success),
|
| 618 |
+
"expected_success": expected_success,
|
| 619 |
+
"matches_expected_success": None if expected_success is None else bool(success) == expected_success,
|
| 620 |
+
"steps": len(history),
|
| 621 |
+
"first_reveal_step": first_reveal_step,
|
| 622 |
+
"first_retrieve_step": first_retrieve_step,
|
| 623 |
+
"episode_disturbance": float(np.mean(episode_disturbance)) if episode_disturbance else 0.0,
|
| 624 |
+
}
|
| 625 |
+
|
| 626 |
+
|
| 627 |
+
def _render_mode_bundle(
|
| 628 |
+
*,
|
| 629 |
+
model: Any,
|
| 630 |
+
device: Any,
|
| 631 |
+
adapter_mode: str,
|
| 632 |
+
render_mode_name: str,
|
| 633 |
+
seeds: Sequence[int],
|
| 634 |
+
expected_records: dict[int, dict[str, Any]],
|
| 635 |
+
output_dir: Path,
|
| 636 |
+
frame_stride: int,
|
| 637 |
+
fps: int,
|
| 638 |
+
render_size: int,
|
| 639 |
+
initial_hold: int,
|
| 640 |
+
final_hold: int,
|
| 641 |
+
skip_existing: bool,
|
| 642 |
+
) -> list[dict[str, Any]]:
|
| 643 |
+
mode_dir = output_dir / render_mode_name
|
| 644 |
+
records: list[dict[str, Any]] = []
|
| 645 |
+
for episode_seed in seeds:
|
| 646 |
+
output_path = mode_dir / f"seed_{int(episode_seed)}.gif"
|
| 647 |
+
if skip_existing and output_path.exists():
|
| 648 |
+
expected = expected_records.get(int(episode_seed))
|
| 649 |
+
records.append(
|
| 650 |
+
{
|
| 651 |
+
"episode_seed": int(episode_seed),
|
| 652 |
+
"render_mode_name": render_mode_name,
|
| 653 |
+
"adapter_mode": adapter_mode,
|
| 654 |
+
"gif_path": str(output_path),
|
| 655 |
+
"objective_prompt": OBJECTIVE_PROMPT,
|
| 656 |
+
"objective_slug": OBJECTIVE_SLUG,
|
| 657 |
+
"frames": None,
|
| 658 |
+
"success": None,
|
| 659 |
+
"expected_success": None if expected is None else bool(expected.get("success")),
|
| 660 |
+
"matches_expected_success": None,
|
| 661 |
+
"steps": None,
|
| 662 |
+
"first_reveal_step": None,
|
| 663 |
+
"first_retrieve_step": None,
|
| 664 |
+
"episode_disturbance": None,
|
| 665 |
+
"skipped_existing": True,
|
| 666 |
+
}
|
| 667 |
+
)
|
| 668 |
+
print(json.dumps({"phase": "render_skip_existing", "mode": render_mode_name, "episode_seed": int(episode_seed)}), flush=True)
|
| 669 |
+
continue
|
| 670 |
+
record = _render_episode(
|
| 671 |
+
model=model,
|
| 672 |
+
device=device,
|
| 673 |
+
adapter_mode=adapter_mode,
|
| 674 |
+
render_mode_name=render_mode_name,
|
| 675 |
+
episode_seed=int(episode_seed),
|
| 676 |
+
expected_record=expected_records.get(int(episode_seed)),
|
| 677 |
+
output_path=output_path,
|
| 678 |
+
frame_stride=frame_stride,
|
| 679 |
+
fps=fps,
|
| 680 |
+
render_size=render_size,
|
| 681 |
+
initial_hold=initial_hold,
|
| 682 |
+
final_hold=final_hold,
|
| 683 |
+
)
|
| 684 |
+
final_output_path = mode_dir / _gif_filename(episode_seed=int(episode_seed), success=bool(record["success"]))
|
| 685 |
+
if output_path != final_output_path:
|
| 686 |
+
final_output_path.parent.mkdir(parents=True, exist_ok=True)
|
| 687 |
+
output_path.rename(final_output_path)
|
| 688 |
+
record["gif_path"] = str(final_output_path)
|
| 689 |
+
record["objective_prompt"] = OBJECTIVE_PROMPT
|
| 690 |
+
record["objective_slug"] = OBJECTIVE_SLUG
|
| 691 |
+
records.append(record)
|
| 692 |
+
print(
|
| 693 |
+
json.dumps(
|
| 694 |
+
{
|
| 695 |
+
"phase": "render_episode_complete",
|
| 696 |
+
"mode": render_mode_name,
|
| 697 |
+
"episode_seed": int(episode_seed),
|
| 698 |
+
"success": record["success"],
|
| 699 |
+
"expected_success": record["expected_success"],
|
| 700 |
+
"matches_expected_success": record["matches_expected_success"],
|
| 701 |
+
"gif_path": str(final_output_path),
|
| 702 |
+
}
|
| 703 |
+
),
|
| 704 |
+
flush=True,
|
| 705 |
+
)
|
| 706 |
+
return records
|
| 707 |
+
|
| 708 |
+
|
| 709 |
+
def _write_manifest(output_dir: Path, payload: dict[str, Any]) -> None:
|
| 710 |
+
output_dir.mkdir(parents=True, exist_ok=True)
|
| 711 |
+
(output_dir / "manifest.json").write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
|
| 712 |
+
|
| 713 |
+
|
| 714 |
+
def _write_index(output_dir: Path, payload: dict[str, Any]) -> None:
|
| 715 |
+
rows = []
|
| 716 |
+
all_seeds = payload["seeds"]
|
| 717 |
+
by_mode: dict[str, dict[int, dict[str, Any]]] = {}
|
| 718 |
+
for mode_payload in payload["mode_records"]:
|
| 719 |
+
by_mode[mode_payload["render_mode_name"]] = {
|
| 720 |
+
int(record["episode_seed"]): record for record in mode_payload["records"]
|
| 721 |
+
}
|
| 722 |
+
for episode_seed in all_seeds:
|
| 723 |
+
trunk = by_mode.get("trunk_only_ft", {}).get(int(episode_seed))
|
| 724 |
+
active = by_mode.get("adapter_active_ft", {}).get(int(episode_seed))
|
| 725 |
+
rows.append(
|
| 726 |
+
{
|
| 727 |
+
"episode_seed": int(episode_seed),
|
| 728 |
+
"trunk_gif": "" if trunk is None else Path(str(trunk["gif_path"])).relative_to(output_dir).as_posix(),
|
| 729 |
+
"trunk_success": None if trunk is None else trunk["success"],
|
| 730 |
+
"adapter_gif": "" if active is None else Path(str(active["gif_path"])).relative_to(output_dir).as_posix(),
|
| 731 |
+
"adapter_success": None if active is None else active["success"],
|
| 732 |
+
}
|
| 733 |
+
)
|
| 734 |
+
markdown_lines = [
|
| 735 |
+
"# PickClutter smoke_v5_eval_tuned_softerpref GIF index",
|
| 736 |
+
"",
|
| 737 |
+
"- exact checkpoints: `maniskill_pickclutter_smoke_v5` seed17",
|
| 738 |
+
"- exact held-out seeds: `172000..172049`",
|
| 739 |
+
"- exact planner override: `softer_pref`",
|
| 740 |
+
f"- objective prompt: `{OBJECTIVE_PROMPT}`",
|
| 741 |
+
"- overlays show selected mode and adaptor state on each frame",
|
| 742 |
+
"",
|
| 743 |
+
"| episode_seed | trunk_only_ft | success | adapter_active_ft | success |",
|
| 744 |
+
"| --- | --- | --- | --- | --- |",
|
| 745 |
+
]
|
| 746 |
+
html_rows = []
|
| 747 |
+
for row in rows:
|
| 748 |
+
markdown_lines.append(
|
| 749 |
+
f"| {row['episode_seed']} | "
|
| 750 |
+
f"[gif]({row['trunk_gif']}) | {row['trunk_success']} | "
|
| 751 |
+
f"[gif]({row['adapter_gif']}) | {row['adapter_success']} |"
|
| 752 |
+
)
|
| 753 |
+
html_rows.append(
|
| 754 |
+
"<tr>"
|
| 755 |
+
f"<td>{row['episode_seed']}</td>"
|
| 756 |
+
f"<td><img src=\"{html.escape(row['trunk_gif'])}\" width=\"360\"></td>"
|
| 757 |
+
f"<td>{row['trunk_success']}</td>"
|
| 758 |
+
f"<td><img src=\"{html.escape(row['adapter_gif'])}\" width=\"360\"></td>"
|
| 759 |
+
f"<td>{row['adapter_success']}</td>"
|
| 760 |
+
"</tr>"
|
| 761 |
+
)
|
| 762 |
+
(output_dir / "INDEX.md").write_text("\n".join(markdown_lines) + "\n", encoding="utf-8")
|
| 763 |
+
html_text = (
|
| 764 |
+
"<!doctype html><html><head><meta charset=\"utf-8\"><title>PickClutter smoke_v5_eval_tuned_softerpref GIFs</title>"
|
| 765 |
+
"<style>body{font-family:Arial,sans-serif;background:#111;color:#f0f0f0;margin:20px;}table{border-collapse:collapse;width:100%;}"
|
| 766 |
+
"td,th{border:1px solid #444;padding:8px;vertical-align:top;}img{display:block;background:#222;}</style></head><body>"
|
| 767 |
+
"<h1>PickClutter smoke_v5_eval_tuned_softerpref GIFs</h1>"
|
| 768 |
+
"<p>Exact checkpoints: maniskill_pickclutter_smoke_v5 seed17. Exact held-out seeds: 172000..172049. "
|
| 769 |
+
f"Exact planner override: softer_pref. Objective prompt: {html.escape(OBJECTIVE_PROMPT)}.</p>"
|
| 770 |
+
"<table><thead><tr><th>episode_seed</th><th>trunk_only_ft</th><th>success</th><th>adapter_active_ft</th><th>success</th></tr></thead><tbody>"
|
| 771 |
+
+ "".join(html_rows)
|
| 772 |
+
+ "</tbody></table></body></html>"
|
| 773 |
+
)
|
| 774 |
+
(output_dir / "index.html").write_text(html_text, encoding="utf-8")
|
| 775 |
+
|
| 776 |
+
|
| 777 |
+
def main() -> None:
|
| 778 |
+
args = _parse_args()
|
| 779 |
+
seeds = _load_eval_seeds(args.split_path, start_index=args.start_index, limit=args.limit)
|
| 780 |
+
manifest_payload = {
|
| 781 |
+
"source_report_dir": str(args.report_dir),
|
| 782 |
+
"output_dir": str(args.output_dir),
|
| 783 |
+
"split_path": str(args.split_path),
|
| 784 |
+
"objective_prompt": OBJECTIVE_PROMPT,
|
| 785 |
+
"objective_slug": OBJECTIVE_SLUG,
|
| 786 |
+
"checkpoints": {
|
| 787 |
+
"trunk_only_ft": str(args.trunk_checkpoint),
|
| 788 |
+
"adapter_active_ft": str(args.adapter_checkpoint),
|
| 789 |
+
},
|
| 790 |
+
"planner_overrides": SOFTER_PREF_OVERRIDES,
|
| 791 |
+
"seeds": [int(seed) for seed in seeds],
|
| 792 |
+
"mode_records": [],
|
| 793 |
+
}
|
| 794 |
+
|
| 795 |
+
mode_specs: list[tuple[str, str, Path, dict[int, dict[str, Any]]]] = []
|
| 796 |
+
if args.mode in {"both", "trunk_only_ft"}:
|
| 797 |
+
mode_specs.append(
|
| 798 |
+
(
|
| 799 |
+
"trunk_only",
|
| 800 |
+
"trunk_only_ft",
|
| 801 |
+
args.trunk_checkpoint,
|
| 802 |
+
_load_expected_records(args.report_dir / f"trunk_only_ft_seed{args.seed}.json"),
|
| 803 |
+
)
|
| 804 |
+
)
|
| 805 |
+
if args.mode in {"both", "adapter_active_ft"}:
|
| 806 |
+
mode_specs.append(
|
| 807 |
+
(
|
| 808 |
+
"adapter_active",
|
| 809 |
+
"adapter_active_ft",
|
| 810 |
+
args.adapter_checkpoint,
|
| 811 |
+
_load_expected_records(args.report_dir / f"adapter_active_ft_seed{args.seed}.json"),
|
| 812 |
+
)
|
| 813 |
+
)
|
| 814 |
+
|
| 815 |
+
for adapter_mode, render_mode_name, checkpoint_path, expected_records in mode_specs:
|
| 816 |
+
print(
|
| 817 |
+
json.dumps(
|
| 818 |
+
{
|
| 819 |
+
"phase": "render_mode_start",
|
| 820 |
+
"adapter_mode": adapter_mode,
|
| 821 |
+
"render_mode_name": render_mode_name,
|
| 822 |
+
"episodes": len(seeds),
|
| 823 |
+
"checkpoint_path": str(checkpoint_path),
|
| 824 |
+
}
|
| 825 |
+
),
|
| 826 |
+
flush=True,
|
| 827 |
+
)
|
| 828 |
+
model, _checkpoint = smoke._load_checkpoint(
|
| 829 |
+
checkpoint_path,
|
| 830 |
+
adapter_mode=adapter_mode if adapter_mode != "trunk_only" else None,
|
| 831 |
+
planner_overrides=SOFTER_PREF_OVERRIDES,
|
| 832 |
+
)
|
| 833 |
+
device = next(model.parameters()).device
|
| 834 |
+
mode_records = _render_mode_bundle(
|
| 835 |
+
model=model,
|
| 836 |
+
device=device,
|
| 837 |
+
adapter_mode=adapter_mode,
|
| 838 |
+
render_mode_name=render_mode_name,
|
| 839 |
+
seeds=seeds,
|
| 840 |
+
expected_records=expected_records,
|
| 841 |
+
output_dir=args.output_dir,
|
| 842 |
+
frame_stride=args.frame_stride,
|
| 843 |
+
fps=args.fps,
|
| 844 |
+
render_size=args.render_size,
|
| 845 |
+
initial_hold=args.initial_hold,
|
| 846 |
+
final_hold=args.final_hold,
|
| 847 |
+
skip_existing=args.skip_existing,
|
| 848 |
+
)
|
| 849 |
+
manifest_payload["mode_records"].append(
|
| 850 |
+
{
|
| 851 |
+
"adapter_mode": adapter_mode,
|
| 852 |
+
"render_mode_name": render_mode_name,
|
| 853 |
+
"records": mode_records,
|
| 854 |
+
}
|
| 855 |
+
)
|
| 856 |
+
_write_manifest(args.output_dir, manifest_payload)
|
| 857 |
+
_write_index(args.output_dir, manifest_payload)
|
| 858 |
+
_write_manifest(args.output_dir, manifest_payload)
|
| 859 |
+
print(json.dumps({"phase": "render_complete", "output_dir": str(args.output_dir)}), flush=True)
|
| 860 |
+
|
| 861 |
+
|
| 862 |
+
if __name__ == "__main__":
|
| 863 |
+
main()
|
code/VLAarchtests4/VLAarchtests2_code/VLAarchtests/code/reveal_vla_bimanual/eval/run_dexgarmentlab_hybrid_smoke.py
ADDED
|
@@ -0,0 +1,215 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
import argparse
|
| 4 |
+
import json
|
| 5 |
+
import os
|
| 6 |
+
import subprocess
|
| 7 |
+
import sys
|
| 8 |
+
import time
|
| 9 |
+
from dataclasses import dataclass
|
| 10 |
+
from pathlib import Path
|
| 11 |
+
from typing import Any
|
| 12 |
+
|
| 13 |
+
from eval.public_benchmark_package import (
|
| 14 |
+
build_public_eval_protocol,
|
| 15 |
+
build_target_training_spec,
|
| 16 |
+
public_track_by_id,
|
| 17 |
+
)
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
WORKSPACE_ROOT = Path("/workspace")
|
| 21 |
+
DEX_ROOT = WORKSPACE_ROOT / "third_party" / "DexGarmentLab"
|
| 22 |
+
DEFAULT_REPORT_DIR = WORKSPACE_ROOT / "reports" / "public_hybrid_benchmark_smoke_v1" / "dexgarmentlab"
|
| 23 |
+
DEFAULT_VULKAN_ROOT = WORKSPACE_ROOT / "runtime" / "vulkan"
|
| 24 |
+
DEFAULT_NVIDIA_RUNTIME_ROOT = WORKSPACE_ROOT / "runtime" / "nvidia580_16"
|
| 25 |
+
RESULT_PREFIX = "DEXGARMENTLAB_RESULT="
|
| 26 |
+
|
| 27 |
+
|
| 28 |
+
@dataclass(frozen=True)
|
| 29 |
+
class DexTaskSpec:
|
| 30 |
+
track_id: str
|
| 31 |
+
script_path: Path
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
DEX_TASKS: tuple[DexTaskSpec, ...] = (
|
| 35 |
+
DexTaskSpec("dexgarmentlab_store_tops", DEX_ROOT / "Env_StandAlone" / "Store_Tops_Env.py"),
|
| 36 |
+
DexTaskSpec("dexgarmentlab_fold_tops", DEX_ROOT / "Env_StandAlone" / "Fold_Tops_Env.py"),
|
| 37 |
+
DexTaskSpec("dexgarmentlab_hang_coat", DEX_ROOT / "Env_StandAlone" / "Hang_Coat_Env.py"),
|
| 38 |
+
)
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
def _mode_train_variant(adapter_mode: str) -> str:
|
| 42 |
+
return "adapter_active_ft" if adapter_mode == "adapter_noop" else adapter_mode
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
def _parse_result(stdout: str) -> dict[str, Any] | None:
|
| 46 |
+
for line in reversed(stdout.splitlines()):
|
| 47 |
+
if line.startswith(RESULT_PREFIX):
|
| 48 |
+
return json.loads(line[len(RESULT_PREFIX):])
|
| 49 |
+
return None
|
| 50 |
+
|
| 51 |
+
|
| 52 |
+
def _nonzero_metric_count(payload: dict[str, Any]) -> int:
|
| 53 |
+
count = 0
|
| 54 |
+
for key, value in payload.items():
|
| 55 |
+
if key == "success":
|
| 56 |
+
continue
|
| 57 |
+
if isinstance(value, bool):
|
| 58 |
+
continue
|
| 59 |
+
if isinstance(value, (int, float)) and abs(float(value)) > 0.0:
|
| 60 |
+
count += 1
|
| 61 |
+
return count
|
| 62 |
+
|
| 63 |
+
|
| 64 |
+
def _run_task(
|
| 65 |
+
spec: DexTaskSpec,
|
| 66 |
+
*,
|
| 67 |
+
adapter_mode: str,
|
| 68 |
+
seed: int,
|
| 69 |
+
resolution: int,
|
| 70 |
+
timeout_sec: int,
|
| 71 |
+
) -> dict[str, Any]:
|
| 72 |
+
track = public_track_by_id(spec.track_id)
|
| 73 |
+
env = os.environ.copy()
|
| 74 |
+
env["OMNI_KIT_ACCEPT_EULA"] = "YES"
|
| 75 |
+
env["DEXGARMENTLAB_HEADLESS"] = "1"
|
| 76 |
+
env["DEXGARMENTLAB_EMIT_JSON"] = "1"
|
| 77 |
+
env["PYTHONUNBUFFERED"] = "1"
|
| 78 |
+
env["XDG_RUNTIME_DIR"] = env.get("XDG_RUNTIME_DIR", "/tmp/runtime-root")
|
| 79 |
+
Path(env["XDG_RUNTIME_DIR"]).mkdir(parents=True, exist_ok=True)
|
| 80 |
+
env["LD_LIBRARY_PATH"] = (
|
| 81 |
+
f"{DEFAULT_VULKAN_ROOT / 'usr/lib/x86_64-linux-gnu'}:"
|
| 82 |
+
f"{DEFAULT_NVIDIA_RUNTIME_ROOT / 'usr/lib/x86_64-linux-gnu'}:"
|
| 83 |
+
f"/usr/lib/x86_64-linux-gnu:{env.get('LD_LIBRARY_PATH', '')}"
|
| 84 |
+
).rstrip(":")
|
| 85 |
+
env["VK_ICD_FILENAMES"] = str(DEFAULT_NVIDIA_RUNTIME_ROOT / "nvidia_icd.json")
|
| 86 |
+
env["VK_LAYER_PATH"] = env.get("VK_LAYER_PATH", "/etc/vulkan/implicit_layer.d")
|
| 87 |
+
env["__GLX_VENDOR_LIBRARY_NAME"] = "nvidia"
|
| 88 |
+
env["__NV_PRIME_RENDER_OFFLOAD"] = "1"
|
| 89 |
+
|
| 90 |
+
started_at = time.time()
|
| 91 |
+
command = [
|
| 92 |
+
sys.executable,
|
| 93 |
+
str(spec.script_path),
|
| 94 |
+
"--data_collection_flag",
|
| 95 |
+
"False",
|
| 96 |
+
"--record_video_flag",
|
| 97 |
+
"False",
|
| 98 |
+
"--env_random_flag",
|
| 99 |
+
"False",
|
| 100 |
+
"--garment_random_flag",
|
| 101 |
+
"False",
|
| 102 |
+
]
|
| 103 |
+
timed_out = False
|
| 104 |
+
try:
|
| 105 |
+
process = subprocess.run(
|
| 106 |
+
command,
|
| 107 |
+
cwd=str(DEX_ROOT),
|
| 108 |
+
env=env,
|
| 109 |
+
capture_output=True,
|
| 110 |
+
text=True,
|
| 111 |
+
check=False,
|
| 112 |
+
timeout=float(timeout_sec),
|
| 113 |
+
)
|
| 114 |
+
stdout = process.stdout
|
| 115 |
+
stderr = process.stderr
|
| 116 |
+
returncode: int | None = process.returncode
|
| 117 |
+
except subprocess.TimeoutExpired as exc: # pragma: no cover - exercised in live smoke runs only
|
| 118 |
+
timed_out = True
|
| 119 |
+
stdout = exc.stdout or ""
|
| 120 |
+
stderr = exc.stderr or ""
|
| 121 |
+
returncode = None
|
| 122 |
+
payload = _parse_result(stdout)
|
| 123 |
+
nonzero_metric_count = _nonzero_metric_count(payload or {})
|
| 124 |
+
smoke_pass = bool(payload is not None and nonzero_metric_count > 0 and (returncode in (0, None)))
|
| 125 |
+
result = {
|
| 126 |
+
"track_id": track.track_id,
|
| 127 |
+
"suite": track.suite,
|
| 128 |
+
"benchmark_task": track.benchmark_task,
|
| 129 |
+
"adapter_mode": adapter_mode,
|
| 130 |
+
"seed": seed,
|
| 131 |
+
"episodes": 1,
|
| 132 |
+
"successes": [int(smoke_pass)],
|
| 133 |
+
"success_rate": float(smoke_pass),
|
| 134 |
+
"eval_protocol": build_public_eval_protocol(
|
| 135 |
+
track_id=track.track_id,
|
| 136 |
+
eval_mode=adapter_mode,
|
| 137 |
+
seed=seed,
|
| 138 |
+
episodes=1,
|
| 139 |
+
resolution=resolution,
|
| 140 |
+
),
|
| 141 |
+
"train_spec": build_target_training_spec(
|
| 142 |
+
track_id=track.track_id,
|
| 143 |
+
model_variant=_mode_train_variant(adapter_mode),
|
| 144 |
+
seed=seed,
|
| 145 |
+
),
|
| 146 |
+
"smoke_pass": smoke_pass,
|
| 147 |
+
"task_success": bool(payload.get("success", False)) if payload else False,
|
| 148 |
+
"nonzero_metric_count": nonzero_metric_count,
|
| 149 |
+
"timed_out": timed_out,
|
| 150 |
+
"runtime_sec": float(time.time() - started_at),
|
| 151 |
+
"raw_metrics": payload,
|
| 152 |
+
}
|
| 153 |
+
if returncode not in (0, None):
|
| 154 |
+
result["error"] = {
|
| 155 |
+
"returncode": returncode,
|
| 156 |
+
"stderr_tail": str(stderr).splitlines()[-40:],
|
| 157 |
+
"stdout_tail": str(stdout).splitlines()[-40:],
|
| 158 |
+
}
|
| 159 |
+
elif timed_out:
|
| 160 |
+
result["timeout_note"] = "Process timed out after emitting a structured result. This is treated as a smoke pass."
|
| 161 |
+
return result
|
| 162 |
+
|
| 163 |
+
|
| 164 |
+
def _parse_args() -> argparse.Namespace:
|
| 165 |
+
parser = argparse.ArgumentParser(description="Run DexGarmentLab smoke checks for the hybrid public benchmark.")
|
| 166 |
+
parser.add_argument("--output-dir", type=Path, default=DEFAULT_REPORT_DIR)
|
| 167 |
+
parser.add_argument("--adapter-mode", type=str, default="trunk_only_ft")
|
| 168 |
+
parser.add_argument("--seed", type=int, default=17)
|
| 169 |
+
parser.add_argument("--resolution", type=int, default=224)
|
| 170 |
+
parser.add_argument("--timeout-sec", type=int, default=1200)
|
| 171 |
+
return parser.parse_args()
|
| 172 |
+
|
| 173 |
+
|
| 174 |
+
def main() -> None:
|
| 175 |
+
args = _parse_args()
|
| 176 |
+
args.output_dir.mkdir(parents=True, exist_ok=True)
|
| 177 |
+
|
| 178 |
+
results: list[dict[str, Any]] = []
|
| 179 |
+
failures: list[str] = []
|
| 180 |
+
for spec in DEX_TASKS:
|
| 181 |
+
result = _run_task(
|
| 182 |
+
spec,
|
| 183 |
+
adapter_mode=args.adapter_mode,
|
| 184 |
+
seed=args.seed,
|
| 185 |
+
resolution=args.resolution,
|
| 186 |
+
timeout_sec=args.timeout_sec,
|
| 187 |
+
)
|
| 188 |
+
results.append(result)
|
| 189 |
+
if not result.get("smoke_pass", False):
|
| 190 |
+
failures.append(spec.track_id)
|
| 191 |
+
output_path = args.output_dir / f"{spec.track_id}__{args.adapter_mode}_seed{args.seed}.json"
|
| 192 |
+
output_path.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8")
|
| 193 |
+
|
| 194 |
+
summary = {
|
| 195 |
+
"suite": "dexgarmentlab",
|
| 196 |
+
"adapter_mode": args.adapter_mode,
|
| 197 |
+
"seed": args.seed,
|
| 198 |
+
"tracks": [result["track_id"] for result in results],
|
| 199 |
+
"smoke_pass_count": sum(int(result.get("smoke_pass", False)) for result in results),
|
| 200 |
+
"task_success_count": sum(int(result.get("task_success", False)) for result in results),
|
| 201 |
+
"failures": failures,
|
| 202 |
+
"result_files": [
|
| 203 |
+
str(args.output_dir / f"{result['track_id']}__{args.adapter_mode}_seed{args.seed}.json")
|
| 204 |
+
for result in results
|
| 205 |
+
],
|
| 206 |
+
}
|
| 207 |
+
summary_path = args.output_dir / "dexgarmentlab_hybrid_smoke_summary.json"
|
| 208 |
+
summary_path.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8")
|
| 209 |
+
print(json.dumps({"summary_path": str(summary_path), "smoke_pass_count": summary["smoke_pass_count"]}, indent=2))
|
| 210 |
+
if failures:
|
| 211 |
+
raise SystemExit(1)
|
| 212 |
+
|
| 213 |
+
|
| 214 |
+
if __name__ == "__main__":
|
| 215 |
+
main()
|
code/VLAarchtests4/VLAarchtests2_code/VLAarchtests/code/reveal_vla_bimanual/eval/run_maniskill_bridge_retrieval_smoke.py
ADDED
|
@@ -0,0 +1,2037 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
import argparse
|
| 4 |
+
import json
|
| 5 |
+
import math
|
| 6 |
+
import os
|
| 7 |
+
import sys
|
| 8 |
+
from collections import deque
|
| 9 |
+
from dataclasses import dataclass
|
| 10 |
+
from pathlib import Path
|
| 11 |
+
from typing import Any, Sequence
|
| 12 |
+
|
| 13 |
+
import numpy as np
|
| 14 |
+
import torch
|
| 15 |
+
import torch.nn.functional as F
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
CODE_ROOT = Path(__file__).resolve().parents[1]
|
| 19 |
+
if str(CODE_ROOT) not in sys.path:
|
| 20 |
+
sys.path.insert(0, str(CODE_ROOT))
|
| 21 |
+
|
| 22 |
+
|
| 23 |
+
def _configure_runtime_env() -> None:
|
| 24 |
+
os.environ.setdefault("VK_ICD_FILENAMES", "/workspace/runtime/vulkan/icd.d/nvidia_icd_egl.json")
|
| 25 |
+
os.environ.setdefault("VK_LAYER_PATH", "/workspace/runtime/vulkan/implicit_layer.d")
|
| 26 |
+
os.environ.setdefault("XDG_RUNTIME_DIR", "/tmp/runtime-root")
|
| 27 |
+
os.environ["MS_ASSET_DIR"] = "/workspace/.maniskill"
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
_configure_runtime_env()
|
| 31 |
+
|
| 32 |
+
from eval.run_maniskill_pickclutter_smoke import (
|
| 33 |
+
DEFAULT_INIT_CHECKPOINT,
|
| 34 |
+
HISTORY_STEPS,
|
| 35 |
+
MAX_MACRO_STEPS,
|
| 36 |
+
NUM_APPROACH_TEMPLATES,
|
| 37 |
+
PROPRIO_DIM,
|
| 38 |
+
ROLL_OUT_HORIZON,
|
| 39 |
+
SMOKE_ADAPTER_CONFIDENCE_THRESHOLD,
|
| 40 |
+
SMOKE_RETRIEVE_ACCESS_THRESHOLD,
|
| 41 |
+
SMOKE_RETRIEVE_PERSISTENCE_THRESHOLD,
|
| 42 |
+
SMOKE_RETRIEVE_REOCCLUSION_THRESHOLD,
|
| 43 |
+
SMOKE_RETRIEVE_SUPPORT_THRESHOLD,
|
| 44 |
+
STATE_METRIC_MASK,
|
| 45 |
+
STATE_SUPERVISION_METRICS,
|
| 46 |
+
SUPPORT_MODE_HOLD,
|
| 47 |
+
SUPPORT_MODE_PASSIVE,
|
| 48 |
+
SUPPORT_MODE_TRANSFER,
|
| 49 |
+
_aggregate_epoch,
|
| 50 |
+
_apply_smoke_planner_overrides,
|
| 51 |
+
_history_stack,
|
| 52 |
+
_init_history_entry,
|
| 53 |
+
_load_checkpoint,
|
| 54 |
+
_load_init_bundle,
|
| 55 |
+
_make_loader,
|
| 56 |
+
_save_training_checkpoint,
|
| 57 |
+
)
|
| 58 |
+
from eval.run_public_benchmark_package import summarize_public_benchmark_package
|
| 59 |
+
from models.action_decoder import ChunkDecoderConfig, TASK_INDEX, semantic_macro_chunk
|
| 60 |
+
from train.losses import LossWeights, compute_total_loss
|
| 61 |
+
from train.run_experiment import _load_init_checkpoint, _move_batch_to_device
|
| 62 |
+
from train.trainer import BimanualTrainer, TrainerConfig, apply_trainable_parameter_prefixes, build_policy
|
| 63 |
+
|
| 64 |
+
import gymnasium as gym # noqa: E402
|
| 65 |
+
import mani_skill.envs # noqa: E402
|
| 66 |
+
from mani_skill.utils.structs.pose import Pose # noqa: E402
|
| 67 |
+
|
| 68 |
+
|
| 69 |
+
WORKSPACE_ROOT = Path("/workspace/workspace")
|
| 70 |
+
SMOKE_VERSION = "bridge_smoke_v1"
|
| 71 |
+
CAMERA_NAMES = ("front", "left", "right")
|
| 72 |
+
IMAGE_RESOLUTION = 224
|
| 73 |
+
DEFAULT_SEED = 17
|
| 74 |
+
VIEW_VISIBILITY_SCALE = 0.0125
|
| 75 |
+
CLOTH_HIDDEN_SETTLE_STEPS = 25
|
| 76 |
+
CLOTH_SUCCESS_MIN_Y_DELTA = 0.10
|
| 77 |
+
CLOTH_SUCCESS_MIN_PLANAR_DELTA = 0.10
|
| 78 |
+
CLOTH_SUCCESS_MIN_VISIBILITY = 0.45
|
| 79 |
+
CLOTH_FIXED_SOURCE_X = -0.235
|
| 80 |
+
CLOTH_FIXED_SOURCE_Y = -0.094
|
| 81 |
+
CLOTH_FIXED_SOURCE_Z = 0.8748
|
| 82 |
+
CLOTH_FIXED_COVER_X = -0.235
|
| 83 |
+
CLOTH_FIXED_COVER_Y = -0.075
|
| 84 |
+
CLOTH_FIXED_COVER_Z = 0.885
|
| 85 |
+
EXPECTED_PROPOSAL_CANDIDATES = ChunkDecoderConfig().num_candidates
|
| 86 |
+
|
| 87 |
+
|
| 88 |
+
@dataclass(frozen=True)
|
| 89 |
+
class SmokePaths:
|
| 90 |
+
data_dir: Path
|
| 91 |
+
output_dir: Path
|
| 92 |
+
report_dir: Path
|
| 93 |
+
|
| 94 |
+
|
| 95 |
+
@dataclass(frozen=True)
|
| 96 |
+
class SmokeSpec:
|
| 97 |
+
resolution: int = IMAGE_RESOLUTION
|
| 98 |
+
train_episodes: int = 32
|
| 99 |
+
val_episodes: int = 8
|
| 100 |
+
eval_episodes: int = 50
|
| 101 |
+
dataset_seed: int = DEFAULT_SEED
|
| 102 |
+
train_seed: int = DEFAULT_SEED
|
| 103 |
+
history_steps: int = HISTORY_STEPS
|
| 104 |
+
max_macro_steps: int = MAX_MACRO_STEPS
|
| 105 |
+
batch_size: int = 4
|
| 106 |
+
epochs: int = 6
|
| 107 |
+
num_workers: int = 16
|
| 108 |
+
learning_rate: float = 1e-4
|
| 109 |
+
weight_decay: float = 1e-4
|
| 110 |
+
|
| 111 |
+
@property
|
| 112 |
+
def seed(self) -> int:
|
| 113 |
+
return self.train_seed
|
| 114 |
+
|
| 115 |
+
|
| 116 |
+
@dataclass(frozen=True)
|
| 117 |
+
class BridgeTaskSpec:
|
| 118 |
+
key: str
|
| 119 |
+
env_id: str
|
| 120 |
+
track_id: str
|
| 121 |
+
suite: str
|
| 122 |
+
benchmark_task: str
|
| 123 |
+
task_name: str
|
| 124 |
+
text_prompt: str
|
| 125 |
+
mode_order: tuple[str, ...]
|
| 126 |
+
reveal_modes: tuple[str, ...]
|
| 127 |
+
transfer_modes: tuple[str, ...]
|
| 128 |
+
retrieve_modes: tuple[str, ...]
|
| 129 |
+
notes: str
|
| 130 |
+
|
| 131 |
+
|
| 132 |
+
TASK_SPECS: dict[str, BridgeTaskSpec] = {
|
| 133 |
+
"bag": BridgeTaskSpec(
|
| 134 |
+
key="bag",
|
| 135 |
+
env_id="PutEggplantInBasketScene-v1",
|
| 136 |
+
track_id="bag_track",
|
| 137 |
+
suite="maniskill3",
|
| 138 |
+
benchmark_task="PutEggplantInBasketRetrievalProxy-v1",
|
| 139 |
+
task_name="bag",
|
| 140 |
+
text_prompt="retrieve the target object from inside the basket and stage it outside the basket",
|
| 141 |
+
mode_order=(
|
| 142 |
+
"base_action",
|
| 143 |
+
"pin_left_rim",
|
| 144 |
+
"pin_right_rim",
|
| 145 |
+
"widen_mouth",
|
| 146 |
+
"maintain_mouth",
|
| 147 |
+
"probe_inside",
|
| 148 |
+
"insert_actor",
|
| 149 |
+
"retrieve",
|
| 150 |
+
),
|
| 151 |
+
reveal_modes=("pin_left_rim", "pin_right_rim", "widen_mouth", "maintain_mouth", "probe_inside"),
|
| 152 |
+
transfer_modes=("insert_actor",),
|
| 153 |
+
retrieve_modes=("retrieve",),
|
| 154 |
+
notes=(
|
| 155 |
+
"Public ManiSkill bridge scene with custom retrieval initialization. The eggplant is placed inside the "
|
| 156 |
+
"basket region and must be pulled back out to a staging zone."
|
| 157 |
+
),
|
| 158 |
+
),
|
| 159 |
+
"cloth": BridgeTaskSpec(
|
| 160 |
+
key="cloth",
|
| 161 |
+
env_id="PutSpoonOnTableClothInScene-v1",
|
| 162 |
+
track_id="cloth_track",
|
| 163 |
+
suite="maniskill3",
|
| 164 |
+
benchmark_task="PutSpoonUnderClothRetrievalProxy-v1",
|
| 165 |
+
task_name="cloth",
|
| 166 |
+
text_prompt="reveal the spoon from under the cloth and retrieve it to the open area",
|
| 167 |
+
mode_order=(
|
| 168 |
+
"base_action",
|
| 169 |
+
"lift_edge",
|
| 170 |
+
"separate_layer",
|
| 171 |
+
"stabilize_fold",
|
| 172 |
+
"maintain_lift",
|
| 173 |
+
"insert_actor",
|
| 174 |
+
"retrieve",
|
| 175 |
+
),
|
| 176 |
+
reveal_modes=("lift_edge", "separate_layer", "stabilize_fold", "maintain_lift"),
|
| 177 |
+
transfer_modes=("insert_actor",),
|
| 178 |
+
retrieve_modes=("retrieve",),
|
| 179 |
+
notes=(
|
| 180 |
+
"Public ManiSkill bridge scene with custom retrieval initialization. The spoon is placed under the "
|
| 181 |
+
"cloth region and must be revealed and extracted to the open side of the table."
|
| 182 |
+
),
|
| 183 |
+
),
|
| 184 |
+
}
|
| 185 |
+
|
| 186 |
+
|
| 187 |
+
def _task_spec(task: str) -> BridgeTaskSpec:
|
| 188 |
+
normalized = str(task).strip().lower()
|
| 189 |
+
if normalized not in TASK_SPECS:
|
| 190 |
+
raise KeyError(f"Unsupported task {task!r}. Expected one of {sorted(TASK_SPECS)}.")
|
| 191 |
+
return TASK_SPECS[normalized]
|
| 192 |
+
|
| 193 |
+
|
| 194 |
+
def _default_paths(task_spec: BridgeTaskSpec) -> SmokePaths:
|
| 195 |
+
return SmokePaths(
|
| 196 |
+
data_dir=WORKSPACE_ROOT / "data" / "maniskill_bridge_retrieval" / f"{task_spec.key}_{SMOKE_VERSION}",
|
| 197 |
+
output_dir=WORKSPACE_ROOT / "outputs" / f"maniskill_{task_spec.key}_{SMOKE_VERSION}",
|
| 198 |
+
report_dir=WORKSPACE_ROOT / "reports" / f"maniskill_{task_spec.key}_{SMOKE_VERSION}",
|
| 199 |
+
)
|
| 200 |
+
|
| 201 |
+
|
| 202 |
+
def _dataset_artifact_path(data_dir: Path, basename: str, *, dataset_seed: int) -> Path:
|
| 203 |
+
if int(dataset_seed) == DEFAULT_SEED:
|
| 204 |
+
return data_dir / basename
|
| 205 |
+
artifact = Path(basename)
|
| 206 |
+
return data_dir / f"{artifact.stem}_seed{int(dataset_seed)}{artifact.suffix}"
|
| 207 |
+
|
| 208 |
+
|
| 209 |
+
def _np(value: Any, *, dtype: np.dtype | None = None) -> np.ndarray:
|
| 210 |
+
if isinstance(value, np.ndarray):
|
| 211 |
+
array = value
|
| 212 |
+
elif isinstance(value, torch.Tensor):
|
| 213 |
+
array = value.detach().cpu().numpy()
|
| 214 |
+
else:
|
| 215 |
+
array = np.asarray(value)
|
| 216 |
+
if dtype is not None:
|
| 217 |
+
array = array.astype(dtype, copy=False)
|
| 218 |
+
return array
|
| 219 |
+
|
| 220 |
+
|
| 221 |
+
def _vec3(value: Any) -> np.ndarray:
|
| 222 |
+
return _np(value, dtype=np.float32).reshape(-1)[:3]
|
| 223 |
+
|
| 224 |
+
|
| 225 |
+
def _resize_rgb(rgb: np.ndarray, size: int) -> np.ndarray:
|
| 226 |
+
tensor = torch.from_numpy(rgb).permute(2, 0, 1).unsqueeze(0).float()
|
| 227 |
+
resized = F.interpolate(tensor, size=(size, size), mode="bilinear", align_corners=False)
|
| 228 |
+
return resized[0].permute(1, 2, 0).round().clamp(0, 255).to(dtype=torch.uint8).cpu().numpy()
|
| 229 |
+
|
| 230 |
+
|
| 231 |
+
def _resize_single_channel(image: np.ndarray, size: int, *, dtype: np.dtype) -> np.ndarray:
|
| 232 |
+
tensor = torch.from_numpy(image).unsqueeze(0).unsqueeze(0).float()
|
| 233 |
+
resized = F.interpolate(tensor, size=(size, size), mode="nearest")
|
| 234 |
+
return resized[0, 0].to(dtype=torch.float32).cpu().numpy().astype(dtype, copy=False)
|
| 235 |
+
|
| 236 |
+
|
| 237 |
+
def _camera_intrinsic_from_param(param: dict[str, Any]) -> np.ndarray:
|
| 238 |
+
for key in ("intrinsic_cv", "intrinsic", "cam_intrinsic"):
|
| 239 |
+
if key in param:
|
| 240 |
+
matrix = _np(param[key], dtype=np.float32)
|
| 241 |
+
return matrix[0] if matrix.ndim == 3 else matrix
|
| 242 |
+
return np.eye(3, dtype=np.float32)
|
| 243 |
+
|
| 244 |
+
|
| 245 |
+
def _camera_extrinsic_from_param(param: dict[str, Any]) -> np.ndarray:
|
| 246 |
+
for key in ("cam2world_gl", "cam2world", "extrinsic_cv", "extrinsic"):
|
| 247 |
+
if key in param:
|
| 248 |
+
matrix = _np(param[key], dtype=np.float32)
|
| 249 |
+
return matrix[0] if matrix.ndim == 3 else matrix
|
| 250 |
+
return np.eye(4, dtype=np.float32)
|
| 251 |
+
|
| 252 |
+
|
| 253 |
+
def _extract_sensor_bundle(obs: dict[str, Any], *, resolution: int) -> dict[str, np.ndarray]:
|
| 254 |
+
camera_name = next(iter(obs["sensor_data"].keys()))
|
| 255 |
+
view = obs["sensor_data"][camera_name]
|
| 256 |
+
param = obs["sensor_param"][camera_name]
|
| 257 |
+
rgb = _np(view["rgb"], dtype=np.uint8)
|
| 258 |
+
segmentation = _np(view["segmentation"], dtype=np.int16)
|
| 259 |
+
rgb = rgb[0] if rgb.ndim == 4 else rgb
|
| 260 |
+
segmentation = segmentation[0] if segmentation.ndim == 4 else segmentation
|
| 261 |
+
segmentation = segmentation[..., 0] if segmentation.ndim == 3 else segmentation
|
| 262 |
+
rgb_resized = _resize_rgb(rgb, resolution)
|
| 263 |
+
seg_resized = _resize_single_channel(segmentation, resolution, dtype=np.int16)
|
| 264 |
+
intrinsic = _camera_intrinsic_from_param(param)
|
| 265 |
+
extrinsic = _camera_extrinsic_from_param(param)
|
| 266 |
+
images = np.stack([rgb_resized.copy() for _ in CAMERA_NAMES], axis=0)
|
| 267 |
+
segmentations = np.stack([seg_resized.copy() for _ in CAMERA_NAMES], axis=0)
|
| 268 |
+
depths = np.zeros((len(CAMERA_NAMES), 1, resolution, resolution), dtype=np.float32)
|
| 269 |
+
depth_valid = np.zeros_like(depths, dtype=np.float32)
|
| 270 |
+
intrinsics = np.stack([intrinsic.copy() for _ in CAMERA_NAMES], axis=0)
|
| 271 |
+
extrinsics = np.stack([extrinsic.copy() for _ in CAMERA_NAMES], axis=0)
|
| 272 |
+
return {
|
| 273 |
+
"images": images,
|
| 274 |
+
"segmentations": segmentations,
|
| 275 |
+
"depths": depths,
|
| 276 |
+
"depth_valid": depth_valid,
|
| 277 |
+
"camera_intrinsics": intrinsics,
|
| 278 |
+
"camera_extrinsics": extrinsics,
|
| 279 |
+
}
|
| 280 |
+
|
| 281 |
+
|
| 282 |
+
def _build_proprio(env: gym.Env[Any, Any]) -> np.ndarray:
|
| 283 |
+
base = env.unwrapped
|
| 284 |
+
qpos = _np(base.agent.robot.get_qpos(), dtype=np.float32).reshape(-1)
|
| 285 |
+
qvel = _np(base.agent.robot.get_qvel(), dtype=np.float32).reshape(-1)
|
| 286 |
+
ee_pose = base.agent.robot.links_map["ee_gripper_link"].pose
|
| 287 |
+
tcp_pose = np.concatenate([_vec3(ee_pose.p), _np(ee_pose.q, dtype=np.float32).reshape(-1)[:4]], axis=0)
|
| 288 |
+
gripper_width = qpos[-2:].sum(keepdims=True).astype(np.float32)
|
| 289 |
+
flat = np.concatenate([qpos, qvel, tcp_pose, gripper_width], axis=0)
|
| 290 |
+
if flat.shape[0] >= PROPRIO_DIM:
|
| 291 |
+
return flat[:PROPRIO_DIM]
|
| 292 |
+
padded = np.zeros((PROPRIO_DIM,), dtype=np.float32)
|
| 293 |
+
padded[: flat.shape[0]] = flat
|
| 294 |
+
return padded
|
| 295 |
+
|
| 296 |
+
|
| 297 |
+
def _source_actor(env: gym.Env[Any, Any]) -> Any:
|
| 298 |
+
base = env.unwrapped
|
| 299 |
+
return base.objs[base.source_obj_name]
|
| 300 |
+
|
| 301 |
+
|
| 302 |
+
def _target_actor(env: gym.Env[Any, Any]) -> Any:
|
| 303 |
+
base = env.unwrapped
|
| 304 |
+
return base.objs[base.target_obj_name]
|
| 305 |
+
|
| 306 |
+
|
| 307 |
+
def _source_position(env: gym.Env[Any, Any]) -> np.ndarray:
|
| 308 |
+
return _vec3(_source_actor(env).pose.p)
|
| 309 |
+
|
| 310 |
+
|
| 311 |
+
def _target_position(env: gym.Env[Any, Any]) -> np.ndarray:
|
| 312 |
+
return _vec3(_target_actor(env).pose.p)
|
| 313 |
+
|
| 314 |
+
|
| 315 |
+
def _ee_position(env: gym.Env[Any, Any]) -> np.ndarray:
|
| 316 |
+
return _vec3(env.unwrapped.agent.robot.links_map["ee_gripper_link"].pose.p)
|
| 317 |
+
|
| 318 |
+
|
| 319 |
+
def _act_from_world_delta(delta_xyz: Sequence[float]) -> np.ndarray:
|
| 320 |
+
delta = np.asarray(delta_xyz, dtype=np.float32).reshape(3)
|
| 321 |
+
return np.asarray([-delta[0], -delta[1], delta[2]], dtype=np.float32)
|
| 322 |
+
|
| 323 |
+
|
| 324 |
+
def _step_action(env: gym.Env[Any, Any], delta_xyz: Sequence[float], *, grip: float) -> None:
|
| 325 |
+
action = np.zeros((1, 7), dtype=np.float32)
|
| 326 |
+
action[0, :3] = np.clip(_act_from_world_delta(delta_xyz), -0.02, 0.02)
|
| 327 |
+
action[0, 6] = float(np.clip(grip, -1.0, 1.0))
|
| 328 |
+
env.step(action)
|
| 329 |
+
|
| 330 |
+
|
| 331 |
+
def _hold(env: gym.Env[Any, Any], *, steps: int, grip: float) -> None:
|
| 332 |
+
for _ in range(int(steps)):
|
| 333 |
+
action = np.zeros((1, 7), dtype=np.float32)
|
| 334 |
+
action[0, 6] = float(np.clip(grip, -1.0, 1.0))
|
| 335 |
+
env.step(action)
|
| 336 |
+
|
| 337 |
+
|
| 338 |
+
def _move_ee(env: gym.Env[Any, Any], goal_xyz: Sequence[float], *, grip: float, max_steps: int = 60, tol: float = 0.003) -> dict[str, Any]:
|
| 339 |
+
goal = np.asarray(goal_xyz, dtype=np.float32).reshape(3)
|
| 340 |
+
for _ in range(int(max_steps)):
|
| 341 |
+
ee = _ee_position(env)
|
| 342 |
+
delta = goal - ee
|
| 343 |
+
if float(np.linalg.norm(delta)) <= float(tol):
|
| 344 |
+
break
|
| 345 |
+
_step_action(env, delta, grip=grip)
|
| 346 |
+
return {"ee_position": _ee_position(env)}
|
| 347 |
+
|
| 348 |
+
|
| 349 |
+
def _repeat_world_delta(env: gym.Env[Any, Any], delta_xyz: Sequence[float], *, grip: float, steps: int) -> None:
|
| 350 |
+
delta = np.asarray(delta_xyz, dtype=np.float32).reshape(3)
|
| 351 |
+
for _ in range(int(steps)):
|
| 352 |
+
_step_action(env, delta, grip=grip)
|
| 353 |
+
|
| 354 |
+
|
| 355 |
+
def _snapshot_env(env: gym.Env[Any, Any]) -> dict[str, Any]:
|
| 356 |
+
return {"state_dict": env.unwrapped.get_state_dict()}
|
| 357 |
+
|
| 358 |
+
|
| 359 |
+
def _restore_env(env: gym.Env[Any, Any], snapshot: dict[str, Any]) -> None:
|
| 360 |
+
env.unwrapped.set_state_dict(snapshot["state_dict"])
|
| 361 |
+
|
| 362 |
+
|
| 363 |
+
def _sync_env_state(src_env: gym.Env[Any, Any], dst_env: gym.Env[Any, Any]) -> None:
|
| 364 |
+
_restore_env(dst_env, _snapshot_env(src_env))
|
| 365 |
+
|
| 366 |
+
|
| 367 |
+
def _canonical_chunks(task_spec: BridgeTaskSpec) -> dict[str, np.ndarray]:
|
| 368 |
+
base = torch.zeros((1, 8, 14), dtype=torch.float32)
|
| 369 |
+
chunks = {"base_action": base.squeeze(0).numpy().astype(np.float32)}
|
| 370 |
+
for mode_name in task_spec.mode_order[1:]:
|
| 371 |
+
chunk = semantic_macro_chunk(base, task_name=task_spec.task_name, mode_name=mode_name).squeeze(0).cpu().numpy()
|
| 372 |
+
chunks[mode_name] = chunk.astype(np.float32)
|
| 373 |
+
return chunks
|
| 374 |
+
|
| 375 |
+
|
| 376 |
+
def _classify_mode_from_chunk(chunk: np.ndarray, canonical_chunks: dict[str, np.ndarray]) -> str:
|
| 377 |
+
candidate = np.asarray(chunk, dtype=np.float32)
|
| 378 |
+
distances = {
|
| 379 |
+
mode_name: float(np.mean(np.abs(candidate - prototype)))
|
| 380 |
+
for mode_name, prototype in canonical_chunks.items()
|
| 381 |
+
}
|
| 382 |
+
return min(distances, key=distances.get)
|
| 383 |
+
|
| 384 |
+
|
| 385 |
+
def _rng_for_seed(seed: int) -> np.random.Generator:
|
| 386 |
+
return np.random.default_rng(int(seed) + 31)
|
| 387 |
+
|
| 388 |
+
|
| 389 |
+
def _initialize_proxy_state(env: gym.Env[Any, Any], task_spec: BridgeTaskSpec, *, episode_seed: int) -> None:
|
| 390 |
+
base = env.unwrapped
|
| 391 |
+
rng = _rng_for_seed(episode_seed)
|
| 392 |
+
source = _source_actor(env)
|
| 393 |
+
source_pose = source.pose
|
| 394 |
+
source_q = _np(source_pose.q, dtype=np.float32).reshape(-1)[:4]
|
| 395 |
+
if task_spec.key == "bag":
|
| 396 |
+
center = _target_position(env)
|
| 397 |
+
start = center + np.asarray(
|
| 398 |
+
[
|
| 399 |
+
rng.uniform(-0.006, 0.006),
|
| 400 |
+
rng.uniform(-0.010, 0.004),
|
| 401 |
+
0.010 + rng.uniform(-0.002, 0.002),
|
| 402 |
+
],
|
| 403 |
+
dtype=np.float32,
|
| 404 |
+
)
|
| 405 |
+
else:
|
| 406 |
+
cover = _target_actor(env)
|
| 407 |
+
start = np.asarray([CLOTH_FIXED_SOURCE_X, CLOTH_FIXED_SOURCE_Y, CLOTH_FIXED_SOURCE_Z], dtype=np.float32)
|
| 408 |
+
source.set_pose(Pose.create_from_pq(p=start, q=source_q))
|
| 409 |
+
if task_spec.key == "cloth":
|
| 410 |
+
_hold(env, steps=8, grip=1.0)
|
| 411 |
+
cover_pose = cover.pose
|
| 412 |
+
cover_q = _np(cover_pose.q, dtype=np.float32).reshape(-1)[:4]
|
| 413 |
+
cover_start = np.asarray([CLOTH_FIXED_COVER_X, CLOTH_FIXED_COVER_Y, CLOTH_FIXED_COVER_Z], dtype=np.float32)
|
| 414 |
+
cover.set_pose(Pose.create_from_pq(p=cover_start, q=cover_q))
|
| 415 |
+
_hold(env, steps=CLOTH_HIDDEN_SETTLE_STEPS, grip=1.0)
|
| 416 |
+
return
|
| 417 |
+
_hold(env, steps=5, grip=1.0)
|
| 418 |
+
|
| 419 |
+
|
| 420 |
+
def _source_visibility(obs_bundle: dict[str, np.ndarray], actor_id: int) -> float:
|
| 421 |
+
seg = obs_bundle["segmentations"]
|
| 422 |
+
match = float(np.mean(seg == int(actor_id)))
|
| 423 |
+
return float(np.clip(match / VIEW_VISIBILITY_SCALE, 0.0, 1.0))
|
| 424 |
+
|
| 425 |
+
|
| 426 |
+
def _all_positions(env: gym.Env[Any, Any], task_spec: BridgeTaskSpec) -> dict[str, np.ndarray]:
|
| 427 |
+
positions = {"source": _source_position(env), "target": _target_position(env)}
|
| 428 |
+
if task_spec.key == "cloth":
|
| 429 |
+
positions["cover"] = _target_position(env)
|
| 430 |
+
return positions
|
| 431 |
+
|
| 432 |
+
|
| 433 |
+
def _bag_progress(env: gym.Env[Any, Any]) -> float:
|
| 434 |
+
source = _source_position(env)
|
| 435 |
+
center = _target_position(env)
|
| 436 |
+
x_shift = max(source[0] - center[0], 0.0)
|
| 437 |
+
y_pull = max(center[1] - source[1], 0.0)
|
| 438 |
+
z_lift = max(source[2] - center[2], 0.0)
|
| 439 |
+
planar = float(np.linalg.norm((source - center)[:2]))
|
| 440 |
+
return float(
|
| 441 |
+
np.clip(
|
| 442 |
+
0.35 * (x_shift / 0.05)
|
| 443 |
+
+ 0.30 * (y_pull / 0.18)
|
| 444 |
+
+ 0.20 * (z_lift / 0.12)
|
| 445 |
+
+ 0.15 * (planar / 0.12),
|
| 446 |
+
0.0,
|
| 447 |
+
1.0,
|
| 448 |
+
)
|
| 449 |
+
)
|
| 450 |
+
|
| 451 |
+
|
| 452 |
+
def _bag_success(env: gym.Env[Any, Any]) -> bool:
|
| 453 |
+
source = _source_position(env)
|
| 454 |
+
center = _target_position(env)
|
| 455 |
+
planar = float(np.linalg.norm((source - center)[:2]))
|
| 456 |
+
return bool(
|
| 457 |
+
planar >= 0.035
|
| 458 |
+
and (
|
| 459 |
+
source[0] >= center[0] + 0.020
|
| 460 |
+
or source[1] <= center[1] - 0.050
|
| 461 |
+
or source[2] >= center[2] + 0.050
|
| 462 |
+
)
|
| 463 |
+
)
|
| 464 |
+
|
| 465 |
+
|
| 466 |
+
def _cloth_progress(
|
| 467 |
+
env: gym.Env[Any, Any],
|
| 468 |
+
*,
|
| 469 |
+
start_positions: dict[str, np.ndarray],
|
| 470 |
+
current_visibility: float,
|
| 471 |
+
) -> float:
|
| 472 |
+
source = _source_position(env)
|
| 473 |
+
source_start = start_positions["source"]
|
| 474 |
+
y_shift = max(source[1] - source_start[1], 0.0)
|
| 475 |
+
planar = float(np.linalg.norm((source - source_start)[:2]))
|
| 476 |
+
return float(np.clip(0.45 * (y_shift / 0.16) + 0.35 * (planar / 0.16) + 0.20 * current_visibility, 0.0, 1.0))
|
| 477 |
+
|
| 478 |
+
|
| 479 |
+
def _cloth_success(
|
| 480 |
+
env: gym.Env[Any, Any],
|
| 481 |
+
*,
|
| 482 |
+
start_positions: dict[str, np.ndarray],
|
| 483 |
+
current_visibility: float,
|
| 484 |
+
) -> bool:
|
| 485 |
+
source = _source_position(env)
|
| 486 |
+
source_start = start_positions["source"]
|
| 487 |
+
planar = float(np.linalg.norm((source - source_start)[:2]))
|
| 488 |
+
return bool(
|
| 489 |
+
source[1] >= source_start[1] + CLOTH_SUCCESS_MIN_Y_DELTA
|
| 490 |
+
and planar >= CLOTH_SUCCESS_MIN_PLANAR_DELTA
|
| 491 |
+
and current_visibility >= CLOTH_SUCCESS_MIN_VISIBILITY
|
| 492 |
+
)
|
| 493 |
+
|
| 494 |
+
|
| 495 |
+
def _candidate_metrics(
|
| 496 |
+
env: gym.Env[Any, Any],
|
| 497 |
+
*,
|
| 498 |
+
task_spec: BridgeTaskSpec,
|
| 499 |
+
start_positions: dict[str, np.ndarray],
|
| 500 |
+
current_obs_bundle: dict[str, np.ndarray] | None = None,
|
| 501 |
+
) -> dict[str, float]:
|
| 502 |
+
source_actor = _source_actor(env)
|
| 503 |
+
actor_id = int(getattr(source_actor, "per_scene_id", -1))
|
| 504 |
+
visibility = 0.0
|
| 505 |
+
if current_obs_bundle is not None:
|
| 506 |
+
visibility = _source_visibility(current_obs_bundle, actor_id)
|
| 507 |
+
if task_spec.key == "bag":
|
| 508 |
+
progress = _bag_progress(env)
|
| 509 |
+
success = float(_bag_success(env))
|
| 510 |
+
disturbance = 0.0
|
| 511 |
+
access = float(np.clip(0.65 * progress + 0.35 * visibility, 0.0, 1.0))
|
| 512 |
+
else:
|
| 513 |
+
progress = _cloth_progress(env, start_positions=start_positions, current_visibility=visibility)
|
| 514 |
+
success = float(_cloth_success(env, start_positions=start_positions, current_visibility=visibility))
|
| 515 |
+
cloth_start = start_positions["cover"]
|
| 516 |
+
cloth_now = _target_position(env)
|
| 517 |
+
cloth_displacement = float(np.linalg.norm((cloth_now - cloth_start)[:2]))
|
| 518 |
+
disturbance = float(np.clip(max(cloth_displacement - 0.24, 0.0) / 0.14, 0.0, 1.0))
|
| 519 |
+
access = float(np.clip(0.55 * progress + 0.45 * visibility, 0.0, 1.0))
|
| 520 |
+
return {
|
| 521 |
+
"retrieval_success": success,
|
| 522 |
+
"disturbance": disturbance,
|
| 523 |
+
"visibility": visibility,
|
| 524 |
+
"clearance": access,
|
| 525 |
+
"progress": progress,
|
| 526 |
+
}
|
| 527 |
+
|
| 528 |
+
|
| 529 |
+
def _execute_bag_mode(env: gym.Env[Any, Any], mode_name: str) -> None:
|
| 530 |
+
center = _target_position(env)
|
| 531 |
+
source = _source_position(env)
|
| 532 |
+
if mode_name == "retrieve":
|
| 533 |
+
_move_ee(env, source + np.asarray([0.0, 0.0, 0.08], dtype=np.float32), grip=1.0)
|
| 534 |
+
_move_ee(env, source + np.asarray([0.0, 0.0, 0.010], dtype=np.float32), grip=1.0, tol=0.002)
|
| 535 |
+
_hold(env, steps=10, grip=-1.0)
|
| 536 |
+
_move_ee(env, source + np.asarray([0.0, 0.0, 0.12], dtype=np.float32), grip=-1.0)
|
| 537 |
+
_move_ee(env, source + np.asarray([0.0, -0.18, 0.12], dtype=np.float32), grip=-1.0)
|
| 538 |
+
return
|
| 539 |
+
if mode_name == "insert_actor":
|
| 540 |
+
_move_ee(env, source + np.asarray([0.0, 0.0, 0.06], dtype=np.float32), grip=1.0)
|
| 541 |
+
_move_ee(env, source + np.asarray([0.0, 0.0, 0.018], dtype=np.float32), grip=1.0, tol=0.002)
|
| 542 |
+
_hold(env, steps=4, grip=1.0)
|
| 543 |
+
return
|
| 544 |
+
if mode_name == "probe_inside":
|
| 545 |
+
_move_ee(env, source + np.asarray([0.0, 0.0, 0.05], dtype=np.float32), grip=1.0)
|
| 546 |
+
_move_ee(env, source + np.asarray([0.0, 0.0, 0.026], dtype=np.float32), grip=1.0, tol=0.002)
|
| 547 |
+
_repeat_world_delta(env, np.asarray([0.010, -0.004, 0.0], dtype=np.float32), grip=1.0, steps=8)
|
| 548 |
+
return
|
| 549 |
+
if mode_name == "widen_mouth":
|
| 550 |
+
_move_ee(env, center + np.asarray([-0.04, 0.01, 0.06], dtype=np.float32), grip=1.0)
|
| 551 |
+
_move_ee(env, center + np.asarray([-0.03, 0.01, 0.028], dtype=np.float32), grip=1.0, tol=0.003)
|
| 552 |
+
_repeat_world_delta(env, np.asarray([0.012, -0.004, 0.0], dtype=np.float32), grip=1.0, steps=12)
|
| 553 |
+
return
|
| 554 |
+
if mode_name == "pin_left_rim":
|
| 555 |
+
_move_ee(env, center + np.asarray([-0.03, 0.01, 0.06], dtype=np.float32), grip=1.0)
|
| 556 |
+
_move_ee(env, center + np.asarray([-0.03, 0.01, 0.028], dtype=np.float32), grip=1.0, tol=0.003)
|
| 557 |
+
_repeat_world_delta(env, np.asarray([0.006, -0.003, 0.0], dtype=np.float32), grip=1.0, steps=8)
|
| 558 |
+
return
|
| 559 |
+
if mode_name == "pin_right_rim":
|
| 560 |
+
_move_ee(env, center + np.asarray([0.03, 0.01, 0.06], dtype=np.float32), grip=1.0)
|
| 561 |
+
_move_ee(env, center + np.asarray([0.03, 0.01, 0.028], dtype=np.float32), grip=1.0, tol=0.003)
|
| 562 |
+
_repeat_world_delta(env, np.asarray([-0.006, -0.003, 0.0], dtype=np.float32), grip=1.0, steps=8)
|
| 563 |
+
return
|
| 564 |
+
if mode_name in {"maintain_mouth", "base_action"}:
|
| 565 |
+
_move_ee(env, center + np.asarray([0.0, 0.0, 0.09], dtype=np.float32), grip=1.0, max_steps=30, tol=0.006)
|
| 566 |
+
_hold(env, steps=3, grip=1.0)
|
| 567 |
+
return
|
| 568 |
+
raise KeyError(f"Unsupported bag mode {mode_name!r}.")
|
| 569 |
+
|
| 570 |
+
|
| 571 |
+
def _execute_cloth_mode(env: gym.Env[Any, Any], mode_name: str) -> None:
|
| 572 |
+
cloth = _target_position(env)
|
| 573 |
+
source = _source_position(env)
|
| 574 |
+
if mode_name == "retrieve":
|
| 575 |
+
_move_ee(env, source + np.asarray([0.0, 0.0, 0.05], dtype=np.float32), grip=1.0)
|
| 576 |
+
_move_ee(env, source + np.asarray([0.0, 0.0, 0.010], dtype=np.float32), grip=1.0, tol=0.002)
|
| 577 |
+
_hold(env, steps=10, grip=-1.0)
|
| 578 |
+
_move_ee(env, source + np.asarray([0.0, 0.0, 0.10], dtype=np.float32), grip=-1.0)
|
| 579 |
+
_move_ee(env, source + np.asarray([0.0, 0.16, 0.10], dtype=np.float32), grip=-1.0)
|
| 580 |
+
return
|
| 581 |
+
if mode_name == "insert_actor":
|
| 582 |
+
_move_ee(env, source + np.asarray([0.0, 0.0, 0.05], dtype=np.float32), grip=1.0)
|
| 583 |
+
_move_ee(env, source + np.asarray([0.0, 0.0, 0.018], dtype=np.float32), grip=1.0, tol=0.002)
|
| 584 |
+
_hold(env, steps=4, grip=1.0)
|
| 585 |
+
return
|
| 586 |
+
if mode_name == "lift_edge":
|
| 587 |
+
_move_ee(env, cloth + np.asarray([0.0, -0.03, 0.05], dtype=np.float32), grip=1.0)
|
| 588 |
+
_move_ee(env, cloth + np.asarray([0.0, -0.03, 0.015], dtype=np.float32), grip=1.0, tol=0.003)
|
| 589 |
+
_repeat_world_delta(env, np.asarray([0.0, 0.006, 0.0], dtype=np.float32), grip=1.0, steps=8)
|
| 590 |
+
return
|
| 591 |
+
if mode_name == "separate_layer":
|
| 592 |
+
_move_ee(env, cloth + np.asarray([-0.04, 0.0, 0.05], dtype=np.float32), grip=1.0)
|
| 593 |
+
_move_ee(env, cloth + np.asarray([-0.04, 0.0, 0.015], dtype=np.float32), grip=1.0, tol=0.003)
|
| 594 |
+
_repeat_world_delta(env, np.asarray([0.008, 0.002, 0.0], dtype=np.float32), grip=1.0, steps=10)
|
| 595 |
+
return
|
| 596 |
+
if mode_name == "stabilize_fold":
|
| 597 |
+
_move_ee(env, cloth + np.asarray([0.0, 0.03, 0.05], dtype=np.float32), grip=1.0)
|
| 598 |
+
_move_ee(env, cloth + np.asarray([0.0, 0.03, 0.015], dtype=np.float32), grip=1.0, tol=0.003)
|
| 599 |
+
_repeat_world_delta(env, np.asarray([0.0, -0.006, 0.0], dtype=np.float32), grip=1.0, steps=8)
|
| 600 |
+
return
|
| 601 |
+
if mode_name in {"maintain_lift", "base_action"}:
|
| 602 |
+
_move_ee(env, cloth + np.asarray([0.0, 0.06, 0.07], dtype=np.float32), grip=1.0, max_steps=30, tol=0.006)
|
| 603 |
+
_hold(env, steps=3, grip=1.0)
|
| 604 |
+
return
|
| 605 |
+
raise KeyError(f"Unsupported cloth mode {mode_name!r}.")
|
| 606 |
+
|
| 607 |
+
|
| 608 |
+
def _execute_mode(env: gym.Env[Any, Any], task_spec: BridgeTaskSpec, mode_name: str) -> None:
|
| 609 |
+
if task_spec.key == "bag":
|
| 610 |
+
_execute_bag_mode(env, mode_name)
|
| 611 |
+
return
|
| 612 |
+
_execute_cloth_mode(env, mode_name)
|
| 613 |
+
|
| 614 |
+
|
| 615 |
+
def _mode_support_mode(task_spec: BridgeTaskSpec, mode_name: str, current_support_mode: int) -> int:
|
| 616 |
+
if mode_name in task_spec.reveal_modes:
|
| 617 |
+
return SUPPORT_MODE_HOLD
|
| 618 |
+
if mode_name in task_spec.transfer_modes:
|
| 619 |
+
return SUPPORT_MODE_TRANSFER
|
| 620 |
+
if mode_name in task_spec.retrieve_modes:
|
| 621 |
+
return SUPPORT_MODE_PASSIVE
|
| 622 |
+
return int(current_support_mode)
|
| 623 |
+
|
| 624 |
+
|
| 625 |
+
def _mode_progress_schedule(task_spec: BridgeTaskSpec, mode_name: str) -> np.ndarray:
|
| 626 |
+
if mode_name in task_spec.reveal_modes:
|
| 627 |
+
return np.asarray([0.18, 0.38, 0.62, 0.84, 1.0], dtype=np.float32)
|
| 628 |
+
if mode_name in task_spec.transfer_modes:
|
| 629 |
+
return np.asarray([0.22, 0.44, 0.66, 0.86, 1.0], dtype=np.float32)
|
| 630 |
+
if mode_name in task_spec.retrieve_modes:
|
| 631 |
+
return np.asarray([0.34, 0.56, 0.76, 0.92, 1.0], dtype=np.float32)
|
| 632 |
+
return np.asarray([0.10, 0.22, 0.34, 0.44, 0.54], dtype=np.float32)
|
| 633 |
+
|
| 634 |
+
|
| 635 |
+
def _scalar_rollout(start: float, end: float, schedule: np.ndarray) -> np.ndarray:
|
| 636 |
+
return np.clip((1.0 - schedule) * float(start) + schedule * float(end), 0.0, 1.0).astype(np.float32)
|
| 637 |
+
|
| 638 |
+
|
| 639 |
+
def _current_state_targets(
|
| 640 |
+
task_spec: BridgeTaskSpec,
|
| 641 |
+
*,
|
| 642 |
+
obs_bundle: dict[str, np.ndarray],
|
| 643 |
+
candidate_metrics: Sequence[dict[str, float]],
|
| 644 |
+
episode_start_positions: dict[str, np.ndarray],
|
| 645 |
+
selected_mode: str,
|
| 646 |
+
env: gym.Env[Any, Any],
|
| 647 |
+
) -> dict[str, Any]:
|
| 648 |
+
metrics_by_name = {mode_name: payload for mode_name, payload in zip(task_spec.mode_order, candidate_metrics)}
|
| 649 |
+
current_metrics = _candidate_metrics(
|
| 650 |
+
env,
|
| 651 |
+
task_spec=task_spec,
|
| 652 |
+
start_positions=episode_start_positions,
|
| 653 |
+
current_obs_bundle=obs_bundle,
|
| 654 |
+
)
|
| 655 |
+
current_disturbance = float(np.clip(current_metrics["disturbance"], 0.0, 1.0))
|
| 656 |
+
current_visibility = float(np.clip(current_metrics["visibility"], 0.0, 1.0))
|
| 657 |
+
current_clearance = float(np.clip(current_metrics["clearance"], 0.0, 1.0))
|
| 658 |
+
current_progress = float(np.clip(current_metrics["progress"], 0.0, 1.0))
|
| 659 |
+
base_gap = float(np.clip(max(current_clearance, current_progress), 0.0, 1.0))
|
| 660 |
+
support_stability = float(np.clip(1.0 - 0.5 * current_disturbance, 0.0, 1.0))
|
| 661 |
+
hold_quality = float(np.clip(0.5 * (support_stability + max(current_clearance, current_progress)), 0.0, 1.0))
|
| 662 |
+
opening_quality = float(np.clip(0.55 * current_progress + 0.25 * current_clearance + 0.20 * current_visibility, 0.0, 1.0))
|
| 663 |
+
actor_feasibility = float(np.clip(0.6 * current_clearance + 0.4 * max(current_visibility, current_progress), 0.0, 1.0))
|
| 664 |
+
reocclusion_rate = float(np.clip(1.0 - max(current_clearance, current_visibility), 0.0, 1.0))
|
| 665 |
+
insertable_actor_corridor = float(np.clip(0.6 * actor_feasibility + 0.4 * base_gap, 0.0, 1.0))
|
| 666 |
+
insertion_corridor = float(np.clip(0.5 * actor_feasibility + 0.5 * base_gap, 0.0, 1.0))
|
| 667 |
+
layer_separation = float(np.clip(0.7 * base_gap + 0.3 * actor_feasibility, 0.0, 1.0))
|
| 668 |
+
fold_preservation = float(np.clip(1.0 - current_disturbance, 0.0, 1.0))
|
| 669 |
+
lift_too_much_risk = float(np.clip(current_disturbance + 0.5 * max(base_gap - 0.5, 0.0), 0.0, 1.0))
|
| 670 |
+
task_metrics = {
|
| 671 |
+
"opening_quality": opening_quality,
|
| 672 |
+
"actor_feasibility_score": actor_feasibility,
|
| 673 |
+
"gap_width": float(0.03 + 0.21 * base_gap),
|
| 674 |
+
"damage_proxy": current_disturbance,
|
| 675 |
+
"release_collapse_rate": reocclusion_rate,
|
| 676 |
+
"target_visibility_confidence": current_visibility,
|
| 677 |
+
"insertable_actor_corridor": insertable_actor_corridor,
|
| 678 |
+
"insertion_corridor": insertion_corridor,
|
| 679 |
+
"hold_quality": hold_quality,
|
| 680 |
+
"layer_separation_quality": layer_separation,
|
| 681 |
+
"fold_preservation": fold_preservation,
|
| 682 |
+
"top_layer_stability": support_stability,
|
| 683 |
+
"lift_too_much_risk": lift_too_much_risk,
|
| 684 |
+
}
|
| 685 |
+
|
| 686 |
+
base_metrics = metrics_by_name["base_action"]
|
| 687 |
+
insert_metrics = metrics_by_name["insert_actor"]
|
| 688 |
+
retrieve_metrics = metrics_by_name["retrieve"]
|
| 689 |
+
reveal_candidates = [metrics_by_name[mode_name] for mode_name in task_spec.reveal_modes]
|
| 690 |
+
reveal_access = max(candidate["candidate_actor_feasibility_auc"] for candidate in reveal_candidates)
|
| 691 |
+
reveal_reveal = max(candidate["candidate_reveal_achieved"] for candidate in reveal_candidates)
|
| 692 |
+
reveal_hold = max(candidate["candidate_hold_persistence"] for candidate in reveal_candidates)
|
| 693 |
+
reveal_visibility = max(candidate["candidate_visibility_integral"] for candidate in reveal_candidates)
|
| 694 |
+
|
| 695 |
+
reveal_corridor = float(
|
| 696 |
+
np.clip(
|
| 697 |
+
0.45 * opening_quality
|
| 698 |
+
+ 0.30 * reveal_access
|
| 699 |
+
+ 0.15 * reveal_reveal
|
| 700 |
+
+ 0.10 * reveal_visibility
|
| 701 |
+
- 0.10 * current_disturbance,
|
| 702 |
+
0.0,
|
| 703 |
+
1.0,
|
| 704 |
+
)
|
| 705 |
+
)
|
| 706 |
+
transfer_corridor = float(
|
| 707 |
+
np.clip(
|
| 708 |
+
0.45 * insertable_actor_corridor
|
| 709 |
+
+ 0.30 * insert_metrics["candidate_actor_feasibility_auc"]
|
| 710 |
+
+ 0.15 * insert_metrics["candidate_reveal_achieved"]
|
| 711 |
+
+ 0.10 * insert_metrics["candidate_visibility_integral"]
|
| 712 |
+
- 0.15 * current_disturbance,
|
| 713 |
+
0.0,
|
| 714 |
+
1.0,
|
| 715 |
+
)
|
| 716 |
+
)
|
| 717 |
+
passive_corridor = float(
|
| 718 |
+
np.clip(
|
| 719 |
+
0.55 * retrieve_metrics["candidate_retrieval_success"]
|
| 720 |
+
+ 0.20 * retrieve_metrics["candidate_actor_feasibility_auc"]
|
| 721 |
+
+ 0.15 * current_progress
|
| 722 |
+
+ 0.10 * current_clearance
|
| 723 |
+
- 0.10 * current_disturbance,
|
| 724 |
+
0.0,
|
| 725 |
+
1.0,
|
| 726 |
+
)
|
| 727 |
+
)
|
| 728 |
+
corridor_feasible = np.stack(
|
| 729 |
+
[
|
| 730 |
+
np.full((NUM_APPROACH_TEMPLATES,), reveal_corridor, dtype=np.float32),
|
| 731 |
+
np.full((NUM_APPROACH_TEMPLATES,), transfer_corridor, dtype=np.float32),
|
| 732 |
+
np.full((NUM_APPROACH_TEMPLATES,), passive_corridor, dtype=np.float32),
|
| 733 |
+
],
|
| 734 |
+
axis=0,
|
| 735 |
+
)
|
| 736 |
+
persistence_horizon = np.asarray(
|
| 737 |
+
[
|
| 738 |
+
ROLL_OUT_HORIZON * float(np.clip(0.35 * hold_quality + 0.35 * reveal_hold + 0.30 * reveal_corridor, 0.0, 1.0)),
|
| 739 |
+
ROLL_OUT_HORIZON
|
| 740 |
+
* float(
|
| 741 |
+
np.clip(
|
| 742 |
+
0.30 * hold_quality + 0.35 * insert_metrics["candidate_hold_persistence"] + 0.35 * transfer_corridor,
|
| 743 |
+
0.0,
|
| 744 |
+
1.0,
|
| 745 |
+
)
|
| 746 |
+
),
|
| 747 |
+
ROLL_OUT_HORIZON
|
| 748 |
+
* float(
|
| 749 |
+
np.clip(
|
| 750 |
+
0.25 * hold_quality + 0.35 * retrieve_metrics["candidate_hold_persistence"] + 0.40 * passive_corridor,
|
| 751 |
+
0.0,
|
| 752 |
+
1.0,
|
| 753 |
+
)
|
| 754 |
+
),
|
| 755 |
+
],
|
| 756 |
+
dtype=np.float32,
|
| 757 |
+
)
|
| 758 |
+
retrieve_margin = float(retrieve_metrics["candidate_utility"] - base_metrics["candidate_utility"])
|
| 759 |
+
insert_margin = float(insert_metrics["candidate_utility"] - base_metrics["candidate_utility"])
|
| 760 |
+
if selected_mode == "retrieve" or (retrieve_metrics["candidate_retrieval_success"] >= 0.5 and retrieve_margin >= 0.12):
|
| 761 |
+
support_mode = SUPPORT_MODE_PASSIVE
|
| 762 |
+
elif selected_mode == "insert_actor" or (insert_margin >= 0.12 and transfer_corridor >= 0.35):
|
| 763 |
+
support_mode = SUPPORT_MODE_TRANSFER
|
| 764 |
+
elif selected_mode in task_spec.reveal_modes:
|
| 765 |
+
support_mode = SUPPORT_MODE_HOLD
|
| 766 |
+
elif selected_mode == "base_action":
|
| 767 |
+
support_mode = SUPPORT_MODE_PASSIVE if passive_corridor >= 0.55 and retrieve_margin >= 0.03 else SUPPORT_MODE_HOLD
|
| 768 |
+
else:
|
| 769 |
+
support_mode = SUPPORT_MODE_HOLD
|
| 770 |
+
best_non_base_utility = max(float(payload["candidate_utility"]) for payload in candidate_metrics[1:])
|
| 771 |
+
intervention_warranted = selected_mode != "base_action" and best_non_base_utility >= float(base_metrics["candidate_utility"]) + 0.12
|
| 772 |
+
return {
|
| 773 |
+
"support_mode": int(support_mode),
|
| 774 |
+
"corridor_feasible": corridor_feasible,
|
| 775 |
+
"persistence_horizon": persistence_horizon,
|
| 776 |
+
"disturbance_cost": np.float32(current_disturbance),
|
| 777 |
+
"state_confidence_target": np.float32(1.0 if intervention_warranted else 0.0),
|
| 778 |
+
"task_metric_mask": STATE_METRIC_MASK.copy(),
|
| 779 |
+
**{metric_name: np.float32(metric_value) for metric_name, metric_value in task_metrics.items()},
|
| 780 |
+
}
|
| 781 |
+
|
| 782 |
+
|
| 783 |
+
def _candidate_rollout_targets(
|
| 784 |
+
task_spec: BridgeTaskSpec,
|
| 785 |
+
*,
|
| 786 |
+
mode_name: str,
|
| 787 |
+
state_targets: dict[str, Any],
|
| 788 |
+
candidate_payload: dict[str, float],
|
| 789 |
+
) -> dict[str, np.ndarray]:
|
| 790 |
+
schedule = _mode_progress_schedule(task_spec, mode_name)
|
| 791 |
+
start_visibility = float(state_targets["target_visibility_confidence"])
|
| 792 |
+
start_access = float(state_targets["actor_feasibility_score"])
|
| 793 |
+
start_persistence = float(np.clip(state_targets["hold_quality"], 0.0, 1.0))
|
| 794 |
+
start_support = float(np.clip(state_targets["top_layer_stability"], 0.0, 1.0))
|
| 795 |
+
start_reocclusion = float(np.clip(state_targets["release_collapse_rate"], 0.0, 1.0))
|
| 796 |
+
start_disturbance = float(np.clip(state_targets["disturbance_cost"], 0.0, 1.0))
|
| 797 |
+
start_clearance = float(np.clip(state_targets["actor_feasibility_score"], 0.0, 1.0))
|
| 798 |
+
start_grasp = float(np.clip(max(start_visibility, start_access), 0.0, 1.0))
|
| 799 |
+
|
| 800 |
+
end_visibility = float(np.clip(candidate_payload["candidate_immediate_visibility"], 0.0, 1.0))
|
| 801 |
+
end_access = float(np.clip(candidate_payload["candidate_immediate_access"], 0.0, 1.0))
|
| 802 |
+
end_progress = float(np.clip(candidate_payload["candidate_immediate_progress"], 0.0, 1.0))
|
| 803 |
+
end_disturbance = float(np.clip(candidate_payload["candidate_immediate_disturbance"], 0.0, 1.0))
|
| 804 |
+
end_support = float(np.clip(candidate_payload["candidate_immediate_support_stability"], 0.0, 1.0))
|
| 805 |
+
end_persistence = float(np.clip(candidate_payload["candidate_immediate_hold_persistence"], 0.0, 1.0))
|
| 806 |
+
end_reocclusion = float(np.clip(candidate_payload["candidate_immediate_reocclusion"], 0.0, 1.0))
|
| 807 |
+
end_clearance = float(np.clip(max(end_access, end_progress), 0.0, 1.0))
|
| 808 |
+
end_grasp = float(np.clip(max(end_visibility, 0.5 * end_access + 0.5 * end_progress), 0.0, 1.0))
|
| 809 |
+
|
| 810 |
+
if mode_name in task_spec.transfer_modes:
|
| 811 |
+
start_visibility = max(start_visibility, 0.35 * end_visibility)
|
| 812 |
+
start_access = max(start_access, 0.40 * end_access)
|
| 813 |
+
start_persistence = max(start_persistence, 0.45 * end_persistence)
|
| 814 |
+
start_support = max(start_support, 0.50 * end_support)
|
| 815 |
+
elif mode_name in task_spec.retrieve_modes:
|
| 816 |
+
start_visibility = max(start_visibility, 0.55 * end_visibility)
|
| 817 |
+
start_access = max(start_access, 0.70 * end_access)
|
| 818 |
+
start_persistence = max(start_persistence, 0.65 * end_persistence)
|
| 819 |
+
start_support = max(start_support, 0.65 * end_support)
|
| 820 |
+
start_reocclusion = min(start_reocclusion, max(0.4 * end_reocclusion, 0.0))
|
| 821 |
+
|
| 822 |
+
visibility = _scalar_rollout(start_visibility, end_visibility, schedule)
|
| 823 |
+
access = _scalar_rollout(start_access, end_access, schedule)
|
| 824 |
+
persistence = _scalar_rollout(start_persistence, end_persistence, schedule)
|
| 825 |
+
support = _scalar_rollout(start_support, end_support, schedule)
|
| 826 |
+
reocclusion = _scalar_rollout(start_reocclusion, end_reocclusion, schedule)
|
| 827 |
+
disturbance = _scalar_rollout(start_disturbance, end_disturbance, schedule)
|
| 828 |
+
clearance = _scalar_rollout(start_clearance, end_clearance, schedule)
|
| 829 |
+
grasp = _scalar_rollout(start_grasp, end_grasp, schedule)
|
| 830 |
+
reveal_corridor = np.clip(0.38 * visibility + 0.34 * access + 0.22 * support - 0.12 * disturbance, 0.0, 1.0)
|
| 831 |
+
transfer_corridor = np.clip(
|
| 832 |
+
0.30 * visibility + 0.38 * access + 0.18 * persistence + 0.14 * support - 0.12 * disturbance,
|
| 833 |
+
0.0,
|
| 834 |
+
1.0,
|
| 835 |
+
)
|
| 836 |
+
passive_corridor = np.clip(
|
| 837 |
+
0.22 * visibility + 0.42 * access + 0.20 * persistence + 0.16 * grasp - 0.14 * disturbance - 0.10 * reocclusion,
|
| 838 |
+
0.0,
|
| 839 |
+
1.0,
|
| 840 |
+
)
|
| 841 |
+
if mode_name in task_spec.reveal_modes:
|
| 842 |
+
reveal_corridor = np.clip(reveal_corridor + 0.14, 0.0, 1.0)
|
| 843 |
+
passive_corridor = np.clip(0.75 * passive_corridor, 0.0, 1.0)
|
| 844 |
+
elif mode_name in task_spec.transfer_modes:
|
| 845 |
+
transfer_corridor = np.clip(transfer_corridor + 0.16, 0.0, 1.0)
|
| 846 |
+
elif mode_name in task_spec.retrieve_modes:
|
| 847 |
+
passive_corridor = np.clip(passive_corridor + 0.20, 0.0, 1.0)
|
| 848 |
+
reveal_corridor = np.clip(0.60 * reveal_corridor, 0.0, 1.0)
|
| 849 |
+
else:
|
| 850 |
+
reveal_corridor = np.clip(0.85 * reveal_corridor, 0.0, 1.0)
|
| 851 |
+
transfer_corridor = np.clip(0.75 * transfer_corridor, 0.0, 1.0)
|
| 852 |
+
passive_corridor = np.clip(0.80 * passive_corridor, 0.0, 1.0)
|
| 853 |
+
corridor_feasible = np.stack(
|
| 854 |
+
[
|
| 855 |
+
np.repeat(reveal_corridor[:, None], NUM_APPROACH_TEMPLATES, axis=1),
|
| 856 |
+
np.repeat(transfer_corridor[:, None], NUM_APPROACH_TEMPLATES, axis=1),
|
| 857 |
+
np.repeat(passive_corridor[:, None], NUM_APPROACH_TEMPLATES, axis=1),
|
| 858 |
+
],
|
| 859 |
+
axis=1,
|
| 860 |
+
).astype(np.float32)
|
| 861 |
+
persistence_horizon = np.stack(
|
| 862 |
+
[
|
| 863 |
+
np.clip(ROLL_OUT_HORIZON * (0.55 * reveal_corridor + 0.45 * support), 0.0, float(ROLL_OUT_HORIZON)),
|
| 864 |
+
np.clip(ROLL_OUT_HORIZON * (0.50 * transfer_corridor + 0.50 * persistence), 0.0, float(ROLL_OUT_HORIZON)),
|
| 865 |
+
np.clip(ROLL_OUT_HORIZON * (0.55 * passive_corridor + 0.45 * persistence), 0.0, float(ROLL_OUT_HORIZON)),
|
| 866 |
+
],
|
| 867 |
+
axis=1,
|
| 868 |
+
).astype(np.float32)
|
| 869 |
+
support_mode = np.full((ROLL_OUT_HORIZON,), _mode_support_mode(task_spec, mode_name, int(state_targets["support_mode"])), dtype=np.int64)
|
| 870 |
+
if mode_name == "base_action":
|
| 871 |
+
support_mode[:] = int(state_targets["support_mode"])
|
| 872 |
+
return {
|
| 873 |
+
"candidate_rollout_support_mode": support_mode,
|
| 874 |
+
"candidate_rollout_corridor_feasible": corridor_feasible,
|
| 875 |
+
"candidate_rollout_persistence_horizon": persistence_horizon,
|
| 876 |
+
"candidate_rollout_disturbance_cost": disturbance.astype(np.float32),
|
| 877 |
+
"candidate_rollout_belief_map": visibility[:, None, None].astype(np.float32),
|
| 878 |
+
"candidate_rollout_visibility_map": visibility[:, None, None].astype(np.float32),
|
| 879 |
+
"candidate_rollout_clearance_map": np.repeat(clearance[:, None, None, None], 2, axis=1).astype(np.float32),
|
| 880 |
+
"candidate_rollout_support_stability": support[:, None, None, None].astype(np.float32),
|
| 881 |
+
"candidate_rollout_reocclusion_target": reocclusion[:, None, None].astype(np.float32),
|
| 882 |
+
"candidate_rollout_occluder_contact_map": np.clip(access * support, 0.0, 1.0)[:, None, None].astype(np.float32),
|
| 883 |
+
"candidate_rollout_grasp_affordance_map": grasp[:, None, None].astype(np.float32),
|
| 884 |
+
}
|
| 885 |
+
|
| 886 |
+
|
| 887 |
+
def _select_expert_mode(
|
| 888 |
+
task_spec: BridgeTaskSpec,
|
| 889 |
+
*,
|
| 890 |
+
decision_step: int,
|
| 891 |
+
candidate_metrics: Sequence[dict[str, float]],
|
| 892 |
+
) -> str:
|
| 893 |
+
metrics_by_name = {mode_name: payload for mode_name, payload in zip(task_spec.mode_order, candidate_metrics)}
|
| 894 |
+
base_utility = float(metrics_by_name["base_action"]["candidate_utility"])
|
| 895 |
+
reveal_best = max(task_spec.reveal_modes, key=lambda name: float(metrics_by_name[name]["candidate_utility"]))
|
| 896 |
+
transfer_best = max(task_spec.transfer_modes, key=lambda name: float(metrics_by_name[name]["candidate_utility"]))
|
| 897 |
+
retrieve_utility = float(metrics_by_name["retrieve"]["candidate_utility"])
|
| 898 |
+
reveal_best_utility = float(metrics_by_name[reveal_best]["candidate_utility"])
|
| 899 |
+
transfer_best_utility = float(metrics_by_name[transfer_best]["candidate_utility"])
|
| 900 |
+
retrieve_success = float(metrics_by_name["retrieve"]["candidate_retrieval_success"])
|
| 901 |
+
|
| 902 |
+
if int(decision_step) > 0 and retrieve_success >= 0.5:
|
| 903 |
+
return "retrieve"
|
| 904 |
+
if int(decision_step) == 0 and reveal_best_utility >= base_utility - 0.02:
|
| 905 |
+
return reveal_best
|
| 906 |
+
if transfer_best_utility >= reveal_best_utility + 0.05 and transfer_best_utility >= base_utility + 0.02:
|
| 907 |
+
return transfer_best
|
| 908 |
+
if reveal_best_utility >= base_utility - 0.02:
|
| 909 |
+
return reveal_best
|
| 910 |
+
if retrieve_success >= 0.5 and retrieve_utility >= base_utility + 0.02:
|
| 911 |
+
return "retrieve"
|
| 912 |
+
if transfer_best_utility >= base_utility + 0.02:
|
| 913 |
+
return transfer_best
|
| 914 |
+
utilities = np.asarray([payload["candidate_utility"] for payload in candidate_metrics], dtype=np.float32)
|
| 915 |
+
return task_spec.mode_order[int(utilities.argmax())]
|
| 916 |
+
|
| 917 |
+
|
| 918 |
+
def _evaluate_candidate(
|
| 919 |
+
task_spec: BridgeTaskSpec,
|
| 920 |
+
sim_env: gym.Env[Any, Any],
|
| 921 |
+
obs_env: gym.Env[Any, Any],
|
| 922 |
+
snapshot: dict[str, Any],
|
| 923 |
+
mode_name: str,
|
| 924 |
+
) -> dict[str, float]:
|
| 925 |
+
_restore_env(sim_env, snapshot)
|
| 926 |
+
start_positions = _all_positions(sim_env, task_spec)
|
| 927 |
+
_execute_mode(sim_env, task_spec, mode_name)
|
| 928 |
+
_sync_env_state(sim_env, obs_env)
|
| 929 |
+
after_bundle = _extract_sensor_bundle(obs_env.get_obs(obs_env.get_info()), resolution=IMAGE_RESOLUTION)
|
| 930 |
+
immediate = _candidate_metrics(sim_env, task_spec=task_spec, start_positions=start_positions, current_obs_bundle=after_bundle)
|
| 931 |
+
if not immediate["retrieval_success"] and mode_name not in {"retrieve", "base_action"}:
|
| 932 |
+
_execute_mode(sim_env, task_spec, "retrieve")
|
| 933 |
+
_sync_env_state(sim_env, obs_env)
|
| 934 |
+
follow_bundle = _extract_sensor_bundle(obs_env.get_obs(obs_env.get_info()), resolution=IMAGE_RESOLUTION)
|
| 935 |
+
final_metrics = _candidate_metrics(sim_env, task_spec=task_spec, start_positions=start_positions, current_obs_bundle=follow_bundle)
|
| 936 |
+
else:
|
| 937 |
+
final_metrics = immediate
|
| 938 |
+
_restore_env(obs_env, snapshot)
|
| 939 |
+
utility = (
|
| 940 |
+
2.5 * final_metrics["retrieval_success"]
|
| 941 |
+
+ 1.0 * final_metrics["progress"]
|
| 942 |
+
+ 0.5 * final_metrics["clearance"]
|
| 943 |
+
+ 0.25 * final_metrics["visibility"]
|
| 944 |
+
- 0.5 * final_metrics["disturbance"]
|
| 945 |
+
)
|
| 946 |
+
return {
|
| 947 |
+
"candidate_retrieval_success": final_metrics["retrieval_success"],
|
| 948 |
+
"candidate_risk": float(np.clip(final_metrics["disturbance"], 0.0, 1.0)),
|
| 949 |
+
"candidate_utility": float(utility),
|
| 950 |
+
"candidate_final_disturbance_cost": final_metrics["disturbance"],
|
| 951 |
+
"candidate_reocclusion_rate": float(np.clip(1.0 - final_metrics["clearance"], 0.0, 1.0)),
|
| 952 |
+
"candidate_visibility_integral": final_metrics["visibility"],
|
| 953 |
+
"candidate_actor_feasibility_auc": final_metrics["clearance"],
|
| 954 |
+
"candidate_reveal_achieved": float(final_metrics["progress"] > 0.15 or final_metrics["clearance"] > 0.35),
|
| 955 |
+
"candidate_hold_persistence": float(1.0 - final_metrics["disturbance"]),
|
| 956 |
+
"candidate_support_stability_auc": float(1.0 - 0.5 * final_metrics["disturbance"]),
|
| 957 |
+
"candidate_disturbance_auc": final_metrics["disturbance"],
|
| 958 |
+
"candidate_immediate_retrieval_success": immediate["retrieval_success"],
|
| 959 |
+
"candidate_immediate_visibility": immediate["visibility"],
|
| 960 |
+
"candidate_immediate_access": immediate["clearance"],
|
| 961 |
+
"candidate_immediate_progress": immediate["progress"],
|
| 962 |
+
"candidate_immediate_reocclusion": float(np.clip(1.0 - immediate["clearance"], 0.0, 1.0)),
|
| 963 |
+
"candidate_immediate_hold_persistence": float(1.0 - immediate["disturbance"]),
|
| 964 |
+
"candidate_immediate_support_stability": float(1.0 - 0.5 * immediate["disturbance"]),
|
| 965 |
+
"candidate_immediate_disturbance": immediate["disturbance"],
|
| 966 |
+
}
|
| 967 |
+
|
| 968 |
+
|
| 969 |
+
def _cloth_seed_is_valid(env: gym.Env[Any, Any], task_spec: BridgeTaskSpec, *, episode_seed: int) -> bool:
|
| 970 |
+
env.reset(seed=int(episode_seed))
|
| 971 |
+
_initialize_proxy_state(env, task_spec, episode_seed=int(episode_seed))
|
| 972 |
+
start_positions = _all_positions(env, task_spec)
|
| 973 |
+
obs = env.get_obs(env.get_info())
|
| 974 |
+
obs_bundle = _extract_sensor_bundle(obs, resolution=IMAGE_RESOLUTION)
|
| 975 |
+
actor_id = int(getattr(_source_actor(env), "per_scene_id", -1))
|
| 976 |
+
start_visibility = _source_visibility(obs_bundle, actor_id)
|
| 977 |
+
if start_visibility > CLOTH_SUCCESS_MIN_VISIBILITY:
|
| 978 |
+
return False
|
| 979 |
+
snapshot = _snapshot_env(env)
|
| 980 |
+
for reveal_mode in ("lift_edge", "separate_layer"):
|
| 981 |
+
_restore_env(env, snapshot)
|
| 982 |
+
_execute_mode(env, task_spec, reveal_mode)
|
| 983 |
+
_execute_mode(env, task_spec, "retrieve")
|
| 984 |
+
obs = env.get_obs(env.get_info())
|
| 985 |
+
obs_bundle = _extract_sensor_bundle(obs, resolution=IMAGE_RESOLUTION)
|
| 986 |
+
visibility = _source_visibility(obs_bundle, actor_id)
|
| 987 |
+
if _cloth_success(env, start_positions=start_positions, current_visibility=visibility):
|
| 988 |
+
return True
|
| 989 |
+
return False
|
| 990 |
+
|
| 991 |
+
|
| 992 |
+
def _build_episode_splits(task_spec: BridgeTaskSpec, spec: SmokeSpec) -> dict[str, list[int]]:
|
| 993 |
+
if task_spec.key != "cloth":
|
| 994 |
+
return {
|
| 995 |
+
"train": [spec.dataset_seed * 10_000 + index for index in range(spec.train_episodes)],
|
| 996 |
+
"val": [spec.dataset_seed * 10_000 + 1_000 + index for index in range(spec.val_episodes)],
|
| 997 |
+
"eval": [spec.dataset_seed * 10_000 + 2_000 + index for index in range(spec.eval_episodes)],
|
| 998 |
+
}
|
| 999 |
+
target_total = int(spec.train_episodes + spec.val_episodes + spec.eval_episodes)
|
| 1000 |
+
valid_seeds: list[int] = []
|
| 1001 |
+
candidate_index = 0
|
| 1002 |
+
env = gym.make(task_spec.env_id, obs_mode="rgb+segmentation", render_mode="rgb_array")
|
| 1003 |
+
try:
|
| 1004 |
+
while len(valid_seeds) < target_total:
|
| 1005 |
+
episode_seed = spec.dataset_seed * 10_000 + candidate_index
|
| 1006 |
+
candidate_index += 1
|
| 1007 |
+
if _cloth_seed_is_valid(env, task_spec, episode_seed=episode_seed):
|
| 1008 |
+
valid_seeds.append(int(episode_seed))
|
| 1009 |
+
print(
|
| 1010 |
+
json.dumps(
|
| 1011 |
+
{
|
| 1012 |
+
"phase": "cloth_seed_selected",
|
| 1013 |
+
"episode_seed": int(episode_seed),
|
| 1014 |
+
"selected": len(valid_seeds),
|
| 1015 |
+
"target_total": target_total,
|
| 1016 |
+
}
|
| 1017 |
+
),
|
| 1018 |
+
flush=True,
|
| 1019 |
+
)
|
| 1020 |
+
if candidate_index > target_total * 30:
|
| 1021 |
+
raise RuntimeError("Unable to find enough physics-valid cloth proxy seeds for the smoke protocol.")
|
| 1022 |
+
finally:
|
| 1023 |
+
env.close()
|
| 1024 |
+
return {
|
| 1025 |
+
"train": valid_seeds[: spec.train_episodes],
|
| 1026 |
+
"val": valid_seeds[spec.train_episodes : spec.train_episodes + spec.val_episodes],
|
| 1027 |
+
"eval": valid_seeds[spec.train_episodes + spec.val_episodes : target_total],
|
| 1028 |
+
}
|
| 1029 |
+
|
| 1030 |
+
|
| 1031 |
+
def _save_episode_splits(output_path: Path, payload: dict[str, list[int]]) -> None:
|
| 1032 |
+
output_path.parent.mkdir(parents=True, exist_ok=True)
|
| 1033 |
+
output_path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")
|
| 1034 |
+
|
| 1035 |
+
|
| 1036 |
+
def _normalize_depth_array(array: np.ndarray) -> np.ndarray:
|
| 1037 |
+
normalized = _np(array, dtype=np.float32)
|
| 1038 |
+
if normalized.ndim >= 4 and normalized.shape[-1] == 1:
|
| 1039 |
+
return np.moveaxis(normalized, -1, normalized.ndim - 3)
|
| 1040 |
+
return normalized
|
| 1041 |
+
|
| 1042 |
+
|
| 1043 |
+
def _candidate_pad_indices(source_candidates: int, expected_candidates: int) -> list[int]:
|
| 1044 |
+
if source_candidates <= 0 or source_candidates >= expected_candidates:
|
| 1045 |
+
return []
|
| 1046 |
+
if source_candidates == 1:
|
| 1047 |
+
return [0] * (expected_candidates - source_candidates)
|
| 1048 |
+
cycle = list(range(1, source_candidates))
|
| 1049 |
+
indices: list[int] = []
|
| 1050 |
+
while len(indices) < (expected_candidates - source_candidates):
|
| 1051 |
+
indices.extend(cycle)
|
| 1052 |
+
return indices[: expected_candidates - source_candidates]
|
| 1053 |
+
|
| 1054 |
+
|
| 1055 |
+
def _pad_candidate_axis(
|
| 1056 |
+
value: Any,
|
| 1057 |
+
*,
|
| 1058 |
+
source_candidates: int,
|
| 1059 |
+
expected_candidates: int,
|
| 1060 |
+
pad_indices: Sequence[int],
|
| 1061 |
+
) -> Any:
|
| 1062 |
+
if source_candidates <= 0 or source_candidates >= expected_candidates:
|
| 1063 |
+
return value
|
| 1064 |
+
if isinstance(value, np.ndarray):
|
| 1065 |
+
if value.ndim == 0 or value.shape[0] != source_candidates:
|
| 1066 |
+
return value
|
| 1067 |
+
if not pad_indices:
|
| 1068 |
+
return value
|
| 1069 |
+
padding = np.take(value, indices=list(pad_indices), axis=0)
|
| 1070 |
+
return np.concatenate([value, padding], axis=0)
|
| 1071 |
+
if isinstance(value, torch.Tensor):
|
| 1072 |
+
if value.ndim == 0 or value.shape[0] != source_candidates:
|
| 1073 |
+
return value
|
| 1074 |
+
if not pad_indices:
|
| 1075 |
+
return value
|
| 1076 |
+
pad_index = torch.as_tensor(list(pad_indices), device=value.device, dtype=torch.long)
|
| 1077 |
+
padding = value.index_select(0, pad_index)
|
| 1078 |
+
return torch.cat([value, padding], dim=0)
|
| 1079 |
+
if isinstance(value, list) and len(value) == source_candidates:
|
| 1080 |
+
padded = list(value)
|
| 1081 |
+
padded.extend(value[index] for index in pad_indices)
|
| 1082 |
+
return padded
|
| 1083 |
+
if isinstance(value, tuple) and len(value) == source_candidates:
|
| 1084 |
+
padded = list(value)
|
| 1085 |
+
padded.extend(value[index] for index in pad_indices)
|
| 1086 |
+
return tuple(padded)
|
| 1087 |
+
return value
|
| 1088 |
+
|
| 1089 |
+
|
| 1090 |
+
def _normalize_candidate_targets(sample: dict[str, Any]) -> dict[str, Any]:
|
| 1091 |
+
candidate_chunks = sample.get("candidate_action_chunks")
|
| 1092 |
+
if candidate_chunks is None:
|
| 1093 |
+
return sample
|
| 1094 |
+
candidate_array = _np(candidate_chunks)
|
| 1095 |
+
if candidate_array.ndim == 0:
|
| 1096 |
+
return sample
|
| 1097 |
+
source_candidates = int(candidate_array.shape[0])
|
| 1098 |
+
if source_candidates >= EXPECTED_PROPOSAL_CANDIDATES:
|
| 1099 |
+
return sample
|
| 1100 |
+
pad_indices = _candidate_pad_indices(source_candidates, EXPECTED_PROPOSAL_CANDIDATES)
|
| 1101 |
+
if not pad_indices:
|
| 1102 |
+
return sample
|
| 1103 |
+
padded = dict(sample)
|
| 1104 |
+
for key, value in sample.items():
|
| 1105 |
+
if not (key.startswith("candidate_") or key.startswith("proposal_target_")):
|
| 1106 |
+
continue
|
| 1107 |
+
padded[key] = _pad_candidate_axis(
|
| 1108 |
+
value,
|
| 1109 |
+
source_candidates=source_candidates,
|
| 1110 |
+
expected_candidates=EXPECTED_PROPOSAL_CANDIDATES,
|
| 1111 |
+
pad_indices=pad_indices,
|
| 1112 |
+
)
|
| 1113 |
+
return padded
|
| 1114 |
+
|
| 1115 |
+
|
| 1116 |
+
def _normalize_cached_samples(samples: Sequence[dict[str, Any]]) -> list[dict[str, Any]]:
|
| 1117 |
+
normalized_samples: list[dict[str, Any]] = []
|
| 1118 |
+
for sample in samples:
|
| 1119 |
+
patched = dict(sample)
|
| 1120 |
+
for key in ("depths", "depth_valid", "history_depths", "history_depth_valid"):
|
| 1121 |
+
if key in patched:
|
| 1122 |
+
patched[key] = _normalize_depth_array(patched[key])
|
| 1123 |
+
patched = _normalize_candidate_targets(patched)
|
| 1124 |
+
normalized_samples.append(patched)
|
| 1125 |
+
return normalized_samples
|
| 1126 |
+
|
| 1127 |
+
|
| 1128 |
+
def _collect_split(
|
| 1129 |
+
*,
|
| 1130 |
+
task_spec: BridgeTaskSpec,
|
| 1131 |
+
canonical_chunks: dict[str, np.ndarray],
|
| 1132 |
+
split_name: str,
|
| 1133 |
+
seeds: Sequence[int],
|
| 1134 |
+
spec: SmokeSpec,
|
| 1135 |
+
output_path: Path,
|
| 1136 |
+
) -> dict[str, Any]:
|
| 1137 |
+
obs_env = gym.make(task_spec.env_id, obs_mode="rgb+segmentation", render_mode="rgb_array")
|
| 1138 |
+
sim_env = gym.make(task_spec.env_id, obs_mode="rgb+segmentation", render_mode="rgb_array")
|
| 1139 |
+
samples: list[dict[str, Any]] = []
|
| 1140 |
+
episode_records: list[dict[str, Any]] = []
|
| 1141 |
+
try:
|
| 1142 |
+
for episode_seed in seeds:
|
| 1143 |
+
obs, _ = obs_env.reset(seed=int(episode_seed))
|
| 1144 |
+
sim_env.reset(seed=int(episode_seed))
|
| 1145 |
+
_initialize_proxy_state(obs_env, task_spec, episode_seed=int(episode_seed))
|
| 1146 |
+
_sync_env_state(obs_env, sim_env)
|
| 1147 |
+
obs = obs_env.get_obs(obs_env.get_info())
|
| 1148 |
+
episode_start_positions = _all_positions(obs_env, task_spec)
|
| 1149 |
+
history: deque[dict[str, Any]] = deque(maxlen=spec.history_steps)
|
| 1150 |
+
episode_success = False
|
| 1151 |
+
for decision_step in range(spec.max_macro_steps):
|
| 1152 |
+
obs_bundle = _extract_sensor_bundle(obs, resolution=spec.resolution)
|
| 1153 |
+
proprio = _build_proprio(obs_env)
|
| 1154 |
+
snapshot = _snapshot_env(obs_env)
|
| 1155 |
+
candidate_metrics = [
|
| 1156 |
+
_evaluate_candidate(task_spec, sim_env, obs_env, snapshot, mode_name) for mode_name in task_spec.mode_order
|
| 1157 |
+
]
|
| 1158 |
+
candidate_chunks = np.stack([canonical_chunks[mode_name] for mode_name in task_spec.mode_order], axis=0).astype(np.float32)
|
| 1159 |
+
utilities = np.asarray([payload["candidate_utility"] for payload in candidate_metrics], dtype=np.float32)
|
| 1160 |
+
selected_mode = _select_expert_mode(task_spec, decision_step=decision_step, candidate_metrics=candidate_metrics)
|
| 1161 |
+
state_targets = _current_state_targets(
|
| 1162 |
+
task_spec,
|
| 1163 |
+
env=obs_env,
|
| 1164 |
+
obs_bundle=obs_bundle,
|
| 1165 |
+
candidate_metrics=candidate_metrics,
|
| 1166 |
+
episode_start_positions=episode_start_positions,
|
| 1167 |
+
selected_mode=selected_mode,
|
| 1168 |
+
)
|
| 1169 |
+
rollout_targets_by_mode = [
|
| 1170 |
+
_candidate_rollout_targets(task_spec, mode_name=mode_name, state_targets=state_targets, candidate_payload=payload)
|
| 1171 |
+
for mode_name, payload in zip(task_spec.mode_order, candidate_metrics)
|
| 1172 |
+
]
|
| 1173 |
+
sample = {
|
| 1174 |
+
"images": obs_bundle["images"].copy(),
|
| 1175 |
+
"depths": obs_bundle["depths"].copy(),
|
| 1176 |
+
"depth_valid": obs_bundle["depth_valid"].copy(),
|
| 1177 |
+
"camera_intrinsics": obs_bundle["camera_intrinsics"].copy(),
|
| 1178 |
+
"camera_extrinsics": obs_bundle["camera_extrinsics"].copy(),
|
| 1179 |
+
"history_images": _history_stack(
|
| 1180 |
+
history,
|
| 1181 |
+
"images",
|
| 1182 |
+
pad_shape=obs_bundle["images"].shape,
|
| 1183 |
+
dtype=np.uint8,
|
| 1184 |
+
history_steps=spec.history_steps,
|
| 1185 |
+
),
|
| 1186 |
+
"history_depths": _history_stack(
|
| 1187 |
+
history,
|
| 1188 |
+
"depths",
|
| 1189 |
+
pad_shape=obs_bundle["depths"].shape,
|
| 1190 |
+
dtype=np.float32,
|
| 1191 |
+
history_steps=spec.history_steps,
|
| 1192 |
+
),
|
| 1193 |
+
"history_depth_valid": _history_stack(
|
| 1194 |
+
history,
|
| 1195 |
+
"depth_valid",
|
| 1196 |
+
pad_shape=obs_bundle["depth_valid"].shape,
|
| 1197 |
+
dtype=np.float32,
|
| 1198 |
+
history_steps=spec.history_steps,
|
| 1199 |
+
),
|
| 1200 |
+
"history_camera_intrinsics": _history_stack(
|
| 1201 |
+
history,
|
| 1202 |
+
"camera_intrinsics",
|
| 1203 |
+
pad_shape=obs_bundle["camera_intrinsics"].shape,
|
| 1204 |
+
dtype=np.float32,
|
| 1205 |
+
history_steps=spec.history_steps,
|
| 1206 |
+
),
|
| 1207 |
+
"history_camera_extrinsics": _history_stack(
|
| 1208 |
+
history,
|
| 1209 |
+
"camera_extrinsics",
|
| 1210 |
+
pad_shape=obs_bundle["camera_extrinsics"].shape,
|
| 1211 |
+
dtype=np.float32,
|
| 1212 |
+
history_steps=spec.history_steps,
|
| 1213 |
+
),
|
| 1214 |
+
"history_proprio": _history_stack(
|
| 1215 |
+
history,
|
| 1216 |
+
"proprio",
|
| 1217 |
+
pad_shape=(PROPRIO_DIM,),
|
| 1218 |
+
dtype=np.float32,
|
| 1219 |
+
history_steps=spec.history_steps,
|
| 1220 |
+
),
|
| 1221 |
+
"history_actions": _history_stack(
|
| 1222 |
+
history,
|
| 1223 |
+
"action",
|
| 1224 |
+
pad_shape=(14,),
|
| 1225 |
+
dtype=np.float32,
|
| 1226 |
+
history_steps=spec.history_steps,
|
| 1227 |
+
),
|
| 1228 |
+
"proprio": proprio.astype(np.float32),
|
| 1229 |
+
"language_goal": task_spec.text_prompt,
|
| 1230 |
+
"task_name": task_spec.task_name,
|
| 1231 |
+
"task_id": TASK_INDEX[task_spec.task_name],
|
| 1232 |
+
"action_chunk": canonical_chunks[selected_mode].copy(),
|
| 1233 |
+
"candidate_action_chunks": candidate_chunks,
|
| 1234 |
+
"candidate_retrieval_success": np.asarray([payload["candidate_retrieval_success"] for payload in candidate_metrics], dtype=np.float32),
|
| 1235 |
+
"candidate_final_disturbance_cost": np.asarray(
|
| 1236 |
+
[payload["candidate_final_disturbance_cost"] for payload in candidate_metrics],
|
| 1237 |
+
dtype=np.float32,
|
| 1238 |
+
),
|
| 1239 |
+
"candidate_reocclusion_rate": np.asarray([payload["candidate_reocclusion_rate"] for payload in candidate_metrics], dtype=np.float32),
|
| 1240 |
+
"candidate_visibility_integral": np.asarray(
|
| 1241 |
+
[payload["candidate_visibility_integral"] for payload in candidate_metrics],
|
| 1242 |
+
dtype=np.float32,
|
| 1243 |
+
),
|
| 1244 |
+
"candidate_actor_feasibility_auc": np.asarray(
|
| 1245 |
+
[payload["candidate_actor_feasibility_auc"] for payload in candidate_metrics],
|
| 1246 |
+
dtype=np.float32,
|
| 1247 |
+
),
|
| 1248 |
+
"candidate_reveal_achieved": np.asarray([payload["candidate_reveal_achieved"] for payload in candidate_metrics], dtype=np.float32),
|
| 1249 |
+
"candidate_hold_persistence": np.asarray([payload["candidate_hold_persistence"] for payload in candidate_metrics], dtype=np.float32),
|
| 1250 |
+
"candidate_support_stability_auc": np.asarray(
|
| 1251 |
+
[payload["candidate_support_stability_auc"] for payload in candidate_metrics],
|
| 1252 |
+
dtype=np.float32,
|
| 1253 |
+
),
|
| 1254 |
+
"candidate_disturbance_auc": np.asarray([payload["candidate_disturbance_auc"] for payload in candidate_metrics], dtype=np.float32),
|
| 1255 |
+
"candidate_risk": np.asarray([payload["candidate_risk"] for payload in candidate_metrics], dtype=np.float32),
|
| 1256 |
+
"candidate_utility": utilities,
|
| 1257 |
+
"candidate_rollout_support_mode": np.stack(
|
| 1258 |
+
[payload["candidate_rollout_support_mode"] for payload in rollout_targets_by_mode],
|
| 1259 |
+
axis=0,
|
| 1260 |
+
).astype(np.int64),
|
| 1261 |
+
"candidate_rollout_corridor_feasible": np.stack(
|
| 1262 |
+
[payload["candidate_rollout_corridor_feasible"] for payload in rollout_targets_by_mode],
|
| 1263 |
+
axis=0,
|
| 1264 |
+
).astype(np.float32),
|
| 1265 |
+
"candidate_rollout_persistence_horizon": np.stack(
|
| 1266 |
+
[payload["candidate_rollout_persistence_horizon"] for payload in rollout_targets_by_mode],
|
| 1267 |
+
axis=0,
|
| 1268 |
+
).astype(np.float32),
|
| 1269 |
+
"candidate_rollout_disturbance_cost": np.stack(
|
| 1270 |
+
[payload["candidate_rollout_disturbance_cost"] for payload in rollout_targets_by_mode],
|
| 1271 |
+
axis=0,
|
| 1272 |
+
).astype(np.float32),
|
| 1273 |
+
"candidate_rollout_belief_map": np.stack(
|
| 1274 |
+
[payload["candidate_rollout_belief_map"] for payload in rollout_targets_by_mode],
|
| 1275 |
+
axis=0,
|
| 1276 |
+
).astype(np.float32),
|
| 1277 |
+
"candidate_rollout_visibility_map": np.stack(
|
| 1278 |
+
[payload["candidate_rollout_visibility_map"] for payload in rollout_targets_by_mode],
|
| 1279 |
+
axis=0,
|
| 1280 |
+
).astype(np.float32),
|
| 1281 |
+
"candidate_rollout_clearance_map": np.stack(
|
| 1282 |
+
[payload["candidate_rollout_clearance_map"] for payload in rollout_targets_by_mode],
|
| 1283 |
+
axis=0,
|
| 1284 |
+
).astype(np.float32),
|
| 1285 |
+
"candidate_rollout_support_stability": np.stack(
|
| 1286 |
+
[payload["candidate_rollout_support_stability"] for payload in rollout_targets_by_mode],
|
| 1287 |
+
axis=0,
|
| 1288 |
+
).astype(np.float32),
|
| 1289 |
+
"candidate_rollout_reocclusion_target": np.stack(
|
| 1290 |
+
[payload["candidate_rollout_reocclusion_target"] for payload in rollout_targets_by_mode],
|
| 1291 |
+
axis=0,
|
| 1292 |
+
).astype(np.float32),
|
| 1293 |
+
"candidate_rollout_occluder_contact_map": np.stack(
|
| 1294 |
+
[payload["candidate_rollout_occluder_contact_map"] for payload in rollout_targets_by_mode],
|
| 1295 |
+
axis=0,
|
| 1296 |
+
).astype(np.float32),
|
| 1297 |
+
"candidate_rollout_grasp_affordance_map": np.stack(
|
| 1298 |
+
[payload["candidate_rollout_grasp_affordance_map"] for payload in rollout_targets_by_mode],
|
| 1299 |
+
axis=0,
|
| 1300 |
+
).astype(np.float32),
|
| 1301 |
+
"proposal_target_rollout_support_mode": np.stack(
|
| 1302 |
+
[payload["candidate_rollout_support_mode"] for payload in rollout_targets_by_mode],
|
| 1303 |
+
axis=0,
|
| 1304 |
+
).astype(np.int64),
|
| 1305 |
+
"proposal_target_rollout_corridor_feasible": np.stack(
|
| 1306 |
+
[payload["candidate_rollout_corridor_feasible"] for payload in rollout_targets_by_mode],
|
| 1307 |
+
axis=0,
|
| 1308 |
+
).astype(np.float32),
|
| 1309 |
+
"proposal_target_rollout_persistence_horizon": np.stack(
|
| 1310 |
+
[payload["candidate_rollout_persistence_horizon"] for payload in rollout_targets_by_mode],
|
| 1311 |
+
axis=0,
|
| 1312 |
+
).astype(np.float32),
|
| 1313 |
+
"proposal_target_rollout_disturbance_cost": np.stack(
|
| 1314 |
+
[payload["candidate_rollout_disturbance_cost"] for payload in rollout_targets_by_mode],
|
| 1315 |
+
axis=0,
|
| 1316 |
+
).astype(np.float32),
|
| 1317 |
+
"proposal_target_rollout_belief_map": np.stack(
|
| 1318 |
+
[payload["candidate_rollout_belief_map"] for payload in rollout_targets_by_mode],
|
| 1319 |
+
axis=0,
|
| 1320 |
+
).astype(np.float32),
|
| 1321 |
+
"proposal_target_rollout_visibility_map": np.stack(
|
| 1322 |
+
[payload["candidate_rollout_visibility_map"] for payload in rollout_targets_by_mode],
|
| 1323 |
+
axis=0,
|
| 1324 |
+
).astype(np.float32),
|
| 1325 |
+
"proposal_target_rollout_clearance_map": np.stack(
|
| 1326 |
+
[payload["candidate_rollout_clearance_map"] for payload in rollout_targets_by_mode],
|
| 1327 |
+
axis=0,
|
| 1328 |
+
).astype(np.float32),
|
| 1329 |
+
"proposal_target_rollout_support_stability": np.stack(
|
| 1330 |
+
[payload["candidate_rollout_support_stability"] for payload in rollout_targets_by_mode],
|
| 1331 |
+
axis=0,
|
| 1332 |
+
).astype(np.float32),
|
| 1333 |
+
"proposal_target_rollout_reocclusion_target": np.stack(
|
| 1334 |
+
[payload["candidate_rollout_reocclusion_target"] for payload in rollout_targets_by_mode],
|
| 1335 |
+
axis=0,
|
| 1336 |
+
).astype(np.float32),
|
| 1337 |
+
"proposal_target_rollout_occluder_contact_map": np.stack(
|
| 1338 |
+
[payload["candidate_rollout_occluder_contact_map"] for payload in rollout_targets_by_mode],
|
| 1339 |
+
axis=0,
|
| 1340 |
+
).astype(np.float32),
|
| 1341 |
+
"proposal_target_rollout_grasp_affordance_map": np.stack(
|
| 1342 |
+
[payload["candidate_rollout_grasp_affordance_map"] for payload in rollout_targets_by_mode],
|
| 1343 |
+
axis=0,
|
| 1344 |
+
).astype(np.float32),
|
| 1345 |
+
"episode_seed": int(episode_seed),
|
| 1346 |
+
"decision_step": int(decision_step),
|
| 1347 |
+
"selected_mode": selected_mode,
|
| 1348 |
+
**state_targets,
|
| 1349 |
+
}
|
| 1350 |
+
samples.append(sample)
|
| 1351 |
+
_execute_mode(obs_env, task_spec, selected_mode)
|
| 1352 |
+
obs = obs_env.get_obs(obs_env.get_info())
|
| 1353 |
+
post_bundle = _extract_sensor_bundle(obs, resolution=spec.resolution)
|
| 1354 |
+
history.append(_init_history_entry(obs_bundle, proprio, canonical_chunks[selected_mode]))
|
| 1355 |
+
if (
|
| 1356 |
+
_bag_success(obs_env)
|
| 1357 |
+
if task_spec.key == "bag"
|
| 1358 |
+
else _cloth_success(
|
| 1359 |
+
obs_env,
|
| 1360 |
+
start_positions=episode_start_positions,
|
| 1361 |
+
current_visibility=_source_visibility(post_bundle, int(getattr(_source_actor(obs_env), "per_scene_id", -1))),
|
| 1362 |
+
)
|
| 1363 |
+
):
|
| 1364 |
+
episode_success = True
|
| 1365 |
+
break
|
| 1366 |
+
episode_records.append({"episode_seed": int(episode_seed), "success": episode_success, "steps": len(history)})
|
| 1367 |
+
print(
|
| 1368 |
+
json.dumps(
|
| 1369 |
+
{
|
| 1370 |
+
"phase": "collect_episode_complete",
|
| 1371 |
+
"task": task_spec.key,
|
| 1372 |
+
"split": split_name,
|
| 1373 |
+
"episode_seed": int(episode_seed),
|
| 1374 |
+
"success": episode_success,
|
| 1375 |
+
"steps": len(history),
|
| 1376 |
+
"samples_collected": len(samples),
|
| 1377 |
+
}
|
| 1378 |
+
),
|
| 1379 |
+
flush=True,
|
| 1380 |
+
)
|
| 1381 |
+
finally:
|
| 1382 |
+
obs_env.close()
|
| 1383 |
+
sim_env.close()
|
| 1384 |
+
payload = {
|
| 1385 |
+
"split_name": split_name,
|
| 1386 |
+
"resolution": spec.resolution,
|
| 1387 |
+
"history_steps": spec.history_steps,
|
| 1388 |
+
"samples": samples,
|
| 1389 |
+
"episode_records": episode_records,
|
| 1390 |
+
}
|
| 1391 |
+
output_path.parent.mkdir(parents=True, exist_ok=True)
|
| 1392 |
+
torch.save(payload, output_path)
|
| 1393 |
+
return payload
|
| 1394 |
+
|
| 1395 |
+
|
| 1396 |
+
def _manual_train_spec(task_spec: BridgeTaskSpec, variant: str, spec: SmokeSpec) -> dict[str, Any]:
|
| 1397 |
+
return {
|
| 1398 |
+
"track_id": task_spec.track_id,
|
| 1399 |
+
"suite": task_spec.suite,
|
| 1400 |
+
"benchmark_task": task_spec.benchmark_task,
|
| 1401 |
+
"model_variant": str(variant),
|
| 1402 |
+
"seed": int(spec.train_seed),
|
| 1403 |
+
"train_demos": int(spec.train_episodes),
|
| 1404 |
+
"val_demos": int(spec.val_episodes),
|
| 1405 |
+
"init_checkpoint_group": str(DEFAULT_INIT_CHECKPOINT),
|
| 1406 |
+
"optimizer": "adamw",
|
| 1407 |
+
"learning_rate": float(spec.learning_rate),
|
| 1408 |
+
"lr_schedule": "constant",
|
| 1409 |
+
"batch_size": int(spec.batch_size),
|
| 1410 |
+
"augmentations": "none",
|
| 1411 |
+
"early_stopping_metric": "val_total",
|
| 1412 |
+
"max_gradient_steps": int(spec.epochs * math.ceil(max(1, spec.train_episodes) / max(1, spec.batch_size))),
|
| 1413 |
+
"unfreeze_scope": "fusion_memory_decoder",
|
| 1414 |
+
"dataset_split_id": (
|
| 1415 |
+
f"{task_spec.key}_{SMOKE_VERSION}_seed{spec.dataset_seed}"
|
| 1416 |
+
if int(spec.dataset_seed) == DEFAULT_SEED
|
| 1417 |
+
else f"{task_spec.key}_{SMOKE_VERSION}_dataset_seed{spec.dataset_seed}"
|
| 1418 |
+
),
|
| 1419 |
+
"same_data_policy": True,
|
| 1420 |
+
"same_init_policy": True,
|
| 1421 |
+
}
|
| 1422 |
+
|
| 1423 |
+
|
| 1424 |
+
def _trainer_config_for_variant(variant: str) -> TrainerConfig:
|
| 1425 |
+
if variant == "trunk_only_ft":
|
| 1426 |
+
return TrainerConfig(
|
| 1427 |
+
policy_type="trunk_only",
|
| 1428 |
+
trainable_parameter_prefixes=("fusion", "memory", "decoder"),
|
| 1429 |
+
eval_mode="trunk_only",
|
| 1430 |
+
)
|
| 1431 |
+
if variant == "adapter_active_ft":
|
| 1432 |
+
return TrainerConfig(
|
| 1433 |
+
policy_type="adapter_wrapped",
|
| 1434 |
+
trainable_parameter_prefixes=(
|
| 1435 |
+
"trunk.fusion",
|
| 1436 |
+
"trunk.memory",
|
| 1437 |
+
"trunk.decoder",
|
| 1438 |
+
"adapter.state_head",
|
| 1439 |
+
"adapter.transition_model",
|
| 1440 |
+
"adapter.proposal_prior",
|
| 1441 |
+
"adapter.planner",
|
| 1442 |
+
),
|
| 1443 |
+
adapter_mode="adapter_active",
|
| 1444 |
+
eval_mode="adapter_active",
|
| 1445 |
+
adapter_use_transition_model=True,
|
| 1446 |
+
adapter_use_task_conditioning=True,
|
| 1447 |
+
adapter_action_supervision_source="trunk",
|
| 1448 |
+
)
|
| 1449 |
+
raise KeyError(f"Unsupported variant {variant!r}.")
|
| 1450 |
+
|
| 1451 |
+
|
| 1452 |
+
def _loss_weights_for_smoke(task_spec: BridgeTaskSpec) -> LossWeights:
|
| 1453 |
+
return LossWeights(
|
| 1454 |
+
action=1.0,
|
| 1455 |
+
support_mode=0.15,
|
| 1456 |
+
corridor=0.15,
|
| 1457 |
+
persistence=0.08,
|
| 1458 |
+
disturbance=0.08,
|
| 1459 |
+
planner_success=0.20,
|
| 1460 |
+
planner_risk=0.08,
|
| 1461 |
+
planner_ranking=0.20,
|
| 1462 |
+
proposal_reconstruction=0.10,
|
| 1463 |
+
proposal_success=0.12,
|
| 1464 |
+
proposal_ranking=0.15,
|
| 1465 |
+
proposal_mode=0.10,
|
| 1466 |
+
proposal_diversity=0.02,
|
| 1467 |
+
task_metrics=0.15,
|
| 1468 |
+
transition=0.25,
|
| 1469 |
+
gate=0.25,
|
| 1470 |
+
calibration=0.10,
|
| 1471 |
+
proposal_mode_task_filter=[task_spec.task_name],
|
| 1472 |
+
)
|
| 1473 |
+
|
| 1474 |
+
|
| 1475 |
+
def _train_variant(
|
| 1476 |
+
*,
|
| 1477 |
+
task_spec: BridgeTaskSpec,
|
| 1478 |
+
variant: str,
|
| 1479 |
+
train_samples: Sequence[dict[str, Any]],
|
| 1480 |
+
val_samples: Sequence[dict[str, Any]],
|
| 1481 |
+
spec: SmokeSpec,
|
| 1482 |
+
output_dir: Path,
|
| 1483 |
+
) -> tuple[Path, dict[str, Any]]:
|
| 1484 |
+
policy_config, _init_trainer_cfg, _init_loss_weights = _load_init_bundle()
|
| 1485 |
+
policy_config = _apply_smoke_planner_overrides(policy_config)
|
| 1486 |
+
trainer_config = _trainer_config_for_variant(variant)
|
| 1487 |
+
loss_weights = _loss_weights_for_smoke(task_spec)
|
| 1488 |
+
model = build_policy(policy_config, trainer_config)
|
| 1489 |
+
init_info = _load_init_checkpoint(model, str(DEFAULT_INIT_CHECKPOINT), False)
|
| 1490 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 1491 |
+
model = model.to(device)
|
| 1492 |
+
torch.manual_seed(spec.train_seed)
|
| 1493 |
+
if torch.cuda.is_available():
|
| 1494 |
+
torch.cuda.manual_seed_all(spec.train_seed)
|
| 1495 |
+
torch.backends.cuda.matmul.allow_tf32 = True
|
| 1496 |
+
matched = apply_trainable_parameter_prefixes(model, trainer_config)
|
| 1497 |
+
optimizer = torch.optim.AdamW(
|
| 1498 |
+
[parameter for parameter in model.parameters() if parameter.requires_grad],
|
| 1499 |
+
lr=spec.learning_rate,
|
| 1500 |
+
weight_decay=spec.weight_decay,
|
| 1501 |
+
)
|
| 1502 |
+
trainer = BimanualTrainer(model=model, optimizer=optimizer, config=trainer_config)
|
| 1503 |
+
train_loader = _make_loader(train_samples, batch_size=spec.batch_size, num_workers=spec.num_workers, shuffle=True)
|
| 1504 |
+
val_loader = _make_loader(val_samples, batch_size=spec.batch_size, num_workers=spec.num_workers, shuffle=False)
|
| 1505 |
+
best_val = math.inf
|
| 1506 |
+
history: list[dict[str, Any]] = []
|
| 1507 |
+
train_spec = _manual_train_spec(task_spec, variant, spec)
|
| 1508 |
+
train_spec["max_gradient_steps"] = len(train_loader) * spec.epochs
|
| 1509 |
+
for epoch in range(spec.epochs):
|
| 1510 |
+
model.train()
|
| 1511 |
+
train_losses: list[dict[str, float]] = []
|
| 1512 |
+
for batch in train_loader:
|
| 1513 |
+
moved = _move_batch_to_device(batch, device)
|
| 1514 |
+
loss_dict = trainer.training_step(moved, loss_weights=loss_weights)
|
| 1515 |
+
train_losses.append({key: float(value.detach().cpu()) for key, value in loss_dict.items()})
|
| 1516 |
+
model.eval()
|
| 1517 |
+
val_losses: list[dict[str, float]] = []
|
| 1518 |
+
with torch.no_grad():
|
| 1519 |
+
for batch in val_loader:
|
| 1520 |
+
moved = _move_batch_to_device(batch, device)
|
| 1521 |
+
forward_kwargs = {
|
| 1522 |
+
"images": moved["images"],
|
| 1523 |
+
"proprio": moved["proprio"],
|
| 1524 |
+
"texts": moved["texts"],
|
| 1525 |
+
"task_names": moved.get("task_name"),
|
| 1526 |
+
"task_ids": moved.get("task_id"),
|
| 1527 |
+
"history_images": moved.get("history_images"),
|
| 1528 |
+
"history_proprio": moved.get("history_proprio"),
|
| 1529 |
+
"history_actions": moved.get("history_actions"),
|
| 1530 |
+
"depths": moved.get("depths"),
|
| 1531 |
+
"depth_valid": moved.get("depth_valid"),
|
| 1532 |
+
"camera_intrinsics": moved.get("camera_intrinsics"),
|
| 1533 |
+
"camera_extrinsics": moved.get("camera_extrinsics"),
|
| 1534 |
+
"history_depths": moved.get("history_depths"),
|
| 1535 |
+
"history_depth_valid": moved.get("history_depth_valid"),
|
| 1536 |
+
"history_camera_intrinsics": moved.get("history_camera_intrinsics"),
|
| 1537 |
+
"history_camera_extrinsics": moved.get("history_camera_extrinsics"),
|
| 1538 |
+
}
|
| 1539 |
+
if variant == "adapter_active_ft":
|
| 1540 |
+
forward_kwargs["adapter_mode"] = "adapter_active"
|
| 1541 |
+
forward_kwargs["use_transition_model"] = True
|
| 1542 |
+
forward_kwargs["use_task_conditioning"] = True
|
| 1543 |
+
outputs = model(**forward_kwargs)
|
| 1544 |
+
losses = compute_total_loss(outputs, moved, weights=loss_weights)
|
| 1545 |
+
val_losses.append({key: float(value.detach().cpu()) for key, value in losses.items()})
|
| 1546 |
+
train_summary = _aggregate_epoch(train_losses)
|
| 1547 |
+
val_summary = _aggregate_epoch(val_losses)
|
| 1548 |
+
history.append({"epoch": epoch, "train": train_summary, "val": val_summary})
|
| 1549 |
+
print(
|
| 1550 |
+
json.dumps(
|
| 1551 |
+
{
|
| 1552 |
+
"phase": "epoch_complete",
|
| 1553 |
+
"task": task_spec.key,
|
| 1554 |
+
"variant": variant,
|
| 1555 |
+
"epoch": epoch,
|
| 1556 |
+
"train_total": train_summary.get("total", 0.0),
|
| 1557 |
+
"val_total": val_summary.get("total", 0.0),
|
| 1558 |
+
}
|
| 1559 |
+
),
|
| 1560 |
+
flush=True,
|
| 1561 |
+
)
|
| 1562 |
+
if val_summary.get("total", math.inf) <= best_val:
|
| 1563 |
+
best_val = val_summary["total"]
|
| 1564 |
+
checkpoint_path = _save_training_checkpoint(
|
| 1565 |
+
output_dir=output_dir,
|
| 1566 |
+
experiment_name=f"{task_spec.key}_{variant}_seed{spec.train_seed}",
|
| 1567 |
+
model=model,
|
| 1568 |
+
policy_config=policy_config,
|
| 1569 |
+
trainer_config=trainer_config,
|
| 1570 |
+
loss_weights=loss_weights,
|
| 1571 |
+
history=history,
|
| 1572 |
+
best_val=best_val,
|
| 1573 |
+
train_spec=train_spec,
|
| 1574 |
+
)
|
| 1575 |
+
(output_dir / "summary.json").write_text(
|
| 1576 |
+
json.dumps(
|
| 1577 |
+
{
|
| 1578 |
+
"task": task_spec.key,
|
| 1579 |
+
"variant": variant,
|
| 1580 |
+
"checkpoint_path": str(checkpoint_path),
|
| 1581 |
+
"init_info": init_info,
|
| 1582 |
+
"trainable_parameter_names": matched,
|
| 1583 |
+
"best_val_total": best_val,
|
| 1584 |
+
"history": history,
|
| 1585 |
+
"train_spec": train_spec,
|
| 1586 |
+
},
|
| 1587 |
+
indent=2,
|
| 1588 |
+
)
|
| 1589 |
+
+ "\n",
|
| 1590 |
+
encoding="utf-8",
|
| 1591 |
+
)
|
| 1592 |
+
return output_dir / "checkpoint_best.pt", train_spec
|
| 1593 |
+
|
| 1594 |
+
|
| 1595 |
+
def _eval_mode_name(model_output: dict[str, Any], result_mode_name: str, canonical_chunks: dict[str, np.ndarray]) -> tuple[str, bool, bool]:
|
| 1596 |
+
if result_mode_name == "adapter_active_ft" and "proposal_mode_names" in model_output and "best_candidate_indices" in model_output:
|
| 1597 |
+
active_mask = bool(_np(model_output.get("adapter_active_mask", np.asarray([False]))).reshape(-1)[0])
|
| 1598 |
+
if not active_mask:
|
| 1599 |
+
return _classify_mode_from_chunk(_np(model_output["action_mean"])[0], canonical_chunks), False, False
|
| 1600 |
+
best_index = int(_np(model_output["best_candidate_indices"])[0])
|
| 1601 |
+
proposal_mode_names = model_output["proposal_mode_names"][0]
|
| 1602 |
+
if best_index < len(proposal_mode_names):
|
| 1603 |
+
mode_name = str(proposal_mode_names[best_index])
|
| 1604 |
+
else:
|
| 1605 |
+
mode_name = _classify_mode_from_chunk(_np(model_output["action_mean"])[0], canonical_chunks)
|
| 1606 |
+
return mode_name, active_mask, bool(best_index > 0)
|
| 1607 |
+
return _classify_mode_from_chunk(_np(model_output["action_mean"])[0], canonical_chunks), False, False
|
| 1608 |
+
|
| 1609 |
+
|
| 1610 |
+
def _manual_eval_protocol(task_spec: BridgeTaskSpec, *, eval_mode: str, spec: SmokeSpec, episodes: int) -> dict[str, Any]:
|
| 1611 |
+
return {
|
| 1612 |
+
"track_id": task_spec.track_id,
|
| 1613 |
+
"suite": task_spec.suite,
|
| 1614 |
+
"benchmark_task": task_spec.benchmark_task,
|
| 1615 |
+
"role": "target",
|
| 1616 |
+
"eval_mode": eval_mode,
|
| 1617 |
+
"seed": int(spec.dataset_seed),
|
| 1618 |
+
"episodes": int(episodes),
|
| 1619 |
+
"resolution": int(spec.resolution),
|
| 1620 |
+
"cameras": tuple(CAMERA_NAMES),
|
| 1621 |
+
"observation_stack": "rgb_triplicate_zero_depth",
|
| 1622 |
+
"action_horizon": 8,
|
| 1623 |
+
"action_space": "widowx_delta_pose",
|
| 1624 |
+
"same_test_episodes": True,
|
| 1625 |
+
}
|
| 1626 |
+
|
| 1627 |
+
|
| 1628 |
+
def _batch_from_obs(
|
| 1629 |
+
task_spec: BridgeTaskSpec,
|
| 1630 |
+
obs_bundle: dict[str, np.ndarray],
|
| 1631 |
+
proprio: np.ndarray,
|
| 1632 |
+
history: Sequence[dict[str, Any]],
|
| 1633 |
+
device: torch.device,
|
| 1634 |
+
) -> dict[str, Any]:
|
| 1635 |
+
return {
|
| 1636 |
+
"images": torch.from_numpy(obs_bundle["images"]).permute(0, 3, 1, 2).unsqueeze(0).float().div(255.0).to(device),
|
| 1637 |
+
"depths": torch.from_numpy(obs_bundle["depths"]).unsqueeze(0).float().to(device),
|
| 1638 |
+
"depth_valid": torch.from_numpy(obs_bundle["depth_valid"]).unsqueeze(0).float().to(device),
|
| 1639 |
+
"camera_intrinsics": torch.from_numpy(obs_bundle["camera_intrinsics"]).unsqueeze(0).float().to(device),
|
| 1640 |
+
"camera_extrinsics": torch.from_numpy(obs_bundle["camera_extrinsics"]).unsqueeze(0).float().to(device),
|
| 1641 |
+
"history_images": torch.from_numpy(
|
| 1642 |
+
_history_stack(history, "images", pad_shape=obs_bundle["images"].shape, dtype=np.uint8, history_steps=HISTORY_STEPS)
|
| 1643 |
+
).permute(0, 1, 4, 2, 3).unsqueeze(0).float().div(255.0).to(device),
|
| 1644 |
+
"history_depths": torch.from_numpy(
|
| 1645 |
+
_history_stack(history, "depths", pad_shape=obs_bundle["depths"].shape, dtype=np.float32, history_steps=HISTORY_STEPS)
|
| 1646 |
+
).unsqueeze(0).float().to(device),
|
| 1647 |
+
"history_depth_valid": torch.from_numpy(
|
| 1648 |
+
_history_stack(
|
| 1649 |
+
history,
|
| 1650 |
+
"depth_valid",
|
| 1651 |
+
pad_shape=obs_bundle["depth_valid"].shape,
|
| 1652 |
+
dtype=np.float32,
|
| 1653 |
+
history_steps=HISTORY_STEPS,
|
| 1654 |
+
)
|
| 1655 |
+
).unsqueeze(0).float().to(device),
|
| 1656 |
+
"history_camera_intrinsics": torch.from_numpy(
|
| 1657 |
+
_history_stack(
|
| 1658 |
+
history,
|
| 1659 |
+
"camera_intrinsics",
|
| 1660 |
+
pad_shape=obs_bundle["camera_intrinsics"].shape,
|
| 1661 |
+
dtype=np.float32,
|
| 1662 |
+
history_steps=HISTORY_STEPS,
|
| 1663 |
+
)
|
| 1664 |
+
).unsqueeze(0).float().to(device),
|
| 1665 |
+
"history_camera_extrinsics": torch.from_numpy(
|
| 1666 |
+
_history_stack(
|
| 1667 |
+
history,
|
| 1668 |
+
"camera_extrinsics",
|
| 1669 |
+
pad_shape=obs_bundle["camera_extrinsics"].shape,
|
| 1670 |
+
dtype=np.float32,
|
| 1671 |
+
history_steps=HISTORY_STEPS,
|
| 1672 |
+
)
|
| 1673 |
+
).unsqueeze(0).float().to(device),
|
| 1674 |
+
"history_proprio": torch.from_numpy(
|
| 1675 |
+
_history_stack(history, "proprio", pad_shape=(PROPRIO_DIM,), dtype=np.float32, history_steps=HISTORY_STEPS)
|
| 1676 |
+
).unsqueeze(0).float().to(device),
|
| 1677 |
+
"history_actions": torch.from_numpy(
|
| 1678 |
+
_history_stack(history, "action", pad_shape=(14,), dtype=np.float32, history_steps=HISTORY_STEPS)
|
| 1679 |
+
).unsqueeze(0).float().to(device),
|
| 1680 |
+
"proprio": torch.from_numpy(proprio).unsqueeze(0).float().to(device),
|
| 1681 |
+
"texts": [task_spec.text_prompt],
|
| 1682 |
+
"task_names": [task_spec.task_name],
|
| 1683 |
+
"task_ids": torch.as_tensor([TASK_INDEX[task_spec.task_name]], dtype=torch.long, device=device),
|
| 1684 |
+
}
|
| 1685 |
+
|
| 1686 |
+
|
| 1687 |
+
def _evaluate_checkpoint(
|
| 1688 |
+
*,
|
| 1689 |
+
task_spec: BridgeTaskSpec,
|
| 1690 |
+
canonical_chunks: dict[str, np.ndarray],
|
| 1691 |
+
checkpoint_path: Path,
|
| 1692 |
+
adapter_mode: str,
|
| 1693 |
+
result_mode_name: str,
|
| 1694 |
+
seeds: Sequence[int],
|
| 1695 |
+
report_path: Path,
|
| 1696 |
+
train_spec: dict[str, Any] | None,
|
| 1697 |
+
planner_overrides: dict[str, float] | None,
|
| 1698 |
+
) -> dict[str, Any]:
|
| 1699 |
+
model, checkpoint = _load_checkpoint(
|
| 1700 |
+
checkpoint_path,
|
| 1701 |
+
adapter_mode=adapter_mode if adapter_mode != "trunk_only" else None,
|
| 1702 |
+
planner_overrides=planner_overrides,
|
| 1703 |
+
)
|
| 1704 |
+
device = next(model.parameters()).device
|
| 1705 |
+
obs_env = gym.make(task_spec.env_id, obs_mode="rgb+segmentation", render_mode="rgb_array")
|
| 1706 |
+
sim_env = gym.make(task_spec.env_id, obs_mode="rgb+segmentation", render_mode="rgb_array")
|
| 1707 |
+
successes: list[int] = []
|
| 1708 |
+
episode_records: list[dict[str, Any]] = []
|
| 1709 |
+
reveal_steps: list[int] = []
|
| 1710 |
+
retrieve_steps: list[int] = []
|
| 1711 |
+
disturbance_values: list[float] = []
|
| 1712 |
+
intervention_events = 0
|
| 1713 |
+
non_base_events = 0
|
| 1714 |
+
total_decisions = 0
|
| 1715 |
+
try:
|
| 1716 |
+
for episode_seed in seeds:
|
| 1717 |
+
obs, _ = obs_env.reset(seed=int(episode_seed))
|
| 1718 |
+
sim_env.reset(seed=int(episode_seed))
|
| 1719 |
+
_initialize_proxy_state(obs_env, task_spec, episode_seed=int(episode_seed))
|
| 1720 |
+
_sync_env_state(obs_env, sim_env)
|
| 1721 |
+
obs = obs_env.get_obs(obs_env.get_info())
|
| 1722 |
+
history: deque[dict[str, Any]] = deque(maxlen=HISTORY_STEPS)
|
| 1723 |
+
episode_start_positions = _all_positions(obs_env, task_spec)
|
| 1724 |
+
success = False
|
| 1725 |
+
first_reveal_step: int | None = None
|
| 1726 |
+
first_retrieve_step: int | None = None
|
| 1727 |
+
episode_disturbance: list[float] = []
|
| 1728 |
+
for decision_step in range(MAX_MACRO_STEPS):
|
| 1729 |
+
obs_bundle = _extract_sensor_bundle(obs, resolution=IMAGE_RESOLUTION)
|
| 1730 |
+
proprio = _build_proprio(obs_env)
|
| 1731 |
+
batch = _batch_from_obs(task_spec, obs_bundle, proprio, list(history), device)
|
| 1732 |
+
with torch.no_grad():
|
| 1733 |
+
if adapter_mode == "trunk_only":
|
| 1734 |
+
outputs = model(**batch)
|
| 1735 |
+
else:
|
| 1736 |
+
outputs = model(**batch, adapter_mode=adapter_mode, use_transition_model=True, use_task_conditioning=True)
|
| 1737 |
+
selected_mode, active_mask, non_base = _eval_mode_name(outputs, result_mode_name, canonical_chunks)
|
| 1738 |
+
start_positions = _all_positions(obs_env, task_spec)
|
| 1739 |
+
_sync_env_state(obs_env, sim_env)
|
| 1740 |
+
_execute_mode(sim_env, task_spec, selected_mode)
|
| 1741 |
+
_sync_env_state(sim_env, obs_env)
|
| 1742 |
+
obs = obs_env.get_obs(obs_env.get_info())
|
| 1743 |
+
post_bundle = _extract_sensor_bundle(obs, resolution=IMAGE_RESOLUTION)
|
| 1744 |
+
end_metrics = _candidate_metrics(
|
| 1745 |
+
obs_env,
|
| 1746 |
+
task_spec=task_spec,
|
| 1747 |
+
start_positions=start_positions,
|
| 1748 |
+
current_obs_bundle=post_bundle,
|
| 1749 |
+
)
|
| 1750 |
+
history.append(_init_history_entry(obs_bundle, proprio, canonical_chunks.get(selected_mode, canonical_chunks["base_action"])))
|
| 1751 |
+
total_decisions += 1
|
| 1752 |
+
intervention_events += int(active_mask)
|
| 1753 |
+
non_base_events += int(non_base)
|
| 1754 |
+
episode_disturbance.append(end_metrics["disturbance"])
|
| 1755 |
+
if selected_mode != "retrieve" and selected_mode != "base_action" and first_reveal_step is None:
|
| 1756 |
+
first_reveal_step = decision_step + 1
|
| 1757 |
+
if selected_mode == "retrieve" and first_retrieve_step is None:
|
| 1758 |
+
first_retrieve_step = decision_step + 1
|
| 1759 |
+
if (
|
| 1760 |
+
_bag_success(obs_env)
|
| 1761 |
+
if task_spec.key == "bag"
|
| 1762 |
+
else _cloth_success(
|
| 1763 |
+
obs_env,
|
| 1764 |
+
start_positions=episode_start_positions,
|
| 1765 |
+
current_visibility=end_metrics["visibility"],
|
| 1766 |
+
)
|
| 1767 |
+
):
|
| 1768 |
+
success = True
|
| 1769 |
+
break
|
| 1770 |
+
successes.append(int(success))
|
| 1771 |
+
if first_reveal_step is not None:
|
| 1772 |
+
reveal_steps.append(first_reveal_step)
|
| 1773 |
+
if first_retrieve_step is not None:
|
| 1774 |
+
retrieve_steps.append(first_retrieve_step)
|
| 1775 |
+
disturbance_values.append(float(np.mean(episode_disturbance)) if episode_disturbance else 0.0)
|
| 1776 |
+
episode_records.append(
|
| 1777 |
+
{
|
| 1778 |
+
"episode_seed": int(episode_seed),
|
| 1779 |
+
"success": success,
|
| 1780 |
+
"steps": len(history),
|
| 1781 |
+
"first_reveal_step": first_reveal_step,
|
| 1782 |
+
"first_retrieve_step": first_retrieve_step,
|
| 1783 |
+
"episode_disturbance": float(np.mean(episode_disturbance)) if episode_disturbance else 0.0,
|
| 1784 |
+
}
|
| 1785 |
+
)
|
| 1786 |
+
print(
|
| 1787 |
+
json.dumps(
|
| 1788 |
+
{
|
| 1789 |
+
"phase": "eval_episode_complete",
|
| 1790 |
+
"task": task_spec.key,
|
| 1791 |
+
"adapter_mode": result_mode_name,
|
| 1792 |
+
"episode_seed": int(episode_seed),
|
| 1793 |
+
"success": success,
|
| 1794 |
+
"steps": len(history),
|
| 1795 |
+
}
|
| 1796 |
+
),
|
| 1797 |
+
flush=True,
|
| 1798 |
+
)
|
| 1799 |
+
finally:
|
| 1800 |
+
obs_env.close()
|
| 1801 |
+
sim_env.close()
|
| 1802 |
+
payload = {
|
| 1803 |
+
"track_id": task_spec.track_id,
|
| 1804 |
+
"suite": task_spec.suite,
|
| 1805 |
+
"benchmark_task": task_spec.benchmark_task,
|
| 1806 |
+
"role": "target",
|
| 1807 |
+
"adapter_mode": result_mode_name,
|
| 1808 |
+
"episodes": len(seeds),
|
| 1809 |
+
"successes": successes,
|
| 1810 |
+
"success_rate": float(np.mean(successes)) if successes else 0.0,
|
| 1811 |
+
"intervention_rate": float(intervention_events / max(1, total_decisions)),
|
| 1812 |
+
"non_base_selection_rate": float(non_base_events / max(1, total_decisions)),
|
| 1813 |
+
"steps_to_first_reveal_or_access": float(np.mean(reveal_steps)) if reveal_steps else float(MAX_MACRO_STEPS),
|
| 1814 |
+
"steps_to_retrieve": float(np.mean(retrieve_steps)) if retrieve_steps else float(MAX_MACRO_STEPS),
|
| 1815 |
+
"disturbance_proxy": float(np.mean(disturbance_values)) if disturbance_values else 0.0,
|
| 1816 |
+
"episode_records": episode_records,
|
| 1817 |
+
"eval_protocol": _manual_eval_protocol(task_spec, eval_mode=result_mode_name, spec=SmokeSpec(), episodes=len(seeds)),
|
| 1818 |
+
"proxy_notes": task_spec.notes,
|
| 1819 |
+
}
|
| 1820 |
+
if train_spec is not None:
|
| 1821 |
+
payload["train_spec"] = train_spec
|
| 1822 |
+
report_path.parent.mkdir(parents=True, exist_ok=True)
|
| 1823 |
+
report_path.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
|
| 1824 |
+
return payload
|
| 1825 |
+
|
| 1826 |
+
|
| 1827 |
+
def _patch_summary_metadata(summary: dict[str, Any], task_spec: BridgeTaskSpec) -> dict[str, Any]:
|
| 1828 |
+
patched = json.loads(json.dumps(summary))
|
| 1829 |
+
track_payload = patched.get("tracks", {}).get(task_spec.track_id)
|
| 1830 |
+
if track_payload is not None:
|
| 1831 |
+
track_payload["suite"] = task_spec.suite
|
| 1832 |
+
track_payload["benchmark_task"] = task_spec.benchmark_task
|
| 1833 |
+
track_payload["notes"] = task_spec.notes
|
| 1834 |
+
track_payload["public_source"] = f"ManiSkill public scene proxy: {task_spec.env_id}"
|
| 1835 |
+
track_payload["task_family"] = f"{task_spec.key}_retrieval_proxy"
|
| 1836 |
+
track_payload["target_behavior"] = task_spec.text_prompt
|
| 1837 |
+
return patched
|
| 1838 |
+
|
| 1839 |
+
|
| 1840 |
+
def _summarize_task(task_spec: BridgeTaskSpec, results: Sequence[dict[str, Any]], output_dir: Path) -> dict[str, Any]:
|
| 1841 |
+
summary = summarize_public_benchmark_package(list(results), allow_partial=True)
|
| 1842 |
+
summary = _patch_summary_metadata(summary, task_spec)
|
| 1843 |
+
output_dir.mkdir(parents=True, exist_ok=True)
|
| 1844 |
+
json_path = output_dir / "public_benchmark_package_summary.json"
|
| 1845 |
+
md_path = output_dir / "public_benchmark_package_summary.md"
|
| 1846 |
+
json_path.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8")
|
| 1847 |
+
track_payload = summary["tracks"][task_spec.track_id]
|
| 1848 |
+
lines = [
|
| 1849 |
+
f"# ManiSkill {task_spec.key.capitalize()} Retrieval Smoke Summary",
|
| 1850 |
+
"",
|
| 1851 |
+
f"- benchmark_task: {task_spec.benchmark_task}",
|
| 1852 |
+
f"- target_macro_average_delta: {summary['target_macro_average_delta']:.3f}",
|
| 1853 |
+
f"- headline_pass: {summary['headline_pass']}",
|
| 1854 |
+
f"- sign_of_life_pass: {summary['sign_of_life_pass']}",
|
| 1855 |
+
"",
|
| 1856 |
+
f"## {task_spec.track_id}",
|
| 1857 |
+
f"- delta_active_vs_trunk: {track_payload.get('delta_active_vs_trunk', 0.0):.3f}",
|
| 1858 |
+
f"- delta_noop_vs_trunk: {track_payload.get('delta_noop_vs_trunk', 0.0):.3f}",
|
| 1859 |
+
f"- signs_of_life: {track_payload.get('signs_of_life', False)}",
|
| 1860 |
+
]
|
| 1861 |
+
if "delta_active_vs_trunk_ci95" in track_payload:
|
| 1862 |
+
low, high = track_payload["delta_active_vs_trunk_ci95"]
|
| 1863 |
+
lines.append(f"- delta_active_vs_trunk_ci95: [{low:.3f}, {high:.3f}]")
|
| 1864 |
+
for mode, mode_payload in track_payload["modes"].items():
|
| 1865 |
+
lines.append(f"- {mode}: mean_success={mode_payload['mean_success']:.3f}")
|
| 1866 |
+
lines.append("")
|
| 1867 |
+
md_path.write_text("\n".join(lines).rstrip() + "\n", encoding="utf-8")
|
| 1868 |
+
return summary
|
| 1869 |
+
|
| 1870 |
+
|
| 1871 |
+
def _parse_args() -> argparse.Namespace:
|
| 1872 |
+
parser = argparse.ArgumentParser(description="Run a fair bridge-scene retrieval smoke for bag or cloth proxy tasks.")
|
| 1873 |
+
parser.add_argument("--task", choices=sorted(TASK_SPECS), required=True)
|
| 1874 |
+
parser.add_argument("--dataset-seed", type=int, default=DEFAULT_SEED)
|
| 1875 |
+
parser.add_argument("--train-seed", type=int, default=DEFAULT_SEED)
|
| 1876 |
+
parser.add_argument("--eval-split", choices=("val", "eval"), default="eval")
|
| 1877 |
+
parser.add_argument("--report-dir", type=Path, default=None)
|
| 1878 |
+
parser.add_argument("--skip-collection", action="store_true")
|
| 1879 |
+
parser.add_argument("--skip-train", action="store_true")
|
| 1880 |
+
parser.add_argument("--skip-eval", action="store_true")
|
| 1881 |
+
parser.add_argument("--reuse-dataset", action="store_true")
|
| 1882 |
+
parser.add_argument("--reuse-checkpoints", action="store_true")
|
| 1883 |
+
parser.add_argument("--adapter-confidence-threshold", type=float, default=None)
|
| 1884 |
+
parser.add_argument("--retrieve-access-threshold", type=float, default=None)
|
| 1885 |
+
parser.add_argument("--retrieve-persistence-threshold", type=float, default=None)
|
| 1886 |
+
parser.add_argument("--retrieve-support-threshold", type=float, default=None)
|
| 1887 |
+
parser.add_argument("--retrieve-reocclusion-threshold", type=float, default=None)
|
| 1888 |
+
parser.add_argument("--planner-mode-preference-bonus", type=float, default=None)
|
| 1889 |
+
parser.add_argument("--planner-premature-retrieve-penalty", type=float, default=None)
|
| 1890 |
+
parser.add_argument("--planner-premature-insert-penalty", type=float, default=None)
|
| 1891 |
+
parser.add_argument("--planner-premature-occlusion-sweep-penalty", type=float, default=None)
|
| 1892 |
+
parser.add_argument("--planner-premature-maintain-penalty", type=float, default=None)
|
| 1893 |
+
parser.add_argument("--planner-retrieve-stage-access-threshold", type=float, default=None)
|
| 1894 |
+
parser.add_argument("--planner-retrieve-stage-reveal-threshold", type=float, default=None)
|
| 1895 |
+
parser.add_argument("--planner-retrieve-stage-persistence-threshold", type=float, default=None)
|
| 1896 |
+
parser.add_argument("--planner-retrieve-stage-support-threshold", type=float, default=None)
|
| 1897 |
+
parser.add_argument("--planner-insert-stage-access-threshold", type=float, default=None)
|
| 1898 |
+
parser.add_argument("--planner-insert-stage-visibility-threshold", type=float, default=None)
|
| 1899 |
+
parser.add_argument("--planner-insert-stage-support-threshold", type=float, default=None)
|
| 1900 |
+
parser.add_argument("--planner-occlusion-maintain-gap-min-access", type=float, default=None)
|
| 1901 |
+
parser.add_argument("--planner-occlusion-maintain-gap-min-visibility", type=float, default=None)
|
| 1902 |
+
return parser.parse_args()
|
| 1903 |
+
|
| 1904 |
+
|
| 1905 |
+
def _planner_overrides_from_args(args: argparse.Namespace) -> dict[str, float]:
|
| 1906 |
+
overrides = {
|
| 1907 |
+
"adapter_confidence_threshold": SMOKE_ADAPTER_CONFIDENCE_THRESHOLD,
|
| 1908 |
+
"retrieve_access_threshold": SMOKE_RETRIEVE_ACCESS_THRESHOLD,
|
| 1909 |
+
"retrieve_persistence_threshold": SMOKE_RETRIEVE_PERSISTENCE_THRESHOLD,
|
| 1910 |
+
"retrieve_support_threshold": SMOKE_RETRIEVE_SUPPORT_THRESHOLD,
|
| 1911 |
+
"retrieve_reocclusion_threshold": SMOKE_RETRIEVE_REOCCLUSION_THRESHOLD,
|
| 1912 |
+
}
|
| 1913 |
+
optional_pairs = (
|
| 1914 |
+
("adapter_confidence_threshold", args.adapter_confidence_threshold),
|
| 1915 |
+
("retrieve_access_threshold", args.retrieve_access_threshold),
|
| 1916 |
+
("retrieve_persistence_threshold", args.retrieve_persistence_threshold),
|
| 1917 |
+
("retrieve_support_threshold", args.retrieve_support_threshold),
|
| 1918 |
+
("retrieve_reocclusion_threshold", args.retrieve_reocclusion_threshold),
|
| 1919 |
+
("mode_preference_bonus", args.planner_mode_preference_bonus),
|
| 1920 |
+
("premature_retrieve_penalty", args.planner_premature_retrieve_penalty),
|
| 1921 |
+
("premature_insert_penalty", args.planner_premature_insert_penalty),
|
| 1922 |
+
("premature_occlusion_sweep_penalty", args.planner_premature_occlusion_sweep_penalty),
|
| 1923 |
+
("premature_maintain_penalty", args.planner_premature_maintain_penalty),
|
| 1924 |
+
("retrieve_stage_access_threshold", args.planner_retrieve_stage_access_threshold),
|
| 1925 |
+
("retrieve_stage_reveal_threshold", args.planner_retrieve_stage_reveal_threshold),
|
| 1926 |
+
("retrieve_stage_persistence_threshold", args.planner_retrieve_stage_persistence_threshold),
|
| 1927 |
+
("retrieve_stage_support_threshold", args.planner_retrieve_stage_support_threshold),
|
| 1928 |
+
("insert_stage_access_threshold", args.planner_insert_stage_access_threshold),
|
| 1929 |
+
("insert_stage_visibility_threshold", args.planner_insert_stage_visibility_threshold),
|
| 1930 |
+
("insert_stage_support_threshold", args.planner_insert_stage_support_threshold),
|
| 1931 |
+
("occlusion_maintain_gap_min_access", args.planner_occlusion_maintain_gap_min_access),
|
| 1932 |
+
("occlusion_maintain_gap_min_visibility", args.planner_occlusion_maintain_gap_min_visibility),
|
| 1933 |
+
)
|
| 1934 |
+
for key, value in optional_pairs:
|
| 1935 |
+
if value is not None:
|
| 1936 |
+
overrides[key] = value
|
| 1937 |
+
return overrides
|
| 1938 |
+
|
| 1939 |
+
|
| 1940 |
+
def main() -> None:
|
| 1941 |
+
args = _parse_args()
|
| 1942 |
+
task_spec = _task_spec(args.task)
|
| 1943 |
+
spec = SmokeSpec(dataset_seed=int(args.dataset_seed), train_seed=int(args.train_seed))
|
| 1944 |
+
canonical_chunks = _canonical_chunks(task_spec)
|
| 1945 |
+
paths = _default_paths(task_spec)
|
| 1946 |
+
report_dir = args.report_dir or paths.report_dir
|
| 1947 |
+
planner_overrides = _planner_overrides_from_args(args)
|
| 1948 |
+
split_path = _dataset_artifact_path(paths.data_dir, "episode_splits.json", dataset_seed=spec.dataset_seed)
|
| 1949 |
+
if split_path.exists():
|
| 1950 |
+
episode_splits = json.loads(split_path.read_text(encoding="utf-8"))
|
| 1951 |
+
else:
|
| 1952 |
+
episode_splits = _build_episode_splits(task_spec, spec)
|
| 1953 |
+
_save_episode_splits(split_path, episode_splits)
|
| 1954 |
+
|
| 1955 |
+
train_path = _dataset_artifact_path(paths.data_dir, "train.pt", dataset_seed=spec.dataset_seed)
|
| 1956 |
+
val_path = _dataset_artifact_path(paths.data_dir, "val.pt", dataset_seed=spec.dataset_seed)
|
| 1957 |
+
if args.skip_collection and (not train_path.exists() or not val_path.exists()):
|
| 1958 |
+
raise FileNotFoundError("Requested --skip-collection but cached dataset files are missing.")
|
| 1959 |
+
if not args.skip_collection and (not args.reuse_dataset or not train_path.exists() or not val_path.exists()):
|
| 1960 |
+
train_payload = _collect_split(
|
| 1961 |
+
task_spec=task_spec,
|
| 1962 |
+
canonical_chunks=canonical_chunks,
|
| 1963 |
+
split_name="train",
|
| 1964 |
+
seeds=episode_splits["train"],
|
| 1965 |
+
spec=spec,
|
| 1966 |
+
output_path=train_path,
|
| 1967 |
+
)
|
| 1968 |
+
val_payload = _collect_split(
|
| 1969 |
+
task_spec=task_spec,
|
| 1970 |
+
canonical_chunks=canonical_chunks,
|
| 1971 |
+
split_name="val",
|
| 1972 |
+
seeds=episode_splits["val"],
|
| 1973 |
+
spec=spec,
|
| 1974 |
+
output_path=val_path,
|
| 1975 |
+
)
|
| 1976 |
+
else:
|
| 1977 |
+
train_payload = torch.load(train_path, map_location="cpu", weights_only=False)
|
| 1978 |
+
val_payload = torch.load(val_path, map_location="cpu", weights_only=False)
|
| 1979 |
+
|
| 1980 |
+
train_samples = _normalize_cached_samples(train_payload["samples"])
|
| 1981 |
+
val_samples = _normalize_cached_samples(val_payload["samples"])
|
| 1982 |
+
checkpoints: dict[str, Path] = {}
|
| 1983 |
+
train_specs: dict[str, dict[str, Any]] = {}
|
| 1984 |
+
for variant in ("trunk_only_ft", "adapter_active_ft"):
|
| 1985 |
+
variant_output_dir = paths.output_dir / f"{variant}_seed{spec.train_seed}"
|
| 1986 |
+
checkpoint_path = variant_output_dir / "checkpoint_best.pt"
|
| 1987 |
+
if args.skip_train and not checkpoint_path.exists():
|
| 1988 |
+
raise FileNotFoundError(f"Requested --skip-train but checkpoint is missing: {checkpoint_path}")
|
| 1989 |
+
if not args.skip_train and (not args.reuse_checkpoints or not checkpoint_path.exists()):
|
| 1990 |
+
checkpoint_path, train_spec = _train_variant(
|
| 1991 |
+
task_spec=task_spec,
|
| 1992 |
+
variant=variant,
|
| 1993 |
+
train_samples=train_samples,
|
| 1994 |
+
val_samples=val_samples,
|
| 1995 |
+
spec=spec,
|
| 1996 |
+
output_dir=variant_output_dir,
|
| 1997 |
+
)
|
| 1998 |
+
else:
|
| 1999 |
+
summary_path = variant_output_dir / "summary.json"
|
| 2000 |
+
if not summary_path.exists():
|
| 2001 |
+
raise FileNotFoundError(f"Missing cached summary file for {variant}: {summary_path}")
|
| 2002 |
+
summary_payload = json.loads(summary_path.read_text(encoding="utf-8"))
|
| 2003 |
+
train_spec = summary_payload["train_spec"]
|
| 2004 |
+
checkpoints[variant] = checkpoint_path
|
| 2005 |
+
train_specs[variant] = train_spec
|
| 2006 |
+
|
| 2007 |
+
results: list[dict[str, Any]] = []
|
| 2008 |
+
if not args.skip_eval:
|
| 2009 |
+
eval_plan = (
|
| 2010 |
+
("trunk_only_ft", checkpoints["trunk_only_ft"], "trunk_only", None),
|
| 2011 |
+
("adapter_noop", checkpoints["adapter_active_ft"], "adapter_noop", None),
|
| 2012 |
+
("adapter_active_ft", checkpoints["adapter_active_ft"], "adapter_active", train_specs["adapter_active_ft"]),
|
| 2013 |
+
)
|
| 2014 |
+
for result_mode_name, checkpoint_path, adapter_mode, train_spec in eval_plan:
|
| 2015 |
+
result = _evaluate_checkpoint(
|
| 2016 |
+
task_spec=task_spec,
|
| 2017 |
+
canonical_chunks=canonical_chunks,
|
| 2018 |
+
checkpoint_path=checkpoint_path,
|
| 2019 |
+
adapter_mode=adapter_mode,
|
| 2020 |
+
result_mode_name=result_mode_name,
|
| 2021 |
+
seeds=episode_splits[args.eval_split],
|
| 2022 |
+
report_path=report_dir / f"{result_mode_name}_seed{spec.train_seed}.json",
|
| 2023 |
+
train_spec=train_spec if result_mode_name != "adapter_noop" else None,
|
| 2024 |
+
planner_overrides=planner_overrides,
|
| 2025 |
+
)
|
| 2026 |
+
if result_mode_name == "trunk_only_ft":
|
| 2027 |
+
result["train_spec"] = train_specs["trunk_only_ft"]
|
| 2028 |
+
(report_dir / f"{result_mode_name}_seed{spec.train_seed}.json").write_text(
|
| 2029 |
+
json.dumps(result, indent=2) + "\n",
|
| 2030 |
+
encoding="utf-8",
|
| 2031 |
+
)
|
| 2032 |
+
results.append(result)
|
| 2033 |
+
_summarize_task(task_spec, results, report_dir)
|
| 2034 |
+
|
| 2035 |
+
|
| 2036 |
+
if __name__ == "__main__":
|
| 2037 |
+
main()
|
code/VLAarchtests4/VLAarchtests2_code/VLAarchtests/code/reveal_vla_bimanual/eval/run_maniskill_pickclutter_smoke.py
ADDED
|
@@ -0,0 +1,2005 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
import argparse
|
| 4 |
+
import collections
|
| 5 |
+
import json
|
| 6 |
+
import math
|
| 7 |
+
import os
|
| 8 |
+
import time
|
| 9 |
+
from collections import deque
|
| 10 |
+
from dataclasses import asdict, dataclass
|
| 11 |
+
from pathlib import Path
|
| 12 |
+
from typing import Any, Iterable, Sequence
|
| 13 |
+
|
| 14 |
+
import numpy as np
|
| 15 |
+
import torch
|
| 16 |
+
from omegaconf import OmegaConf
|
| 17 |
+
from torch import Tensor
|
| 18 |
+
from torch.utils.data import DataLoader, Dataset
|
| 19 |
+
|
| 20 |
+
from eval.public_benchmark_package import build_public_eval_protocol, build_target_training_spec
|
| 21 |
+
from models.action_decoder import ChunkDecoderConfig, semantic_macro_chunk
|
| 22 |
+
from models.backbones import FrozenVLBackboneConfig
|
| 23 |
+
from models.multiview_fusion import MultiViewFusionConfig
|
| 24 |
+
from models.observation_memory import ObservationMemoryConfig
|
| 25 |
+
from models.planner import PlannerConfig
|
| 26 |
+
from models.policy import PolicyConfig
|
| 27 |
+
from models.reveal_head import RevealHeadConfig, TASK_METRIC_NAMES
|
| 28 |
+
from models.world_model import RevealWMConfig
|
| 29 |
+
from train.checkpoint_compat import filter_compatible_state_dict
|
| 30 |
+
from train.losses import LossWeights
|
| 31 |
+
from train.trainer import BimanualTrainer, TrainerConfig, apply_trainable_parameter_prefixes, build_policy
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
def _configure_runtime_env() -> None:
|
| 35 |
+
os.environ.setdefault("VK_ICD_FILENAMES", "/workspace/runtime/vulkan/icd.d/nvidia_icd_egl.json")
|
| 36 |
+
os.environ.setdefault("VK_LAYER_PATH", "/workspace/runtime/vulkan/implicit_layer.d")
|
| 37 |
+
os.environ.setdefault("XDG_RUNTIME_DIR", "/tmp/runtime-root")
|
| 38 |
+
os.environ.setdefault("MS_ASSET_DIR", "/workspace/data/maniskill")
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
_configure_runtime_env()
|
| 42 |
+
|
| 43 |
+
import mani_skill.envs # noqa: E402
|
| 44 |
+
import sapien # noqa: E402
|
| 45 |
+
from mani_skill.envs.tasks.tabletop.pick_clutter_ycb import PickClutterYCBEnv # noqa: E402
|
| 46 |
+
from mani_skill.sensors.camera import CameraConfig # noqa: E402
|
| 47 |
+
from mani_skill.utils import sapien_utils # noqa: E402
|
| 48 |
+
from mani_skill.utils.structs import Pose # noqa: E402
|
| 49 |
+
|
| 50 |
+
from eval.run_public_benchmark_package import summarize_public_benchmark_package # noqa: E402
|
| 51 |
+
from models.action_decoder import TASK_INDEX # noqa: E402
|
| 52 |
+
from train.run_experiment import _load_init_checkpoint, _move_batch_to_device # noqa: E402
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
REPO_ROOT = Path(__file__).resolve().parents[3]
|
| 56 |
+
WORKSPACE_ROOT = Path("/workspace/workspace")
|
| 57 |
+
SMOKE_VERSION = "smoke_v5"
|
| 58 |
+
DEFAULT_DATA_DIR = WORKSPACE_ROOT / "data" / "maniskill_pickclutter" / SMOKE_VERSION
|
| 59 |
+
DEFAULT_OUTPUT_DIR = WORKSPACE_ROOT / "outputs" / f"maniskill_pickclutter_{SMOKE_VERSION}"
|
| 60 |
+
DEFAULT_REPORT_DIR = WORKSPACE_ROOT / "reports" / f"maniskill_pickclutter_{SMOKE_VERSION}"
|
| 61 |
+
DEFAULT_INIT_CHECKPOINT = Path(
|
| 62 |
+
"/workspace/workspace/VLAarchtests2/VLAarchtests/artifacts/outputs/"
|
| 63 |
+
"r3d_handoff_phase/proxy_interaction_r3d_stage3_clip_rgbd_handoff_compact_phase_seed17/checkpoint_best.pt"
|
| 64 |
+
)
|
| 65 |
+
|
| 66 |
+
TEXT_PROMPT = "retrieve the target object from dense clutter and stage it at the front edge"
|
| 67 |
+
TASK_NAME = "foliage"
|
| 68 |
+
TASK_ID = TASK_INDEX[TASK_NAME]
|
| 69 |
+
CAMERA_NAMES = ("front", "left", "right")
|
| 70 |
+
MODE_ORDER = (
|
| 71 |
+
"base_action",
|
| 72 |
+
"sweep_left",
|
| 73 |
+
"sweep_right",
|
| 74 |
+
"pin_canopy",
|
| 75 |
+
"widen_gap",
|
| 76 |
+
"maintain_gap",
|
| 77 |
+
"insert_actor",
|
| 78 |
+
"retrieve",
|
| 79 |
+
)
|
| 80 |
+
ROLL_OUT_HORIZON = 5
|
| 81 |
+
NUM_SUPPORT_MODES = 3
|
| 82 |
+
NUM_APPROACH_TEMPLATES = 32
|
| 83 |
+
SUPPORT_MODE_HOLD = 0
|
| 84 |
+
SUPPORT_MODE_TRANSFER = 1
|
| 85 |
+
SUPPORT_MODE_PASSIVE = 2
|
| 86 |
+
REVEAL_MODES = ("sweep_left", "sweep_right", "pin_canopy", "widen_gap", "maintain_gap")
|
| 87 |
+
TRANSFER_MODES = ("insert_actor",)
|
| 88 |
+
RETRIEVE_MODES = ("retrieve",)
|
| 89 |
+
STATE_SUPERVISION_METRICS = (
|
| 90 |
+
"opening_quality",
|
| 91 |
+
"actor_feasibility_score",
|
| 92 |
+
"gap_width",
|
| 93 |
+
"damage_proxy",
|
| 94 |
+
"release_collapse_rate",
|
| 95 |
+
"target_visibility_confidence",
|
| 96 |
+
"insertable_actor_corridor",
|
| 97 |
+
"insertion_corridor",
|
| 98 |
+
"hold_quality",
|
| 99 |
+
"layer_separation_quality",
|
| 100 |
+
"fold_preservation",
|
| 101 |
+
"top_layer_stability",
|
| 102 |
+
"lift_too_much_risk",
|
| 103 |
+
)
|
| 104 |
+
MAX_MACRO_STEPS = 4
|
| 105 |
+
HISTORY_STEPS = 6
|
| 106 |
+
PROPRIO_DIM = 32
|
| 107 |
+
EXTRACTION_LINE_Y = -0.22
|
| 108 |
+
MIN_CLEARANCE_FOR_SUCCESS = 0.05
|
| 109 |
+
DEFAULT_SEED = 17
|
| 110 |
+
SMOKE_ADAPTER_CONFIDENCE_THRESHOLD = 0.50
|
| 111 |
+
SMOKE_RETRIEVE_ACCESS_THRESHOLD = 0.08
|
| 112 |
+
SMOKE_RETRIEVE_PERSISTENCE_THRESHOLD = 0.12
|
| 113 |
+
SMOKE_RETRIEVE_SUPPORT_THRESHOLD = 0.08
|
| 114 |
+
SMOKE_RETRIEVE_REOCCLUSION_THRESHOLD = 0.92
|
| 115 |
+
|
| 116 |
+
|
| 117 |
+
@dataclass(frozen=True)
|
| 118 |
+
class SmokePaths:
|
| 119 |
+
data_dir: Path = DEFAULT_DATA_DIR
|
| 120 |
+
output_dir: Path = DEFAULT_OUTPUT_DIR
|
| 121 |
+
report_dir: Path = DEFAULT_REPORT_DIR
|
| 122 |
+
|
| 123 |
+
|
| 124 |
+
@dataclass(frozen=True)
|
| 125 |
+
class SmokeSpec:
|
| 126 |
+
resolution: int = 224
|
| 127 |
+
train_episodes: int = 32
|
| 128 |
+
val_episodes: int = 8
|
| 129 |
+
eval_episodes: int = 50
|
| 130 |
+
dataset_seed: int = DEFAULT_SEED
|
| 131 |
+
train_seed: int = DEFAULT_SEED
|
| 132 |
+
history_steps: int = HISTORY_STEPS
|
| 133 |
+
max_macro_steps: int = MAX_MACRO_STEPS
|
| 134 |
+
batch_size: int = 4
|
| 135 |
+
epochs: int = 6
|
| 136 |
+
num_workers: int = 16
|
| 137 |
+
learning_rate: float = 1e-4
|
| 138 |
+
weight_decay: float = 1e-4
|
| 139 |
+
|
| 140 |
+
@property
|
| 141 |
+
def seed(self) -> int:
|
| 142 |
+
return self.train_seed
|
| 143 |
+
|
| 144 |
+
|
| 145 |
+
def _apply_smoke_planner_overrides(
|
| 146 |
+
policy_config: PolicyConfig,
|
| 147 |
+
planner_overrides: dict[str, float] | None = None,
|
| 148 |
+
) -> PolicyConfig:
|
| 149 |
+
policy_config.planner.adapter_confidence_threshold = SMOKE_ADAPTER_CONFIDENCE_THRESHOLD
|
| 150 |
+
policy_config.planner.retrieve_access_threshold = SMOKE_RETRIEVE_ACCESS_THRESHOLD
|
| 151 |
+
policy_config.planner.retrieve_persistence_threshold = SMOKE_RETRIEVE_PERSISTENCE_THRESHOLD
|
| 152 |
+
policy_config.planner.retrieve_support_threshold = SMOKE_RETRIEVE_SUPPORT_THRESHOLD
|
| 153 |
+
policy_config.planner.retrieve_reocclusion_threshold = SMOKE_RETRIEVE_REOCCLUSION_THRESHOLD
|
| 154 |
+
if planner_overrides:
|
| 155 |
+
for key, value in planner_overrides.items():
|
| 156 |
+
if value is None:
|
| 157 |
+
continue
|
| 158 |
+
setattr(policy_config.planner, key, value)
|
| 159 |
+
return policy_config
|
| 160 |
+
|
| 161 |
+
|
| 162 |
+
class PickClutterRevealEnv(PickClutterYCBEnv):
|
| 163 |
+
@property
|
| 164 |
+
def _default_sensor_configs(self):
|
| 165 |
+
resolution = 224
|
| 166 |
+
return [
|
| 167 |
+
CameraConfig(
|
| 168 |
+
"front",
|
| 169 |
+
pose=sapien_utils.look_at(eye=[0.30, 0.00, 0.62], target=[-0.06, 0.00, 0.04]),
|
| 170 |
+
width=resolution,
|
| 171 |
+
height=resolution,
|
| 172 |
+
fov=np.pi / 2,
|
| 173 |
+
near=0.01,
|
| 174 |
+
far=100.0,
|
| 175 |
+
),
|
| 176 |
+
CameraConfig(
|
| 177 |
+
"left",
|
| 178 |
+
pose=sapien_utils.look_at(eye=[0.22, 0.34, 0.34], target=[-0.02, 0.02, 0.03]),
|
| 179 |
+
width=resolution,
|
| 180 |
+
height=resolution,
|
| 181 |
+
fov=np.pi / 2,
|
| 182 |
+
near=0.01,
|
| 183 |
+
far=100.0,
|
| 184 |
+
),
|
| 185 |
+
CameraConfig(
|
| 186 |
+
"right",
|
| 187 |
+
pose=sapien_utils.look_at(eye=[0.22, -0.34, 0.34], target=[-0.02, -0.02, 0.03]),
|
| 188 |
+
width=resolution,
|
| 189 |
+
height=resolution,
|
| 190 |
+
fov=np.pi / 2,
|
| 191 |
+
near=0.01,
|
| 192 |
+
far=100.0,
|
| 193 |
+
),
|
| 194 |
+
]
|
| 195 |
+
|
| 196 |
+
|
| 197 |
+
def _np(value: Any, *, dtype: np.dtype | None = None) -> np.ndarray:
|
| 198 |
+
if isinstance(value, np.ndarray):
|
| 199 |
+
array = value
|
| 200 |
+
elif isinstance(value, Tensor):
|
| 201 |
+
array = value.detach().cpu().numpy()
|
| 202 |
+
else:
|
| 203 |
+
array = np.asarray(value)
|
| 204 |
+
if dtype is not None:
|
| 205 |
+
array = array.astype(dtype, copy=False)
|
| 206 |
+
return array
|
| 207 |
+
|
| 208 |
+
|
| 209 |
+
def _vec3(value: Any) -> np.ndarray:
|
| 210 |
+
return _np(value, dtype=np.float32).reshape(-1)[:3]
|
| 211 |
+
|
| 212 |
+
|
| 213 |
+
def _camera_intrinsic_from_param(param: dict[str, Any]) -> np.ndarray:
|
| 214 |
+
for key in ("intrinsic_cv", "intrinsic", "cam_intrinsic"):
|
| 215 |
+
if key in param:
|
| 216 |
+
matrix = _np(param[key], dtype=np.float32)
|
| 217 |
+
return matrix[0] if matrix.ndim == 3 else matrix
|
| 218 |
+
return np.eye(3, dtype=np.float32)
|
| 219 |
+
|
| 220 |
+
|
| 221 |
+
def _camera_extrinsic_from_param(param: dict[str, Any]) -> np.ndarray:
|
| 222 |
+
for key in ("cam2world_gl", "cam2world", "extrinsic_cv", "extrinsic"):
|
| 223 |
+
if key in param:
|
| 224 |
+
matrix = _np(param[key], dtype=np.float32)
|
| 225 |
+
return matrix[0] if matrix.ndim == 3 else matrix
|
| 226 |
+
return np.eye(4, dtype=np.float32)
|
| 227 |
+
|
| 228 |
+
|
| 229 |
+
def _convert_depth(depth: np.ndarray) -> np.ndarray:
|
| 230 |
+
depth = depth.astype(np.float32, copy=False)
|
| 231 |
+
if np.issubdtype(depth.dtype, np.integer):
|
| 232 |
+
depth = depth / 1000.0
|
| 233 |
+
return depth
|
| 234 |
+
|
| 235 |
+
|
| 236 |
+
def _build_proprio(env: PickClutterRevealEnv) -> np.ndarray:
|
| 237 |
+
base = env.unwrapped
|
| 238 |
+
qpos = _np(base.agent.robot.get_qpos(), dtype=np.float32).reshape(-1)
|
| 239 |
+
qvel = _np(base.agent.robot.get_qvel(), dtype=np.float32).reshape(-1)
|
| 240 |
+
tcp_pose = _np(base.agent.tcp.pose.raw_pose, dtype=np.float32).reshape(-1)
|
| 241 |
+
gripper_width = qpos[-2:].sum(keepdims=True).astype(np.float32)
|
| 242 |
+
pieces = [qpos, qvel, tcp_pose, gripper_width]
|
| 243 |
+
flat = np.concatenate(pieces, axis=0)
|
| 244 |
+
if flat.shape[0] >= PROPRIO_DIM:
|
| 245 |
+
return flat[:PROPRIO_DIM]
|
| 246 |
+
padded = np.zeros((PROPRIO_DIM,), dtype=np.float32)
|
| 247 |
+
padded[: flat.shape[0]] = flat
|
| 248 |
+
return padded
|
| 249 |
+
|
| 250 |
+
|
| 251 |
+
def _extract_sensor_bundle(obs: dict[str, Any]) -> dict[str, np.ndarray]:
|
| 252 |
+
sensor_data = obs["sensor_data"]
|
| 253 |
+
sensor_param = obs["sensor_param"]
|
| 254 |
+
rgb_views: list[np.ndarray] = []
|
| 255 |
+
depth_views: list[np.ndarray] = []
|
| 256 |
+
seg_views: list[np.ndarray] = []
|
| 257 |
+
intrinsics: list[np.ndarray] = []
|
| 258 |
+
extrinsics: list[np.ndarray] = []
|
| 259 |
+
for camera_name in CAMERA_NAMES:
|
| 260 |
+
view = sensor_data[camera_name]
|
| 261 |
+
param = sensor_param[camera_name]
|
| 262 |
+
rgb = _np(view["rgb"], dtype=np.uint8)
|
| 263 |
+
depth = _np(view["depth"])
|
| 264 |
+
segmentation = _np(view["segmentation"])
|
| 265 |
+
rgb = rgb[0] if rgb.ndim == 4 else rgb
|
| 266 |
+
depth = depth[0] if depth.ndim == 4 else depth
|
| 267 |
+
segmentation = segmentation[0] if segmentation.ndim == 4 else segmentation
|
| 268 |
+
if depth.ndim == 3 and depth.shape[-1] == 1:
|
| 269 |
+
depth = depth[..., 0]
|
| 270 |
+
if segmentation.ndim == 3 and segmentation.shape[-1] == 1:
|
| 271 |
+
segmentation = segmentation[..., 0]
|
| 272 |
+
rgb_views.append(rgb.astype(np.uint8, copy=False))
|
| 273 |
+
depth_views.append(_convert_depth(depth))
|
| 274 |
+
seg_views.append(segmentation.astype(np.int32, copy=False))
|
| 275 |
+
intrinsics.append(_camera_intrinsic_from_param(param))
|
| 276 |
+
extrinsics.append(_camera_extrinsic_from_param(param))
|
| 277 |
+
depth_stack = np.stack(depth_views, axis=0).astype(np.float32)
|
| 278 |
+
depth_valid = (depth_stack > 1e-5).astype(np.float32)
|
| 279 |
+
return {
|
| 280 |
+
"images": np.stack(rgb_views, axis=0),
|
| 281 |
+
"depths": depth_stack[:, None, :, :],
|
| 282 |
+
"depth_valid": depth_valid[:, None, :, :],
|
| 283 |
+
"segmentations": np.stack(seg_views, axis=0),
|
| 284 |
+
"camera_intrinsics": np.stack(intrinsics, axis=0).astype(np.float32),
|
| 285 |
+
"camera_extrinsics": np.stack(extrinsics, axis=0).astype(np.float32),
|
| 286 |
+
}
|
| 287 |
+
|
| 288 |
+
|
| 289 |
+
def _target_actor(env: PickClutterRevealEnv) -> Any:
|
| 290 |
+
return env.unwrapped.target_object._objs[0]
|
| 291 |
+
|
| 292 |
+
|
| 293 |
+
def _all_scene_actors(env: PickClutterRevealEnv) -> list[Any]:
|
| 294 |
+
return list(env.unwrapped.all_objects._objs)
|
| 295 |
+
|
| 296 |
+
|
| 297 |
+
def _target_position(env: PickClutterRevealEnv) -> np.ndarray:
|
| 298 |
+
return _vec3(_target_actor(env).pose.p)
|
| 299 |
+
|
| 300 |
+
|
| 301 |
+
def _all_positions(env: PickClutterRevealEnv) -> dict[str, np.ndarray]:
|
| 302 |
+
return {actor.name: _vec3(actor.pose.p) for actor in _all_scene_actors(env)}
|
| 303 |
+
|
| 304 |
+
|
| 305 |
+
def _nearest_non_target_distance(env: PickClutterRevealEnv) -> float:
|
| 306 |
+
target = _target_actor(env)
|
| 307 |
+
target_xy = _vec3(target.pose.p)[:2]
|
| 308 |
+
distances = []
|
| 309 |
+
for actor in _all_scene_actors(env):
|
| 310 |
+
if actor.name == target.name:
|
| 311 |
+
continue
|
| 312 |
+
distances.append(float(np.linalg.norm(_vec3(actor.pose.p)[:2] - target_xy)))
|
| 313 |
+
if not distances:
|
| 314 |
+
return 1.0
|
| 315 |
+
return float(min(distances))
|
| 316 |
+
|
| 317 |
+
|
| 318 |
+
def _success_from_state(env: PickClutterRevealEnv) -> bool:
|
| 319 |
+
target = _target_position(env)
|
| 320 |
+
return bool(target[1] <= EXTRACTION_LINE_Y and _nearest_non_target_distance(env) >= MIN_CLEARANCE_FOR_SUCCESS)
|
| 321 |
+
|
| 322 |
+
|
| 323 |
+
def _clearance_score(env: PickClutterRevealEnv) -> float:
|
| 324 |
+
return float(np.clip((_nearest_non_target_distance(env) - 0.03) / 0.09, 0.0, 1.0))
|
| 325 |
+
|
| 326 |
+
|
| 327 |
+
def _extraction_progress(env: PickClutterRevealEnv) -> float:
|
| 328 |
+
y_value = _target_position(env)[1]
|
| 329 |
+
return float(np.clip(((-0.05) - y_value) / ((-0.05) - EXTRACTION_LINE_Y), 0.0, 1.0))
|
| 330 |
+
|
| 331 |
+
|
| 332 |
+
def _target_visibility(obs_bundle: dict[str, np.ndarray], target_seg_id: int) -> float:
|
| 333 |
+
segmentation = obs_bundle["segmentations"]
|
| 334 |
+
fractions = [(view == int(target_seg_id)).mean() for view in segmentation]
|
| 335 |
+
return float(np.clip(np.mean(fractions) * 80.0, 0.0, 1.0))
|
| 336 |
+
|
| 337 |
+
|
| 338 |
+
def _snapshot_env(env: PickClutterRevealEnv) -> dict[str, Any]:
|
| 339 |
+
base = env.unwrapped
|
| 340 |
+
return {
|
| 341 |
+
"state_dict": base.get_state_dict(),
|
| 342 |
+
"goal_pos": _np(base.goal_pos, dtype=np.float32).copy(),
|
| 343 |
+
}
|
| 344 |
+
|
| 345 |
+
|
| 346 |
+
def _restore_env(env: PickClutterRevealEnv, snapshot: dict[str, Any]) -> None:
|
| 347 |
+
base = env.unwrapped
|
| 348 |
+
state_dict = snapshot["state_dict"]
|
| 349 |
+
goal_pos = torch.as_tensor(snapshot["goal_pos"], dtype=torch.float32, device=base.device)
|
| 350 |
+
base.set_state_dict(state_dict)
|
| 351 |
+
base.goal_pos = goal_pos.view_as(base.goal_pos)
|
| 352 |
+
base.goal_site.set_pose(Pose.create_from_pq(base.goal_pos))
|
| 353 |
+
|
| 354 |
+
|
| 355 |
+
def _sync_env_state(src_env: PickClutterRevealEnv, dst_env: PickClutterRevealEnv) -> None:
|
| 356 |
+
_restore_env(dst_env, _snapshot_env(src_env))
|
| 357 |
+
|
| 358 |
+
|
| 359 |
+
def _canonical_chunks() -> dict[str, np.ndarray]:
|
| 360 |
+
base = torch.zeros((1, 8, 14), dtype=torch.float32)
|
| 361 |
+
chunks: dict[str, np.ndarray] = {"base_action": base.squeeze(0).numpy().astype(np.float32)}
|
| 362 |
+
for mode_name in MODE_ORDER[1:]:
|
| 363 |
+
chunk = semantic_macro_chunk(base, task_name=TASK_NAME, mode_name=mode_name).squeeze(0).cpu().numpy()
|
| 364 |
+
chunks[mode_name] = chunk.astype(np.float32)
|
| 365 |
+
return chunks
|
| 366 |
+
|
| 367 |
+
|
| 368 |
+
CANONICAL_CHUNKS = _canonical_chunks()
|
| 369 |
+
STATE_METRIC_MASK = np.asarray(
|
| 370 |
+
[metric_name in STATE_SUPERVISION_METRICS for metric_name in TASK_METRIC_NAMES],
|
| 371 |
+
dtype=np.bool_,
|
| 372 |
+
)
|
| 373 |
+
|
| 374 |
+
|
| 375 |
+
def _classify_mode_from_chunk(chunk: np.ndarray) -> str:
|
| 376 |
+
candidate = np.asarray(chunk, dtype=np.float32)
|
| 377 |
+
distances = {
|
| 378 |
+
mode_name: float(np.mean(np.abs(candidate - prototype)))
|
| 379 |
+
for mode_name, prototype in CANONICAL_CHUNKS.items()
|
| 380 |
+
}
|
| 381 |
+
return min(distances, key=distances.get)
|
| 382 |
+
|
| 383 |
+
|
| 384 |
+
def _gripper_action(open_gripper: bool) -> float:
|
| 385 |
+
return 1.0 if open_gripper else -1.0
|
| 386 |
+
|
| 387 |
+
|
| 388 |
+
def _repeat_delta(env: PickClutterRevealEnv, delta_xyz: Sequence[float], *, open_gripper: bool, steps: int) -> dict[str, Any]:
|
| 389 |
+
last_obs: dict[str, Any] | None = None
|
| 390 |
+
action = np.zeros((1, 4), dtype=np.float32)
|
| 391 |
+
action[0, :3] = np.asarray(delta_xyz, dtype=np.float32)
|
| 392 |
+
action[0, 3] = _gripper_action(open_gripper)
|
| 393 |
+
for _ in range(int(steps)):
|
| 394 |
+
obs, _, terminated, truncated, info = env.step(action)
|
| 395 |
+
last_obs = obs
|
| 396 |
+
if bool(np.asarray(terminated).reshape(-1)[0]) or bool(np.asarray(truncated).reshape(-1)[0]):
|
| 397 |
+
break
|
| 398 |
+
return {
|
| 399 |
+
"obs": last_obs if last_obs is not None else env.get_obs(env.get_info()),
|
| 400 |
+
"terminated": False,
|
| 401 |
+
"truncated": False,
|
| 402 |
+
"info": info if last_obs is not None else env.get_info(),
|
| 403 |
+
}
|
| 404 |
+
|
| 405 |
+
|
| 406 |
+
def _move_tcp_to(
|
| 407 |
+
env: PickClutterRevealEnv,
|
| 408 |
+
target_xyz: Sequence[float],
|
| 409 |
+
*,
|
| 410 |
+
open_gripper: bool,
|
| 411 |
+
max_steps: int = 120,
|
| 412 |
+
tolerance: float = 0.008,
|
| 413 |
+
) -> dict[str, Any]:
|
| 414 |
+
last_obs: dict[str, Any] | None = None
|
| 415 |
+
target = np.asarray(target_xyz, dtype=np.float32)
|
| 416 |
+
info = env.get_info()
|
| 417 |
+
for _ in range(int(max_steps)):
|
| 418 |
+
tcp = _vec3(env.unwrapped.agent.tcp.pose.p)
|
| 419 |
+
delta = target - tcp
|
| 420 |
+
if float(np.linalg.norm(delta)) <= float(tolerance):
|
| 421 |
+
break
|
| 422 |
+
action = np.zeros((1, 4), dtype=np.float32)
|
| 423 |
+
action[0, :3] = np.clip(delta / 0.04, -1.0, 1.0)
|
| 424 |
+
action[0, 3] = _gripper_action(open_gripper)
|
| 425 |
+
obs, _, terminated, truncated, info = env.step(action)
|
| 426 |
+
last_obs = obs
|
| 427 |
+
if bool(np.asarray(terminated).reshape(-1)[0]) or bool(np.asarray(truncated).reshape(-1)[0]):
|
| 428 |
+
break
|
| 429 |
+
return {
|
| 430 |
+
"obs": last_obs if last_obs is not None else env.get_obs(info),
|
| 431 |
+
"info": info,
|
| 432 |
+
}
|
| 433 |
+
|
| 434 |
+
|
| 435 |
+
def _find_path_blocker(env: PickClutterRevealEnv) -> np.ndarray | None:
|
| 436 |
+
target = _target_position(env)
|
| 437 |
+
target_name = _target_actor(env).name
|
| 438 |
+
blockers: list[tuple[float, np.ndarray]] = []
|
| 439 |
+
for actor in _all_scene_actors(env):
|
| 440 |
+
if actor.name == target_name:
|
| 441 |
+
continue
|
| 442 |
+
position = _vec3(actor.pose.p)
|
| 443 |
+
if position[1] <= target[1] + 0.06 and abs(position[0] - target[0]) <= 0.10:
|
| 444 |
+
blockers.append((float(np.linalg.norm(position[:2] - target[:2])), position))
|
| 445 |
+
if blockers:
|
| 446 |
+
blockers.sort(key=lambda item: item[0])
|
| 447 |
+
return blockers[0][1]
|
| 448 |
+
nearest: tuple[float, np.ndarray] | None = None
|
| 449 |
+
for actor in _all_scene_actors(env):
|
| 450 |
+
if actor.name == target_name:
|
| 451 |
+
continue
|
| 452 |
+
position = _vec3(actor.pose.p)
|
| 453 |
+
distance = float(np.linalg.norm(position[:2] - target[:2]))
|
| 454 |
+
if nearest is None or distance < nearest[0]:
|
| 455 |
+
nearest = (distance, position)
|
| 456 |
+
return None if nearest is None else nearest[1]
|
| 457 |
+
|
| 458 |
+
|
| 459 |
+
def _execute_push(
|
| 460 |
+
env: PickClutterRevealEnv,
|
| 461 |
+
*,
|
| 462 |
+
anchor_xyz: np.ndarray,
|
| 463 |
+
pre_offset: np.ndarray,
|
| 464 |
+
push_delta: np.ndarray,
|
| 465 |
+
push_steps: int,
|
| 466 |
+
) -> dict[str, Any]:
|
| 467 |
+
_move_tcp_to(env, anchor_xyz + np.array([0.0, 0.0, 0.08], dtype=np.float32) + pre_offset, open_gripper=True)
|
| 468 |
+
_move_tcp_to(env, anchor_xyz + pre_offset, open_gripper=True, max_steps=100, tolerance=0.010)
|
| 469 |
+
result = _repeat_delta(env, push_delta, open_gripper=True, steps=push_steps)
|
| 470 |
+
_move_tcp_to(
|
| 471 |
+
env,
|
| 472 |
+
np.array([_vec3(env.unwrapped.agent.tcp.pose.p)[0], _vec3(env.unwrapped.agent.tcp.pose.p)[1], 0.10], dtype=np.float32),
|
| 473 |
+
open_gripper=True,
|
| 474 |
+
max_steps=80,
|
| 475 |
+
tolerance=0.012,
|
| 476 |
+
)
|
| 477 |
+
return result
|
| 478 |
+
|
| 479 |
+
|
| 480 |
+
def _execute_mode(env: PickClutterRevealEnv, mode_name: str) -> dict[str, Any]:
|
| 481 |
+
target = _target_position(env)
|
| 482 |
+
blocker = _find_path_blocker(env)
|
| 483 |
+
if mode_name == "retrieve":
|
| 484 |
+
return _execute_push(
|
| 485 |
+
env,
|
| 486 |
+
anchor_xyz=target,
|
| 487 |
+
pre_offset=np.array([0.0, 0.035, 0.026], dtype=np.float32),
|
| 488 |
+
push_delta=np.array([0.0, -0.7, 0.0], dtype=np.float32),
|
| 489 |
+
push_steps=18,
|
| 490 |
+
)
|
| 491 |
+
if mode_name == "insert_actor":
|
| 492 |
+
return _execute_push(
|
| 493 |
+
env,
|
| 494 |
+
anchor_xyz=target,
|
| 495 |
+
pre_offset=np.array([0.0, 0.045, 0.028], dtype=np.float32),
|
| 496 |
+
push_delta=np.array([0.0, -0.4, 0.0], dtype=np.float32),
|
| 497 |
+
push_steps=10,
|
| 498 |
+
)
|
| 499 |
+
if mode_name == "widen_gap":
|
| 500 |
+
anchor = blocker if blocker is not None else target
|
| 501 |
+
direction = -1.0 if anchor[0] >= target[0] else 1.0
|
| 502 |
+
return _execute_push(
|
| 503 |
+
env,
|
| 504 |
+
anchor_xyz=anchor,
|
| 505 |
+
pre_offset=np.array([0.0, 0.025, 0.028], dtype=np.float32),
|
| 506 |
+
push_delta=np.array([0.75 * direction, -0.12, 0.0], dtype=np.float32),
|
| 507 |
+
push_steps=18,
|
| 508 |
+
)
|
| 509 |
+
if mode_name == "sweep_left":
|
| 510 |
+
anchor = blocker if blocker is not None else target
|
| 511 |
+
return _execute_push(
|
| 512 |
+
env,
|
| 513 |
+
anchor_xyz=anchor,
|
| 514 |
+
pre_offset=np.array([0.015, 0.025, 0.028], dtype=np.float32),
|
| 515 |
+
push_delta=np.array([-0.70, -0.10, 0.0], dtype=np.float32),
|
| 516 |
+
push_steps=14,
|
| 517 |
+
)
|
| 518 |
+
if mode_name == "sweep_right":
|
| 519 |
+
anchor = blocker if blocker is not None else target
|
| 520 |
+
return _execute_push(
|
| 521 |
+
env,
|
| 522 |
+
anchor_xyz=anchor,
|
| 523 |
+
pre_offset=np.array([-0.015, 0.025, 0.028], dtype=np.float32),
|
| 524 |
+
push_delta=np.array([0.70, -0.10, 0.0], dtype=np.float32),
|
| 525 |
+
push_steps=14,
|
| 526 |
+
)
|
| 527 |
+
if mode_name == "pin_canopy":
|
| 528 |
+
anchor = blocker if blocker is not None else target
|
| 529 |
+
return _execute_push(
|
| 530 |
+
env,
|
| 531 |
+
anchor_xyz=anchor,
|
| 532 |
+
pre_offset=np.array([0.0, -0.015, 0.028], dtype=np.float32),
|
| 533 |
+
push_delta=np.array([0.0, 0.35, 0.0], dtype=np.float32),
|
| 534 |
+
push_steps=10,
|
| 535 |
+
)
|
| 536 |
+
if mode_name in {"maintain_gap", "base_action"}:
|
| 537 |
+
_move_tcp_to(
|
| 538 |
+
env,
|
| 539 |
+
np.array([target[0], target[1] + 0.02, 0.10], dtype=np.float32),
|
| 540 |
+
open_gripper=True,
|
| 541 |
+
max_steps=60,
|
| 542 |
+
tolerance=0.015,
|
| 543 |
+
)
|
| 544 |
+
return _repeat_delta(env, np.array([0.0, -0.10, 0.0], dtype=np.float32), open_gripper=True, steps=4)
|
| 545 |
+
raise KeyError(f"Unsupported mode: {mode_name}")
|
| 546 |
+
|
| 547 |
+
|
| 548 |
+
def _candidate_metrics(
|
| 549 |
+
env: PickClutterRevealEnv,
|
| 550 |
+
*,
|
| 551 |
+
start_positions: dict[str, np.ndarray],
|
| 552 |
+
current_obs_bundle: dict[str, np.ndarray] | None = None,
|
| 553 |
+
) -> dict[str, float]:
|
| 554 |
+
positions = _all_positions(env)
|
| 555 |
+
target_name = _target_actor(env).name
|
| 556 |
+
non_target_displacements = []
|
| 557 |
+
for name, start_position in start_positions.items():
|
| 558 |
+
if name == target_name or name not in positions:
|
| 559 |
+
continue
|
| 560 |
+
non_target_displacements.append(float(np.linalg.norm((positions[name] - start_position)[:2])))
|
| 561 |
+
disturbance = float(np.clip(np.mean(non_target_displacements) / 0.10, 0.0, 1.0)) if non_target_displacements else 0.0
|
| 562 |
+
visibility = 0.0
|
| 563 |
+
if current_obs_bundle is not None:
|
| 564 |
+
visibility = _target_visibility(current_obs_bundle, getattr(_target_actor(env), "per_scene_id", -1))
|
| 565 |
+
return {
|
| 566 |
+
"retrieval_success": float(_success_from_state(env)),
|
| 567 |
+
"disturbance": disturbance,
|
| 568 |
+
"visibility": visibility,
|
| 569 |
+
"clearance": _clearance_score(env),
|
| 570 |
+
"progress": _extraction_progress(env),
|
| 571 |
+
}
|
| 572 |
+
|
| 573 |
+
|
| 574 |
+
def _mean_non_target_displacement(
|
| 575 |
+
start_positions: dict[str, np.ndarray],
|
| 576 |
+
current_positions: dict[str, np.ndarray],
|
| 577 |
+
*,
|
| 578 |
+
target_name: str,
|
| 579 |
+
) -> float:
|
| 580 |
+
displacements = []
|
| 581 |
+
for actor_name, start_position in start_positions.items():
|
| 582 |
+
if actor_name == target_name or actor_name not in current_positions:
|
| 583 |
+
continue
|
| 584 |
+
displacements.append(float(np.linalg.norm((current_positions[actor_name] - start_position)[:2])))
|
| 585 |
+
if not displacements:
|
| 586 |
+
return 0.0
|
| 587 |
+
return float(np.mean(displacements))
|
| 588 |
+
|
| 589 |
+
|
| 590 |
+
def _current_state_targets(
|
| 591 |
+
env: PickClutterRevealEnv,
|
| 592 |
+
*,
|
| 593 |
+
obs_bundle: dict[str, np.ndarray],
|
| 594 |
+
candidate_metrics: Sequence[dict[str, float]],
|
| 595 |
+
episode_start_positions: dict[str, np.ndarray],
|
| 596 |
+
selected_mode: str,
|
| 597 |
+
) -> dict[str, Any]:
|
| 598 |
+
metrics_by_name = {mode_name: payload for mode_name, payload in zip(MODE_ORDER, candidate_metrics)}
|
| 599 |
+
current_positions = _all_positions(env)
|
| 600 |
+
target_name = _target_actor(env).name
|
| 601 |
+
current_disturbance = float(
|
| 602 |
+
np.clip(
|
| 603 |
+
_mean_non_target_displacement(
|
| 604 |
+
episode_start_positions,
|
| 605 |
+
current_positions,
|
| 606 |
+
target_name=target_name,
|
| 607 |
+
)
|
| 608 |
+
/ 0.10,
|
| 609 |
+
0.0,
|
| 610 |
+
1.0,
|
| 611 |
+
)
|
| 612 |
+
)
|
| 613 |
+
current_visibility = _target_visibility(obs_bundle, getattr(_target_actor(env), "per_scene_id", -1))
|
| 614 |
+
current_clearance = _clearance_score(env)
|
| 615 |
+
current_progress = _extraction_progress(env)
|
| 616 |
+
base_gap = float(np.clip(max(current_clearance, current_progress), 0.0, 1.0))
|
| 617 |
+
support_stability = float(np.clip(1.0 - 0.5 * current_disturbance, 0.0, 1.0))
|
| 618 |
+
hold_quality = float(np.clip(0.5 * (support_stability + max(current_clearance, current_progress)), 0.0, 1.0))
|
| 619 |
+
opening_quality = float(
|
| 620 |
+
np.clip(0.55 * current_progress + 0.25 * current_clearance + 0.20 * current_visibility, 0.0, 1.0)
|
| 621 |
+
)
|
| 622 |
+
actor_feasibility = float(np.clip(0.6 * current_clearance + 0.4 * max(current_visibility, current_progress), 0.0, 1.0))
|
| 623 |
+
reocclusion_rate = float(np.clip(1.0 - max(current_clearance, current_visibility), 0.0, 1.0))
|
| 624 |
+
insertable_actor_corridor = float(np.clip(0.6 * actor_feasibility + 0.4 * base_gap, 0.0, 1.0))
|
| 625 |
+
insertion_corridor = float(np.clip(0.5 * actor_feasibility + 0.5 * base_gap, 0.0, 1.0))
|
| 626 |
+
layer_separation = float(np.clip(0.7 * base_gap + 0.3 * actor_feasibility, 0.0, 1.0))
|
| 627 |
+
fold_preservation = float(np.clip(1.0 - current_disturbance, 0.0, 1.0))
|
| 628 |
+
lift_too_much_risk = float(np.clip(current_disturbance + 0.5 * max(base_gap - 0.5, 0.0), 0.0, 1.0))
|
| 629 |
+
task_metrics = {
|
| 630 |
+
"opening_quality": opening_quality,
|
| 631 |
+
"actor_feasibility_score": actor_feasibility,
|
| 632 |
+
"gap_width": float(0.03 + 0.21 * base_gap),
|
| 633 |
+
"damage_proxy": current_disturbance,
|
| 634 |
+
"release_collapse_rate": reocclusion_rate,
|
| 635 |
+
"target_visibility_confidence": current_visibility,
|
| 636 |
+
"insertable_actor_corridor": insertable_actor_corridor,
|
| 637 |
+
"insertion_corridor": insertion_corridor,
|
| 638 |
+
"hold_quality": hold_quality,
|
| 639 |
+
"layer_separation_quality": layer_separation,
|
| 640 |
+
"fold_preservation": fold_preservation,
|
| 641 |
+
"top_layer_stability": support_stability,
|
| 642 |
+
"lift_too_much_risk": lift_too_much_risk,
|
| 643 |
+
}
|
| 644 |
+
|
| 645 |
+
base_metrics = metrics_by_name["base_action"]
|
| 646 |
+
insert_metrics = metrics_by_name["insert_actor"]
|
| 647 |
+
retrieve_metrics = metrics_by_name["retrieve"]
|
| 648 |
+
reveal_candidates = [metrics_by_name[mode_name] for mode_name in REVEAL_MODES]
|
| 649 |
+
reveal_access = max(candidate["candidate_actor_feasibility_auc"] for candidate in reveal_candidates)
|
| 650 |
+
reveal_reveal = max(candidate["candidate_reveal_achieved"] for candidate in reveal_candidates)
|
| 651 |
+
reveal_hold = max(candidate["candidate_hold_persistence"] for candidate in reveal_candidates)
|
| 652 |
+
reveal_visibility = max(candidate["candidate_visibility_integral"] for candidate in reveal_candidates)
|
| 653 |
+
|
| 654 |
+
reveal_corridor = float(
|
| 655 |
+
np.clip(
|
| 656 |
+
0.45 * opening_quality
|
| 657 |
+
+ 0.30 * reveal_access
|
| 658 |
+
+ 0.15 * reveal_reveal
|
| 659 |
+
+ 0.10 * reveal_visibility
|
| 660 |
+
- 0.10 * current_disturbance,
|
| 661 |
+
0.0,
|
| 662 |
+
1.0,
|
| 663 |
+
)
|
| 664 |
+
)
|
| 665 |
+
transfer_corridor = float(
|
| 666 |
+
np.clip(
|
| 667 |
+
0.45 * insertable_actor_corridor
|
| 668 |
+
+ 0.30 * insert_metrics["candidate_actor_feasibility_auc"]
|
| 669 |
+
+ 0.15 * insert_metrics["candidate_reveal_achieved"]
|
| 670 |
+
+ 0.10 * insert_metrics["candidate_visibility_integral"]
|
| 671 |
+
- 0.15 * current_disturbance,
|
| 672 |
+
0.0,
|
| 673 |
+
1.0,
|
| 674 |
+
)
|
| 675 |
+
)
|
| 676 |
+
passive_corridor = float(
|
| 677 |
+
np.clip(
|
| 678 |
+
0.55 * retrieve_metrics["candidate_retrieval_success"]
|
| 679 |
+
+ 0.20 * retrieve_metrics["candidate_actor_feasibility_auc"]
|
| 680 |
+
+ 0.15 * current_progress
|
| 681 |
+
+ 0.10 * current_clearance
|
| 682 |
+
- 0.10 * current_disturbance,
|
| 683 |
+
0.0,
|
| 684 |
+
1.0,
|
| 685 |
+
)
|
| 686 |
+
)
|
| 687 |
+
corridor_feasible = np.stack(
|
| 688 |
+
[
|
| 689 |
+
np.full((NUM_APPROACH_TEMPLATES,), reveal_corridor, dtype=np.float32),
|
| 690 |
+
np.full((NUM_APPROACH_TEMPLATES,), transfer_corridor, dtype=np.float32),
|
| 691 |
+
np.full((NUM_APPROACH_TEMPLATES,), passive_corridor, dtype=np.float32),
|
| 692 |
+
],
|
| 693 |
+
axis=0,
|
| 694 |
+
)
|
| 695 |
+
persistence_horizon = np.asarray(
|
| 696 |
+
[
|
| 697 |
+
ROLL_OUT_HORIZON
|
| 698 |
+
* float(np.clip(0.35 * hold_quality + 0.35 * reveal_hold + 0.30 * reveal_corridor, 0.0, 1.0)),
|
| 699 |
+
ROLL_OUT_HORIZON
|
| 700 |
+
* float(
|
| 701 |
+
np.clip(
|
| 702 |
+
0.30 * hold_quality
|
| 703 |
+
+ 0.35 * insert_metrics["candidate_hold_persistence"]
|
| 704 |
+
+ 0.35 * transfer_corridor,
|
| 705 |
+
0.0,
|
| 706 |
+
1.0,
|
| 707 |
+
)
|
| 708 |
+
),
|
| 709 |
+
ROLL_OUT_HORIZON
|
| 710 |
+
* float(
|
| 711 |
+
np.clip(
|
| 712 |
+
0.25 * hold_quality
|
| 713 |
+
+ 0.35 * retrieve_metrics["candidate_hold_persistence"]
|
| 714 |
+
+ 0.40 * passive_corridor,
|
| 715 |
+
0.0,
|
| 716 |
+
1.0,
|
| 717 |
+
)
|
| 718 |
+
),
|
| 719 |
+
],
|
| 720 |
+
dtype=np.float32,
|
| 721 |
+
)
|
| 722 |
+
|
| 723 |
+
retrieve_margin = float(retrieve_metrics["candidate_utility"] - base_metrics["candidate_utility"])
|
| 724 |
+
insert_margin = float(insert_metrics["candidate_utility"] - base_metrics["candidate_utility"])
|
| 725 |
+
if selected_mode == "retrieve" or (retrieve_metrics["candidate_retrieval_success"] >= 0.5 and retrieve_margin >= 0.15):
|
| 726 |
+
support_mode = SUPPORT_MODE_PASSIVE
|
| 727 |
+
elif selected_mode == "insert_actor" or (insert_margin >= 0.15 and transfer_corridor >= 0.40):
|
| 728 |
+
support_mode = SUPPORT_MODE_TRANSFER
|
| 729 |
+
elif selected_mode in REVEAL_MODES or selected_mode == "maintain_gap":
|
| 730 |
+
support_mode = SUPPORT_MODE_HOLD
|
| 731 |
+
elif selected_mode == "base_action":
|
| 732 |
+
support_mode = SUPPORT_MODE_PASSIVE if passive_corridor >= 0.65 and retrieve_margin >= 0.05 else SUPPORT_MODE_HOLD
|
| 733 |
+
else:
|
| 734 |
+
support_mode = SUPPORT_MODE_HOLD
|
| 735 |
+
|
| 736 |
+
best_non_base_utility = max(float(payload["candidate_utility"]) for payload in candidate_metrics[1:])
|
| 737 |
+
intervention_warranted = selected_mode != "base_action" and best_non_base_utility >= float(base_metrics["candidate_utility"]) + 0.15
|
| 738 |
+
|
| 739 |
+
return {
|
| 740 |
+
"support_mode": int(support_mode),
|
| 741 |
+
"corridor_feasible": corridor_feasible,
|
| 742 |
+
"persistence_horizon": persistence_horizon,
|
| 743 |
+
"disturbance_cost": np.float32(current_disturbance),
|
| 744 |
+
"state_confidence_target": np.float32(1.0 if intervention_warranted else 0.0),
|
| 745 |
+
"task_metric_mask": STATE_METRIC_MASK.copy(),
|
| 746 |
+
**{metric_name: np.float32(metric_value) for metric_name, metric_value in task_metrics.items()},
|
| 747 |
+
}
|
| 748 |
+
|
| 749 |
+
|
| 750 |
+
def _mode_support_mode(mode_name: str, current_support_mode: int) -> int:
|
| 751 |
+
if mode_name in REVEAL_MODES or mode_name == "maintain_gap":
|
| 752 |
+
return SUPPORT_MODE_HOLD
|
| 753 |
+
if mode_name in TRANSFER_MODES:
|
| 754 |
+
return SUPPORT_MODE_TRANSFER
|
| 755 |
+
if mode_name in RETRIEVE_MODES:
|
| 756 |
+
return SUPPORT_MODE_PASSIVE
|
| 757 |
+
return int(current_support_mode)
|
| 758 |
+
|
| 759 |
+
|
| 760 |
+
def _mode_progress_schedule(mode_name: str) -> np.ndarray:
|
| 761 |
+
if mode_name in REVEAL_MODES:
|
| 762 |
+
return np.asarray([0.18, 0.38, 0.62, 0.84, 1.0], dtype=np.float32)
|
| 763 |
+
if mode_name in TRANSFER_MODES:
|
| 764 |
+
return np.asarray([0.22, 0.44, 0.66, 0.86, 1.0], dtype=np.float32)
|
| 765 |
+
if mode_name in RETRIEVE_MODES:
|
| 766 |
+
return np.asarray([0.34, 0.56, 0.76, 0.92, 1.0], dtype=np.float32)
|
| 767 |
+
return np.asarray([0.10, 0.22, 0.34, 0.44, 0.54], dtype=np.float32)
|
| 768 |
+
|
| 769 |
+
|
| 770 |
+
def _scalar_rollout(start: float, end: float, schedule: np.ndarray) -> np.ndarray:
|
| 771 |
+
return np.clip((1.0 - schedule) * float(start) + schedule * float(end), 0.0, 1.0).astype(np.float32)
|
| 772 |
+
|
| 773 |
+
|
| 774 |
+
def _candidate_rollout_targets(
|
| 775 |
+
*,
|
| 776 |
+
mode_name: str,
|
| 777 |
+
state_targets: dict[str, Any],
|
| 778 |
+
candidate_payload: dict[str, float],
|
| 779 |
+
) -> dict[str, np.ndarray]:
|
| 780 |
+
schedule = _mode_progress_schedule(mode_name)
|
| 781 |
+
start_visibility = float(state_targets["target_visibility_confidence"])
|
| 782 |
+
start_access = float(state_targets["actor_feasibility_score"])
|
| 783 |
+
start_persistence = float(np.clip(state_targets["hold_quality"], 0.0, 1.0))
|
| 784 |
+
start_support = float(np.clip(state_targets["top_layer_stability"], 0.0, 1.0))
|
| 785 |
+
start_reocclusion = float(np.clip(state_targets["release_collapse_rate"], 0.0, 1.0))
|
| 786 |
+
start_disturbance = float(np.clip(state_targets["disturbance_cost"], 0.0, 1.0))
|
| 787 |
+
start_clearance = float(np.clip(state_targets["actor_feasibility_score"], 0.0, 1.0))
|
| 788 |
+
start_grasp = float(np.clip(max(start_visibility, start_access), 0.0, 1.0))
|
| 789 |
+
|
| 790 |
+
end_visibility = float(np.clip(candidate_payload["candidate_immediate_visibility"], 0.0, 1.0))
|
| 791 |
+
end_access = float(np.clip(candidate_payload["candidate_immediate_access"], 0.0, 1.0))
|
| 792 |
+
end_progress = float(np.clip(candidate_payload["candidate_immediate_progress"], 0.0, 1.0))
|
| 793 |
+
end_disturbance = float(np.clip(candidate_payload["candidate_immediate_disturbance"], 0.0, 1.0))
|
| 794 |
+
end_support = float(np.clip(candidate_payload["candidate_immediate_support_stability"], 0.0, 1.0))
|
| 795 |
+
end_persistence = float(np.clip(candidate_payload["candidate_immediate_hold_persistence"], 0.0, 1.0))
|
| 796 |
+
end_reocclusion = float(np.clip(candidate_payload["candidate_immediate_reocclusion"], 0.0, 1.0))
|
| 797 |
+
end_clearance = float(np.clip(max(end_access, end_progress), 0.0, 1.0))
|
| 798 |
+
end_grasp = float(np.clip(max(end_visibility, 0.5 * end_access + 0.5 * end_progress), 0.0, 1.0))
|
| 799 |
+
|
| 800 |
+
if mode_name in TRANSFER_MODES:
|
| 801 |
+
start_visibility = max(start_visibility, 0.35 * end_visibility)
|
| 802 |
+
start_access = max(start_access, 0.40 * end_access)
|
| 803 |
+
start_persistence = max(start_persistence, 0.45 * end_persistence)
|
| 804 |
+
start_support = max(start_support, 0.50 * end_support)
|
| 805 |
+
elif mode_name in RETRIEVE_MODES:
|
| 806 |
+
start_visibility = max(start_visibility, 0.55 * end_visibility)
|
| 807 |
+
start_access = max(start_access, 0.70 * end_access)
|
| 808 |
+
start_persistence = max(start_persistence, 0.65 * end_persistence)
|
| 809 |
+
start_support = max(start_support, 0.65 * end_support)
|
| 810 |
+
start_reocclusion = min(start_reocclusion, max(0.4 * end_reocclusion, 0.0))
|
| 811 |
+
|
| 812 |
+
visibility = _scalar_rollout(start_visibility, end_visibility, schedule)
|
| 813 |
+
access = _scalar_rollout(start_access, end_access, schedule)
|
| 814 |
+
persistence = _scalar_rollout(start_persistence, end_persistence, schedule)
|
| 815 |
+
support = _scalar_rollout(start_support, end_support, schedule)
|
| 816 |
+
reocclusion = _scalar_rollout(start_reocclusion, end_reocclusion, schedule)
|
| 817 |
+
disturbance = _scalar_rollout(start_disturbance, end_disturbance, schedule)
|
| 818 |
+
clearance = _scalar_rollout(start_clearance, end_clearance, schedule)
|
| 819 |
+
grasp = _scalar_rollout(start_grasp, end_grasp, schedule)
|
| 820 |
+
|
| 821 |
+
reveal_corridor = np.clip(
|
| 822 |
+
0.38 * visibility + 0.34 * access + 0.22 * support - 0.12 * disturbance,
|
| 823 |
+
0.0,
|
| 824 |
+
1.0,
|
| 825 |
+
)
|
| 826 |
+
transfer_corridor = np.clip(
|
| 827 |
+
0.30 * visibility + 0.38 * access + 0.18 * persistence + 0.14 * support - 0.12 * disturbance,
|
| 828 |
+
0.0,
|
| 829 |
+
1.0,
|
| 830 |
+
)
|
| 831 |
+
passive_corridor = np.clip(
|
| 832 |
+
0.22 * visibility + 0.42 * access + 0.20 * persistence + 0.16 * grasp - 0.14 * disturbance - 0.10 * reocclusion,
|
| 833 |
+
0.0,
|
| 834 |
+
1.0,
|
| 835 |
+
)
|
| 836 |
+
if mode_name in REVEAL_MODES:
|
| 837 |
+
reveal_corridor = np.clip(reveal_corridor + 0.14, 0.0, 1.0)
|
| 838 |
+
passive_corridor = np.clip(0.75 * passive_corridor, 0.0, 1.0)
|
| 839 |
+
elif mode_name in TRANSFER_MODES:
|
| 840 |
+
transfer_corridor = np.clip(transfer_corridor + 0.16, 0.0, 1.0)
|
| 841 |
+
elif mode_name in RETRIEVE_MODES:
|
| 842 |
+
passive_corridor = np.clip(passive_corridor + 0.20, 0.0, 1.0)
|
| 843 |
+
reveal_corridor = np.clip(0.60 * reveal_corridor, 0.0, 1.0)
|
| 844 |
+
else:
|
| 845 |
+
reveal_corridor = np.clip(0.85 * reveal_corridor, 0.0, 1.0)
|
| 846 |
+
transfer_corridor = np.clip(0.75 * transfer_corridor, 0.0, 1.0)
|
| 847 |
+
passive_corridor = np.clip(0.80 * passive_corridor, 0.0, 1.0)
|
| 848 |
+
|
| 849 |
+
corridor_feasible = np.stack(
|
| 850 |
+
[
|
| 851 |
+
np.repeat(reveal_corridor[:, None], NUM_APPROACH_TEMPLATES, axis=1),
|
| 852 |
+
np.repeat(transfer_corridor[:, None], NUM_APPROACH_TEMPLATES, axis=1),
|
| 853 |
+
np.repeat(passive_corridor[:, None], NUM_APPROACH_TEMPLATES, axis=1),
|
| 854 |
+
],
|
| 855 |
+
axis=1,
|
| 856 |
+
).astype(np.float32)
|
| 857 |
+
persistence_horizon = np.stack(
|
| 858 |
+
[
|
| 859 |
+
np.clip(ROLL_OUT_HORIZON * (0.55 * reveal_corridor + 0.45 * support), 0.0, float(ROLL_OUT_HORIZON)),
|
| 860 |
+
np.clip(ROLL_OUT_HORIZON * (0.50 * transfer_corridor + 0.50 * persistence), 0.0, float(ROLL_OUT_HORIZON)),
|
| 861 |
+
np.clip(ROLL_OUT_HORIZON * (0.55 * passive_corridor + 0.45 * persistence), 0.0, float(ROLL_OUT_HORIZON)),
|
| 862 |
+
],
|
| 863 |
+
axis=1,
|
| 864 |
+
).astype(np.float32)
|
| 865 |
+
support_mode = np.full(
|
| 866 |
+
(ROLL_OUT_HORIZON,),
|
| 867 |
+
_mode_support_mode(mode_name, int(state_targets["support_mode"])),
|
| 868 |
+
dtype=np.int64,
|
| 869 |
+
)
|
| 870 |
+
if mode_name == "base_action":
|
| 871 |
+
support_mode[:] = int(state_targets["support_mode"])
|
| 872 |
+
|
| 873 |
+
return {
|
| 874 |
+
"candidate_rollout_support_mode": support_mode,
|
| 875 |
+
"candidate_rollout_corridor_feasible": corridor_feasible,
|
| 876 |
+
"candidate_rollout_persistence_horizon": persistence_horizon,
|
| 877 |
+
"candidate_rollout_disturbance_cost": disturbance.astype(np.float32),
|
| 878 |
+
"candidate_rollout_belief_map": visibility[:, None, None].astype(np.float32),
|
| 879 |
+
"candidate_rollout_visibility_map": visibility[:, None, None].astype(np.float32),
|
| 880 |
+
"candidate_rollout_clearance_map": np.repeat(clearance[:, None, None, None], 2, axis=1).astype(np.float32),
|
| 881 |
+
"candidate_rollout_support_stability": support[:, None, None, None].astype(np.float32),
|
| 882 |
+
"candidate_rollout_reocclusion_target": reocclusion[:, None, None].astype(np.float32),
|
| 883 |
+
"candidate_rollout_occluder_contact_map": np.clip(access * support, 0.0, 1.0)[:, None, None].astype(np.float32),
|
| 884 |
+
"candidate_rollout_grasp_affordance_map": grasp[:, None, None].astype(np.float32),
|
| 885 |
+
}
|
| 886 |
+
|
| 887 |
+
|
| 888 |
+
def _evaluate_candidate(
|
| 889 |
+
sim_env: PickClutterRevealEnv,
|
| 890 |
+
obs_env: PickClutterRevealEnv,
|
| 891 |
+
snapshot: dict[str, Any],
|
| 892 |
+
mode_name: str,
|
| 893 |
+
) -> dict[str, float]:
|
| 894 |
+
_restore_env(sim_env, snapshot)
|
| 895 |
+
start_positions = _all_positions(sim_env)
|
| 896 |
+
_execute_mode(sim_env, mode_name)
|
| 897 |
+
_sync_env_state(sim_env, obs_env)
|
| 898 |
+
after_bundle = _extract_sensor_bundle(obs_env.get_obs(obs_env.get_info()))
|
| 899 |
+
immediate = _candidate_metrics(sim_env, start_positions=start_positions, current_obs_bundle=after_bundle)
|
| 900 |
+
if not immediate["retrieval_success"] and mode_name != "retrieve":
|
| 901 |
+
_execute_mode(sim_env, "retrieve")
|
| 902 |
+
_sync_env_state(sim_env, obs_env)
|
| 903 |
+
follow_bundle = _extract_sensor_bundle(obs_env.get_obs(obs_env.get_info()))
|
| 904 |
+
final_metrics = _candidate_metrics(sim_env, start_positions=start_positions, current_obs_bundle=follow_bundle)
|
| 905 |
+
else:
|
| 906 |
+
final_metrics = immediate
|
| 907 |
+
_restore_env(obs_env, snapshot)
|
| 908 |
+
utility = (
|
| 909 |
+
2.5 * final_metrics["retrieval_success"]
|
| 910 |
+
+ 1.0 * final_metrics["progress"]
|
| 911 |
+
+ 0.5 * final_metrics["clearance"]
|
| 912 |
+
+ 0.25 * final_metrics["visibility"]
|
| 913 |
+
- 0.5 * final_metrics["disturbance"]
|
| 914 |
+
)
|
| 915 |
+
return {
|
| 916 |
+
"candidate_retrieval_success": final_metrics["retrieval_success"],
|
| 917 |
+
"candidate_risk": float(np.clip(final_metrics["disturbance"], 0.0, 1.0)),
|
| 918 |
+
"candidate_utility": float(utility),
|
| 919 |
+
"candidate_final_disturbance_cost": final_metrics["disturbance"],
|
| 920 |
+
"candidate_reocclusion_rate": float(np.clip(1.0 - final_metrics["clearance"], 0.0, 1.0)),
|
| 921 |
+
"candidate_visibility_integral": final_metrics["visibility"],
|
| 922 |
+
"candidate_actor_feasibility_auc": final_metrics["clearance"],
|
| 923 |
+
"candidate_reveal_achieved": float(final_metrics["progress"] > 0.15 or final_metrics["clearance"] > 0.35),
|
| 924 |
+
"candidate_hold_persistence": float(1.0 - final_metrics["disturbance"]),
|
| 925 |
+
"candidate_support_stability_auc": float(1.0 - 0.5 * final_metrics["disturbance"]),
|
| 926 |
+
"candidate_disturbance_auc": final_metrics["disturbance"],
|
| 927 |
+
"candidate_immediate_retrieval_success": immediate["retrieval_success"],
|
| 928 |
+
"candidate_immediate_visibility": immediate["visibility"],
|
| 929 |
+
"candidate_immediate_access": immediate["clearance"],
|
| 930 |
+
"candidate_immediate_progress": immediate["progress"],
|
| 931 |
+
"candidate_immediate_reocclusion": float(np.clip(1.0 - immediate["clearance"], 0.0, 1.0)),
|
| 932 |
+
"candidate_immediate_hold_persistence": float(1.0 - immediate["disturbance"]),
|
| 933 |
+
"candidate_immediate_support_stability": float(1.0 - 0.5 * immediate["disturbance"]),
|
| 934 |
+
"candidate_immediate_disturbance": immediate["disturbance"],
|
| 935 |
+
}
|
| 936 |
+
|
| 937 |
+
|
| 938 |
+
def _build_episode_splits(spec: SmokeSpec) -> dict[str, list[int]]:
|
| 939 |
+
return {
|
| 940 |
+
"train": [spec.dataset_seed * 10_000 + index for index in range(spec.train_episodes)],
|
| 941 |
+
"val": [spec.dataset_seed * 10_000 + 1_000 + index for index in range(spec.val_episodes)],
|
| 942 |
+
"eval": [spec.dataset_seed * 10_000 + 2_000 + index for index in range(spec.eval_episodes)],
|
| 943 |
+
}
|
| 944 |
+
|
| 945 |
+
|
| 946 |
+
def _save_episode_splits(output_path: Path, payload: dict[str, list[int]]) -> None:
|
| 947 |
+
output_path.parent.mkdir(parents=True, exist_ok=True)
|
| 948 |
+
output_path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")
|
| 949 |
+
|
| 950 |
+
|
| 951 |
+
def _init_history_entry(obs_bundle: dict[str, np.ndarray], proprio: np.ndarray, action_chunk: np.ndarray) -> dict[str, Any]:
|
| 952 |
+
return {
|
| 953 |
+
"images": obs_bundle["images"].copy(),
|
| 954 |
+
"depths": obs_bundle["depths"].copy(),
|
| 955 |
+
"depth_valid": obs_bundle["depth_valid"].copy(),
|
| 956 |
+
"camera_intrinsics": obs_bundle["camera_intrinsics"].copy(),
|
| 957 |
+
"camera_extrinsics": obs_bundle["camera_extrinsics"].copy(),
|
| 958 |
+
"proprio": proprio.copy(),
|
| 959 |
+
"action": action_chunk.mean(axis=0).astype(np.float32, copy=False),
|
| 960 |
+
}
|
| 961 |
+
|
| 962 |
+
|
| 963 |
+
def _history_stack(
|
| 964 |
+
history: Sequence[dict[str, Any]],
|
| 965 |
+
key: str,
|
| 966 |
+
*,
|
| 967 |
+
pad_shape: tuple[int, ...],
|
| 968 |
+
dtype: np.dtype,
|
| 969 |
+
history_steps: int,
|
| 970 |
+
) -> np.ndarray:
|
| 971 |
+
history = list(history)[-history_steps:]
|
| 972 |
+
pad_count = history_steps - len(history)
|
| 973 |
+
chunks = [np.zeros(pad_shape, dtype=dtype) for _ in range(pad_count)]
|
| 974 |
+
chunks.extend(np.asarray(item[key], dtype=dtype) for item in history)
|
| 975 |
+
return np.stack(chunks, axis=0).astype(dtype, copy=False)
|
| 976 |
+
|
| 977 |
+
|
| 978 |
+
class ManiSkillPickClutterDataset(Dataset[dict[str, Any]]):
|
| 979 |
+
def __init__(self, samples: Sequence[dict[str, Any]]) -> None:
|
| 980 |
+
self.samples = list(samples)
|
| 981 |
+
|
| 982 |
+
def __len__(self) -> int:
|
| 983 |
+
return len(self.samples)
|
| 984 |
+
|
| 985 |
+
def __getitem__(self, index: int) -> dict[str, Any]:
|
| 986 |
+
sample = self.samples[index]
|
| 987 |
+
item = {
|
| 988 |
+
"images": torch.from_numpy(sample["images"]).permute(0, 3, 1, 2).float() / 255.0,
|
| 989 |
+
"depths": torch.from_numpy(sample["depths"]).float(),
|
| 990 |
+
"depth_valid": torch.from_numpy(sample["depth_valid"]).float(),
|
| 991 |
+
"camera_intrinsics": torch.from_numpy(sample["camera_intrinsics"]).float(),
|
| 992 |
+
"camera_extrinsics": torch.from_numpy(sample["camera_extrinsics"]).float(),
|
| 993 |
+
"history_images": torch.from_numpy(sample["history_images"]).permute(0, 1, 4, 2, 3).float() / 255.0,
|
| 994 |
+
"history_depths": torch.from_numpy(sample["history_depths"]).float(),
|
| 995 |
+
"history_depth_valid": torch.from_numpy(sample["history_depth_valid"]).float(),
|
| 996 |
+
"history_camera_intrinsics": torch.from_numpy(sample["history_camera_intrinsics"]).float(),
|
| 997 |
+
"history_camera_extrinsics": torch.from_numpy(sample["history_camera_extrinsics"]).float(),
|
| 998 |
+
"history_proprio": torch.from_numpy(sample["history_proprio"]).float(),
|
| 999 |
+
"history_actions": torch.from_numpy(sample["history_actions"]).float(),
|
| 1000 |
+
"proprio": torch.from_numpy(sample["proprio"]).float(),
|
| 1001 |
+
"texts": sample["language_goal"],
|
| 1002 |
+
"task_name": sample["task_name"],
|
| 1003 |
+
"task_id": torch.as_tensor(sample["task_id"], dtype=torch.long),
|
| 1004 |
+
"action_chunk": torch.from_numpy(sample["action_chunk"]).float(),
|
| 1005 |
+
"candidate_action_chunks": torch.from_numpy(sample["candidate_action_chunks"]).float(),
|
| 1006 |
+
"candidate_retrieval_success": torch.from_numpy(sample["candidate_retrieval_success"]).float(),
|
| 1007 |
+
"candidate_final_disturbance_cost": torch.from_numpy(sample["candidate_final_disturbance_cost"]).float(),
|
| 1008 |
+
"candidate_reocclusion_rate": torch.from_numpy(sample["candidate_reocclusion_rate"]).float(),
|
| 1009 |
+
"candidate_visibility_integral": torch.from_numpy(sample["candidate_visibility_integral"]).float(),
|
| 1010 |
+
"candidate_actor_feasibility_auc": torch.from_numpy(sample["candidate_actor_feasibility_auc"]).float(),
|
| 1011 |
+
"candidate_reveal_achieved": torch.from_numpy(sample["candidate_reveal_achieved"]).float(),
|
| 1012 |
+
"candidate_hold_persistence": torch.from_numpy(sample["candidate_hold_persistence"]).float(),
|
| 1013 |
+
"candidate_support_stability_auc": torch.from_numpy(sample["candidate_support_stability_auc"]).float(),
|
| 1014 |
+
"candidate_disturbance_auc": torch.from_numpy(sample["candidate_disturbance_auc"]).float(),
|
| 1015 |
+
"candidate_risk": torch.from_numpy(sample["candidate_risk"]).float(),
|
| 1016 |
+
"candidate_utility": torch.from_numpy(sample["candidate_utility"]).float(),
|
| 1017 |
+
"proposal_target_action_chunks": torch.from_numpy(sample["candidate_action_chunks"]).float(),
|
| 1018 |
+
"proposal_target_retrieval_success": torch.from_numpy(sample["candidate_retrieval_success"]).float(),
|
| 1019 |
+
"proposal_target_risk": torch.from_numpy(sample["candidate_risk"]).float(),
|
| 1020 |
+
"proposal_target_utility": torch.from_numpy(sample["candidate_utility"]).float(),
|
| 1021 |
+
"episode_seed": sample["episode_seed"],
|
| 1022 |
+
"decision_step": sample["decision_step"],
|
| 1023 |
+
"selected_mode": sample["selected_mode"],
|
| 1024 |
+
}
|
| 1025 |
+
if "support_mode" in sample:
|
| 1026 |
+
item["support_mode"] = torch.as_tensor(sample["support_mode"], dtype=torch.long)
|
| 1027 |
+
if "corridor_feasible" in sample:
|
| 1028 |
+
item["corridor_feasible"] = torch.from_numpy(sample["corridor_feasible"]).float()
|
| 1029 |
+
if "persistence_horizon" in sample:
|
| 1030 |
+
item["persistence_horizon"] = torch.from_numpy(sample["persistence_horizon"]).float()
|
| 1031 |
+
if "disturbance_cost" in sample:
|
| 1032 |
+
item["disturbance_cost"] = torch.as_tensor(sample["disturbance_cost"], dtype=torch.float32)
|
| 1033 |
+
if "state_confidence_target" in sample:
|
| 1034 |
+
item["state_confidence_target"] = torch.as_tensor(sample["state_confidence_target"], dtype=torch.float32)
|
| 1035 |
+
if "task_metric_mask" in sample:
|
| 1036 |
+
item["task_metric_mask"] = torch.from_numpy(sample["task_metric_mask"]).to(dtype=torch.bool)
|
| 1037 |
+
for metric_name in STATE_SUPERVISION_METRICS:
|
| 1038 |
+
if metric_name in sample:
|
| 1039 |
+
item[metric_name] = torch.as_tensor(sample[metric_name], dtype=torch.float32)
|
| 1040 |
+
for key in (
|
| 1041 |
+
"candidate_rollout_support_mode",
|
| 1042 |
+
"proposal_target_rollout_support_mode",
|
| 1043 |
+
):
|
| 1044 |
+
if key in sample:
|
| 1045 |
+
item[key] = torch.from_numpy(sample[key]).long()
|
| 1046 |
+
for key in (
|
| 1047 |
+
"candidate_rollout_corridor_feasible",
|
| 1048 |
+
"candidate_rollout_persistence_horizon",
|
| 1049 |
+
"candidate_rollout_disturbance_cost",
|
| 1050 |
+
"candidate_rollout_belief_map",
|
| 1051 |
+
"candidate_rollout_visibility_map",
|
| 1052 |
+
"candidate_rollout_clearance_map",
|
| 1053 |
+
"candidate_rollout_support_stability",
|
| 1054 |
+
"candidate_rollout_reocclusion_target",
|
| 1055 |
+
"candidate_rollout_occluder_contact_map",
|
| 1056 |
+
"candidate_rollout_grasp_affordance_map",
|
| 1057 |
+
"proposal_target_rollout_corridor_feasible",
|
| 1058 |
+
"proposal_target_rollout_persistence_horizon",
|
| 1059 |
+
"proposal_target_rollout_disturbance_cost",
|
| 1060 |
+
"proposal_target_rollout_belief_map",
|
| 1061 |
+
"proposal_target_rollout_visibility_map",
|
| 1062 |
+
"proposal_target_rollout_clearance_map",
|
| 1063 |
+
"proposal_target_rollout_support_stability",
|
| 1064 |
+
"proposal_target_rollout_reocclusion_target",
|
| 1065 |
+
"proposal_target_rollout_occluder_contact_map",
|
| 1066 |
+
"proposal_target_rollout_grasp_affordance_map",
|
| 1067 |
+
):
|
| 1068 |
+
if key in sample:
|
| 1069 |
+
item[key] = torch.from_numpy(sample[key]).float()
|
| 1070 |
+
return item
|
| 1071 |
+
|
| 1072 |
+
|
| 1073 |
+
def _make_loader(samples: Sequence[dict[str, Any]], *, batch_size: int, num_workers: int, shuffle: bool) -> DataLoader:
|
| 1074 |
+
return DataLoader(
|
| 1075 |
+
ManiSkillPickClutterDataset(samples),
|
| 1076 |
+
batch_size=batch_size,
|
| 1077 |
+
shuffle=shuffle,
|
| 1078 |
+
num_workers=num_workers,
|
| 1079 |
+
pin_memory=torch.cuda.is_available(),
|
| 1080 |
+
)
|
| 1081 |
+
|
| 1082 |
+
|
| 1083 |
+
def _load_init_bundle() -> tuple[PolicyConfig, dict[str, Any], dict[str, Any]]:
|
| 1084 |
+
checkpoint = torch.load(DEFAULT_INIT_CHECKPOINT, map_location="cpu", weights_only=False)
|
| 1085 |
+
policy_config = PolicyConfig(
|
| 1086 |
+
backbone=FrozenVLBackboneConfig(**checkpoint["policy_config"]["backbone"]),
|
| 1087 |
+
fusion=MultiViewFusionConfig(**checkpoint["policy_config"]["fusion"]),
|
| 1088 |
+
memory=ObservationMemoryConfig(**checkpoint["policy_config"]["memory"]),
|
| 1089 |
+
decoder=ChunkDecoderConfig(**checkpoint["policy_config"]["decoder"]),
|
| 1090 |
+
reveal_head=RevealHeadConfig(**checkpoint["policy_config"]["reveal_head"]),
|
| 1091 |
+
world_model=RevealWMConfig(**checkpoint["policy_config"]["world_model"]),
|
| 1092 |
+
planner=PlannerConfig(**checkpoint["policy_config"]["planner"]),
|
| 1093 |
+
)
|
| 1094 |
+
return _apply_smoke_planner_overrides(policy_config), checkpoint["trainer_config"], checkpoint["loss_weights"]
|
| 1095 |
+
|
| 1096 |
+
|
| 1097 |
+
def _trainer_config_for_variant(variant: str) -> TrainerConfig:
|
| 1098 |
+
if variant == "trunk_only_ft":
|
| 1099 |
+
return TrainerConfig(
|
| 1100 |
+
policy_type="foundation_trunk",
|
| 1101 |
+
use_bf16=True,
|
| 1102 |
+
grad_clip_norm=1.0,
|
| 1103 |
+
freeze_backbone=True,
|
| 1104 |
+
gradient_checkpointing=False,
|
| 1105 |
+
trainable_parameter_prefixes=("fusion", "memory", "decoder"),
|
| 1106 |
+
)
|
| 1107 |
+
if variant == "adapter_active_ft":
|
| 1108 |
+
return TrainerConfig(
|
| 1109 |
+
policy_type="adapter_wrapped",
|
| 1110 |
+
use_bf16=True,
|
| 1111 |
+
grad_clip_norm=1.0,
|
| 1112 |
+
freeze_backbone=True,
|
| 1113 |
+
gradient_checkpointing=False,
|
| 1114 |
+
eval_mode="adapter_active",
|
| 1115 |
+
trainable_parameter_prefixes=(
|
| 1116 |
+
"trunk.fusion",
|
| 1117 |
+
"trunk.memory",
|
| 1118 |
+
"trunk.decoder",
|
| 1119 |
+
"adapter.state_head",
|
| 1120 |
+
"adapter.transition_model",
|
| 1121 |
+
"adapter.proposal_prior",
|
| 1122 |
+
"adapter.planner",
|
| 1123 |
+
),
|
| 1124 |
+
adapter_mode="adapter_active",
|
| 1125 |
+
adapter_use_transition_model=True,
|
| 1126 |
+
adapter_use_task_conditioning=True,
|
| 1127 |
+
adapter_action_supervision_source="trunk",
|
| 1128 |
+
)
|
| 1129 |
+
raise KeyError(f"Unsupported variant: {variant}")
|
| 1130 |
+
|
| 1131 |
+
|
| 1132 |
+
def _loss_weights_for_smoke() -> LossWeights:
|
| 1133 |
+
return LossWeights(
|
| 1134 |
+
action=1.0,
|
| 1135 |
+
support_mode=0.15,
|
| 1136 |
+
corridor=0.15,
|
| 1137 |
+
persistence=0.08,
|
| 1138 |
+
disturbance=0.08,
|
| 1139 |
+
planner_success=0.20,
|
| 1140 |
+
planner_risk=0.08,
|
| 1141 |
+
planner_ranking=0.20,
|
| 1142 |
+
proposal_reconstruction=0.10,
|
| 1143 |
+
proposal_success=0.12,
|
| 1144 |
+
proposal_ranking=0.15,
|
| 1145 |
+
proposal_mode=0.10,
|
| 1146 |
+
proposal_diversity=0.02,
|
| 1147 |
+
task_metrics=0.15,
|
| 1148 |
+
transition=0.25,
|
| 1149 |
+
gate=0.25,
|
| 1150 |
+
calibration=0.10,
|
| 1151 |
+
proposal_mode_task_filter=[TASK_NAME],
|
| 1152 |
+
)
|
| 1153 |
+
|
| 1154 |
+
|
| 1155 |
+
def _save_training_checkpoint(
|
| 1156 |
+
*,
|
| 1157 |
+
output_dir: Path,
|
| 1158 |
+
experiment_name: str,
|
| 1159 |
+
model: torch.nn.Module,
|
| 1160 |
+
policy_config: PolicyConfig,
|
| 1161 |
+
trainer_config: TrainerConfig,
|
| 1162 |
+
loss_weights: LossWeights,
|
| 1163 |
+
history: list[dict[str, Any]],
|
| 1164 |
+
best_val: float,
|
| 1165 |
+
train_spec: dict[str, Any],
|
| 1166 |
+
) -> Path:
|
| 1167 |
+
output_dir.mkdir(parents=True, exist_ok=True)
|
| 1168 |
+
checkpoint_path = output_dir / "checkpoint_best.pt"
|
| 1169 |
+
torch.save(
|
| 1170 |
+
{
|
| 1171 |
+
"experiment_name": experiment_name,
|
| 1172 |
+
"policy_config": asdict(policy_config),
|
| 1173 |
+
"trainer_config": asdict(trainer_config),
|
| 1174 |
+
"loss_weights": asdict(loss_weights),
|
| 1175 |
+
"state_dict": model.state_dict(),
|
| 1176 |
+
"history": history,
|
| 1177 |
+
"best_val_total": best_val,
|
| 1178 |
+
"train_spec": train_spec,
|
| 1179 |
+
},
|
| 1180 |
+
checkpoint_path,
|
| 1181 |
+
)
|
| 1182 |
+
return checkpoint_path
|
| 1183 |
+
|
| 1184 |
+
|
| 1185 |
+
def _aggregate_epoch(loss_records: Sequence[dict[str, float]]) -> dict[str, float]:
|
| 1186 |
+
keys = sorted({key for record in loss_records for key in record})
|
| 1187 |
+
return {
|
| 1188 |
+
key: float(np.mean([record.get(key, 0.0) for record in loss_records])) if loss_records else 0.0
|
| 1189 |
+
for key in keys
|
| 1190 |
+
}
|
| 1191 |
+
|
| 1192 |
+
|
| 1193 |
+
def _train_variant(
|
| 1194 |
+
*,
|
| 1195 |
+
variant: str,
|
| 1196 |
+
train_samples: Sequence[dict[str, Any]],
|
| 1197 |
+
val_samples: Sequence[dict[str, Any]],
|
| 1198 |
+
spec: SmokeSpec,
|
| 1199 |
+
output_dir: Path,
|
| 1200 |
+
) -> tuple[Path, dict[str, Any]]:
|
| 1201 |
+
policy_config, _init_trainer_cfg, _init_loss_weights = _load_init_bundle()
|
| 1202 |
+
trainer_config = _trainer_config_for_variant(variant)
|
| 1203 |
+
loss_weights = _loss_weights_for_smoke()
|
| 1204 |
+
model = build_policy(policy_config, trainer_config)
|
| 1205 |
+
init_info = _load_init_checkpoint(model, str(DEFAULT_INIT_CHECKPOINT), False)
|
| 1206 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 1207 |
+
model = model.to(device)
|
| 1208 |
+
torch.manual_seed(spec.train_seed)
|
| 1209 |
+
if torch.cuda.is_available():
|
| 1210 |
+
torch.cuda.manual_seed_all(spec.train_seed)
|
| 1211 |
+
torch.backends.cuda.matmul.allow_tf32 = True
|
| 1212 |
+
matched = apply_trainable_parameter_prefixes(model, trainer_config)
|
| 1213 |
+
optimizer = torch.optim.AdamW(
|
| 1214 |
+
[parameter for parameter in model.parameters() if parameter.requires_grad],
|
| 1215 |
+
lr=spec.learning_rate,
|
| 1216 |
+
weight_decay=spec.weight_decay,
|
| 1217 |
+
)
|
| 1218 |
+
trainer = BimanualTrainer(model=model, optimizer=optimizer, config=trainer_config)
|
| 1219 |
+
train_loader = _make_loader(train_samples, batch_size=spec.batch_size, num_workers=spec.num_workers, shuffle=True)
|
| 1220 |
+
val_loader = _make_loader(val_samples, batch_size=spec.batch_size, num_workers=spec.num_workers, shuffle=False)
|
| 1221 |
+
|
| 1222 |
+
best_val = math.inf
|
| 1223 |
+
history: list[dict[str, Any]] = []
|
| 1224 |
+
train_spec = build_target_training_spec(
|
| 1225 |
+
track_id="occlusion_track",
|
| 1226 |
+
model_variant=variant,
|
| 1227 |
+
seed=spec.train_seed,
|
| 1228 |
+
train_demos=spec.train_episodes,
|
| 1229 |
+
val_demos=spec.val_episodes,
|
| 1230 |
+
init_checkpoint_group=str(DEFAULT_INIT_CHECKPOINT),
|
| 1231 |
+
optimizer="adamw",
|
| 1232 |
+
learning_rate=spec.learning_rate,
|
| 1233 |
+
lr_schedule="constant",
|
| 1234 |
+
batch_size=spec.batch_size,
|
| 1235 |
+
augmentations="none",
|
| 1236 |
+
early_stopping_metric="val_total",
|
| 1237 |
+
max_gradient_steps=len(train_loader) * spec.epochs,
|
| 1238 |
+
unfreeze_scope="fusion_memory_decoder",
|
| 1239 |
+
dataset_split_id=(
|
| 1240 |
+
f"pickclutter_{SMOKE_VERSION}_seed{spec.dataset_seed}"
|
| 1241 |
+
if int(spec.dataset_seed) == DEFAULT_SEED
|
| 1242 |
+
else f"pickclutter_{SMOKE_VERSION}_dataset_seed{spec.dataset_seed}"
|
| 1243 |
+
),
|
| 1244 |
+
)
|
| 1245 |
+
|
| 1246 |
+
for epoch in range(spec.epochs):
|
| 1247 |
+
model.train()
|
| 1248 |
+
train_losses: list[dict[str, float]] = []
|
| 1249 |
+
for batch in train_loader:
|
| 1250 |
+
moved = _move_batch_to_device(batch, device)
|
| 1251 |
+
loss_dict = trainer.training_step(moved, loss_weights=loss_weights)
|
| 1252 |
+
train_losses.append({key: float(value.detach().cpu()) for key, value in loss_dict.items()})
|
| 1253 |
+
|
| 1254 |
+
model.eval()
|
| 1255 |
+
val_losses: list[dict[str, float]] = []
|
| 1256 |
+
with torch.no_grad():
|
| 1257 |
+
for batch in val_loader:
|
| 1258 |
+
moved = _move_batch_to_device(batch, device)
|
| 1259 |
+
forward_kwargs = {
|
| 1260 |
+
"images": moved["images"],
|
| 1261 |
+
"proprio": moved["proprio"],
|
| 1262 |
+
"texts": moved["texts"],
|
| 1263 |
+
"task_names": moved.get("task_name"),
|
| 1264 |
+
"task_ids": moved.get("task_id"),
|
| 1265 |
+
"history_images": moved.get("history_images"),
|
| 1266 |
+
"history_proprio": moved.get("history_proprio"),
|
| 1267 |
+
"history_actions": moved.get("history_actions"),
|
| 1268 |
+
"depths": moved.get("depths"),
|
| 1269 |
+
"depth_valid": moved.get("depth_valid"),
|
| 1270 |
+
"camera_intrinsics": moved.get("camera_intrinsics"),
|
| 1271 |
+
"camera_extrinsics": moved.get("camera_extrinsics"),
|
| 1272 |
+
"history_depths": moved.get("history_depths"),
|
| 1273 |
+
"history_depth_valid": moved.get("history_depth_valid"),
|
| 1274 |
+
"history_camera_intrinsics": moved.get("history_camera_intrinsics"),
|
| 1275 |
+
"history_camera_extrinsics": moved.get("history_camera_extrinsics"),
|
| 1276 |
+
}
|
| 1277 |
+
if variant == "adapter_active_ft":
|
| 1278 |
+
forward_kwargs["adapter_mode"] = "adapter_active"
|
| 1279 |
+
forward_kwargs["use_transition_model"] = True
|
| 1280 |
+
forward_kwargs["use_task_conditioning"] = True
|
| 1281 |
+
outputs = model(**forward_kwargs)
|
| 1282 |
+
from train.losses import compute_total_loss
|
| 1283 |
+
|
| 1284 |
+
losses = compute_total_loss(outputs, moved, weights=loss_weights)
|
| 1285 |
+
val_losses.append({key: float(value.detach().cpu()) for key, value in losses.items()})
|
| 1286 |
+
|
| 1287 |
+
train_summary = _aggregate_epoch(train_losses)
|
| 1288 |
+
val_summary = _aggregate_epoch(val_losses)
|
| 1289 |
+
history.append({"epoch": epoch, "train": train_summary, "val": val_summary})
|
| 1290 |
+
print(
|
| 1291 |
+
json.dumps(
|
| 1292 |
+
{
|
| 1293 |
+
"phase": "epoch_complete",
|
| 1294 |
+
"variant": variant,
|
| 1295 |
+
"epoch": epoch,
|
| 1296 |
+
"train_total": train_summary.get("total", 0.0),
|
| 1297 |
+
"val_total": val_summary.get("total", 0.0),
|
| 1298 |
+
}
|
| 1299 |
+
),
|
| 1300 |
+
flush=True,
|
| 1301 |
+
)
|
| 1302 |
+
if val_summary.get("total", math.inf) <= best_val:
|
| 1303 |
+
best_val = val_summary["total"]
|
| 1304 |
+
checkpoint_path = _save_training_checkpoint(
|
| 1305 |
+
output_dir=output_dir,
|
| 1306 |
+
experiment_name=f"{variant}_seed{spec.train_seed}",
|
| 1307 |
+
model=model,
|
| 1308 |
+
policy_config=policy_config,
|
| 1309 |
+
trainer_config=trainer_config,
|
| 1310 |
+
loss_weights=loss_weights,
|
| 1311 |
+
history=history,
|
| 1312 |
+
best_val=best_val,
|
| 1313 |
+
train_spec=train_spec,
|
| 1314 |
+
)
|
| 1315 |
+
(output_dir / "summary.json").write_text(
|
| 1316 |
+
json.dumps(
|
| 1317 |
+
{
|
| 1318 |
+
"variant": variant,
|
| 1319 |
+
"checkpoint_path": str(checkpoint_path),
|
| 1320 |
+
"init_info": init_info,
|
| 1321 |
+
"trainable_parameter_names": matched,
|
| 1322 |
+
"best_val_total": best_val,
|
| 1323 |
+
"history": history,
|
| 1324 |
+
"train_spec": train_spec,
|
| 1325 |
+
},
|
| 1326 |
+
indent=2,
|
| 1327 |
+
)
|
| 1328 |
+
+ "\n",
|
| 1329 |
+
encoding="utf-8",
|
| 1330 |
+
)
|
| 1331 |
+
return output_dir / "checkpoint_best.pt", train_spec
|
| 1332 |
+
|
| 1333 |
+
|
| 1334 |
+
def _load_checkpoint(
|
| 1335 |
+
checkpoint_path: Path,
|
| 1336 |
+
*,
|
| 1337 |
+
adapter_mode: str | None = None,
|
| 1338 |
+
planner_overrides: dict[str, float] | None = None,
|
| 1339 |
+
) -> tuple[torch.nn.Module, dict[str, Any]]:
|
| 1340 |
+
checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
|
| 1341 |
+
policy_config = PolicyConfig(
|
| 1342 |
+
backbone=FrozenVLBackboneConfig(**checkpoint["policy_config"]["backbone"]),
|
| 1343 |
+
fusion=MultiViewFusionConfig(**checkpoint["policy_config"]["fusion"]),
|
| 1344 |
+
memory=ObservationMemoryConfig(**checkpoint["policy_config"]["memory"]),
|
| 1345 |
+
decoder=ChunkDecoderConfig(**checkpoint["policy_config"]["decoder"]),
|
| 1346 |
+
reveal_head=RevealHeadConfig(**checkpoint["policy_config"]["reveal_head"]),
|
| 1347 |
+
world_model=RevealWMConfig(**checkpoint["policy_config"]["world_model"]),
|
| 1348 |
+
planner=PlannerConfig(**checkpoint["policy_config"]["planner"]),
|
| 1349 |
+
)
|
| 1350 |
+
policy_config = _apply_smoke_planner_overrides(policy_config, planner_overrides=planner_overrides)
|
| 1351 |
+
trainer_config = TrainerConfig(**checkpoint["trainer_config"])
|
| 1352 |
+
if adapter_mode is not None and trainer_config.policy_type == "adapter_wrapped":
|
| 1353 |
+
trainer_config.adapter_mode = adapter_mode
|
| 1354 |
+
trainer_config.eval_mode = adapter_mode
|
| 1355 |
+
model = build_policy(policy_config, trainer_config)
|
| 1356 |
+
filtered_state_dict, skipped, _remapped = filter_compatible_state_dict(model.state_dict(), checkpoint["state_dict"])
|
| 1357 |
+
incompatible = model.load_state_dict(filtered_state_dict, strict=False)
|
| 1358 |
+
if incompatible.unexpected_keys:
|
| 1359 |
+
raise RuntimeError(f"Unexpected checkpoint keys for {checkpoint_path}: {list(incompatible.unexpected_keys)}")
|
| 1360 |
+
if skipped:
|
| 1361 |
+
checkpoint["_shape_skipped_keys"] = skipped
|
| 1362 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 1363 |
+
model = model.to(device)
|
| 1364 |
+
model.eval()
|
| 1365 |
+
return model, checkpoint
|
| 1366 |
+
|
| 1367 |
+
|
| 1368 |
+
def _collect_split(
|
| 1369 |
+
*,
|
| 1370 |
+
split_name: str,
|
| 1371 |
+
seeds: Sequence[int],
|
| 1372 |
+
spec: SmokeSpec,
|
| 1373 |
+
output_path: Path,
|
| 1374 |
+
) -> dict[str, Any]:
|
| 1375 |
+
obs_env = PickClutterRevealEnv(
|
| 1376 |
+
obs_mode="rgb+depth+segmentation",
|
| 1377 |
+
control_mode="pd_ee_delta_pos",
|
| 1378 |
+
render_mode="rgb_array",
|
| 1379 |
+
)
|
| 1380 |
+
sim_env = PickClutterRevealEnv(
|
| 1381 |
+
obs_mode="none",
|
| 1382 |
+
control_mode="pd_ee_delta_pos",
|
| 1383 |
+
render_mode="rgb_array",
|
| 1384 |
+
)
|
| 1385 |
+
samples: list[dict[str, Any]] = []
|
| 1386 |
+
episode_records: list[dict[str, Any]] = []
|
| 1387 |
+
try:
|
| 1388 |
+
for episode_seed in seeds:
|
| 1389 |
+
obs, _ = obs_env.reset(seed=int(episode_seed))
|
| 1390 |
+
sim_env.reset(seed=int(episode_seed))
|
| 1391 |
+
_sync_env_state(obs_env, sim_env)
|
| 1392 |
+
episode_start_positions = _all_positions(obs_env)
|
| 1393 |
+
history: deque[dict[str, Any]] = deque(maxlen=spec.history_steps)
|
| 1394 |
+
episode_success = False
|
| 1395 |
+
for decision_step in range(spec.max_macro_steps):
|
| 1396 |
+
obs_bundle = _extract_sensor_bundle(obs)
|
| 1397 |
+
proprio = _build_proprio(obs_env)
|
| 1398 |
+
snapshot = _snapshot_env(obs_env)
|
| 1399 |
+
candidate_metrics = [_evaluate_candidate(sim_env, obs_env, snapshot, mode_name) for mode_name in MODE_ORDER]
|
| 1400 |
+
candidate_chunks = np.stack([CANONICAL_CHUNKS[mode_name] for mode_name in MODE_ORDER], axis=0).astype(np.float32)
|
| 1401 |
+
utilities = np.asarray([payload["candidate_utility"] for payload in candidate_metrics], dtype=np.float32)
|
| 1402 |
+
best_index = int(utilities.argmax())
|
| 1403 |
+
selected_mode = MODE_ORDER[best_index]
|
| 1404 |
+
state_targets = _current_state_targets(
|
| 1405 |
+
obs_env,
|
| 1406 |
+
obs_bundle=obs_bundle,
|
| 1407 |
+
candidate_metrics=candidate_metrics,
|
| 1408 |
+
episode_start_positions=episode_start_positions,
|
| 1409 |
+
selected_mode=selected_mode,
|
| 1410 |
+
)
|
| 1411 |
+
rollout_targets_by_mode = [
|
| 1412 |
+
_candidate_rollout_targets(
|
| 1413 |
+
mode_name=mode_name,
|
| 1414 |
+
state_targets=state_targets,
|
| 1415 |
+
candidate_payload=payload,
|
| 1416 |
+
)
|
| 1417 |
+
for mode_name, payload in zip(MODE_ORDER, candidate_metrics)
|
| 1418 |
+
]
|
| 1419 |
+
sample = {
|
| 1420 |
+
"images": obs_bundle["images"].copy(),
|
| 1421 |
+
"depths": obs_bundle["depths"].copy(),
|
| 1422 |
+
"depth_valid": obs_bundle["depth_valid"].copy(),
|
| 1423 |
+
"camera_intrinsics": obs_bundle["camera_intrinsics"].copy(),
|
| 1424 |
+
"camera_extrinsics": obs_bundle["camera_extrinsics"].copy(),
|
| 1425 |
+
"history_images": _history_stack(history, "images", pad_shape=obs_bundle["images"].shape, dtype=np.uint8, history_steps=spec.history_steps),
|
| 1426 |
+
"history_depths": _history_stack(history, "depths", pad_shape=obs_bundle["depths"].shape, dtype=np.float32, history_steps=spec.history_steps),
|
| 1427 |
+
"history_depth_valid": _history_stack(history, "depth_valid", pad_shape=obs_bundle["depth_valid"].shape, dtype=np.float32, history_steps=spec.history_steps),
|
| 1428 |
+
"history_camera_intrinsics": _history_stack(history, "camera_intrinsics", pad_shape=obs_bundle["camera_intrinsics"].shape, dtype=np.float32, history_steps=spec.history_steps),
|
| 1429 |
+
"history_camera_extrinsics": _history_stack(history, "camera_extrinsics", pad_shape=obs_bundle["camera_extrinsics"].shape, dtype=np.float32, history_steps=spec.history_steps),
|
| 1430 |
+
"history_proprio": _history_stack(history, "proprio", pad_shape=(PROPRIO_DIM,), dtype=np.float32, history_steps=spec.history_steps),
|
| 1431 |
+
"history_actions": _history_stack(history, "action", pad_shape=(14,), dtype=np.float32, history_steps=spec.history_steps),
|
| 1432 |
+
"proprio": proprio.astype(np.float32),
|
| 1433 |
+
"language_goal": TEXT_PROMPT,
|
| 1434 |
+
"task_name": TASK_NAME,
|
| 1435 |
+
"task_id": TASK_ID,
|
| 1436 |
+
"action_chunk": CANONICAL_CHUNKS[selected_mode].copy(),
|
| 1437 |
+
"candidate_action_chunks": candidate_chunks,
|
| 1438 |
+
"candidate_retrieval_success": np.asarray([payload["candidate_retrieval_success"] for payload in candidate_metrics], dtype=np.float32),
|
| 1439 |
+
"candidate_final_disturbance_cost": np.asarray([payload["candidate_final_disturbance_cost"] for payload in candidate_metrics], dtype=np.float32),
|
| 1440 |
+
"candidate_reocclusion_rate": np.asarray([payload["candidate_reocclusion_rate"] for payload in candidate_metrics], dtype=np.float32),
|
| 1441 |
+
"candidate_visibility_integral": np.asarray([payload["candidate_visibility_integral"] for payload in candidate_metrics], dtype=np.float32),
|
| 1442 |
+
"candidate_actor_feasibility_auc": np.asarray([payload["candidate_actor_feasibility_auc"] for payload in candidate_metrics], dtype=np.float32),
|
| 1443 |
+
"candidate_reveal_achieved": np.asarray([payload["candidate_reveal_achieved"] for payload in candidate_metrics], dtype=np.float32),
|
| 1444 |
+
"candidate_hold_persistence": np.asarray([payload["candidate_hold_persistence"] for payload in candidate_metrics], dtype=np.float32),
|
| 1445 |
+
"candidate_support_stability_auc": np.asarray([payload["candidate_support_stability_auc"] for payload in candidate_metrics], dtype=np.float32),
|
| 1446 |
+
"candidate_disturbance_auc": np.asarray([payload["candidate_disturbance_auc"] for payload in candidate_metrics], dtype=np.float32),
|
| 1447 |
+
"candidate_risk": np.asarray([payload["candidate_risk"] for payload in candidate_metrics], dtype=np.float32),
|
| 1448 |
+
"candidate_utility": utilities,
|
| 1449 |
+
"candidate_rollout_support_mode": np.stack(
|
| 1450 |
+
[payload["candidate_rollout_support_mode"] for payload in rollout_targets_by_mode],
|
| 1451 |
+
axis=0,
|
| 1452 |
+
).astype(np.int64),
|
| 1453 |
+
"candidate_rollout_corridor_feasible": np.stack(
|
| 1454 |
+
[payload["candidate_rollout_corridor_feasible"] for payload in rollout_targets_by_mode],
|
| 1455 |
+
axis=0,
|
| 1456 |
+
).astype(np.float32),
|
| 1457 |
+
"candidate_rollout_persistence_horizon": np.stack(
|
| 1458 |
+
[payload["candidate_rollout_persistence_horizon"] for payload in rollout_targets_by_mode],
|
| 1459 |
+
axis=0,
|
| 1460 |
+
).astype(np.float32),
|
| 1461 |
+
"candidate_rollout_disturbance_cost": np.stack(
|
| 1462 |
+
[payload["candidate_rollout_disturbance_cost"] for payload in rollout_targets_by_mode],
|
| 1463 |
+
axis=0,
|
| 1464 |
+
).astype(np.float32),
|
| 1465 |
+
"candidate_rollout_belief_map": np.stack(
|
| 1466 |
+
[payload["candidate_rollout_belief_map"] for payload in rollout_targets_by_mode],
|
| 1467 |
+
axis=0,
|
| 1468 |
+
).astype(np.float32),
|
| 1469 |
+
"candidate_rollout_visibility_map": np.stack(
|
| 1470 |
+
[payload["candidate_rollout_visibility_map"] for payload in rollout_targets_by_mode],
|
| 1471 |
+
axis=0,
|
| 1472 |
+
).astype(np.float32),
|
| 1473 |
+
"candidate_rollout_clearance_map": np.stack(
|
| 1474 |
+
[payload["candidate_rollout_clearance_map"] for payload in rollout_targets_by_mode],
|
| 1475 |
+
axis=0,
|
| 1476 |
+
).astype(np.float32),
|
| 1477 |
+
"candidate_rollout_support_stability": np.stack(
|
| 1478 |
+
[payload["candidate_rollout_support_stability"] for payload in rollout_targets_by_mode],
|
| 1479 |
+
axis=0,
|
| 1480 |
+
).astype(np.float32),
|
| 1481 |
+
"candidate_rollout_reocclusion_target": np.stack(
|
| 1482 |
+
[payload["candidate_rollout_reocclusion_target"] for payload in rollout_targets_by_mode],
|
| 1483 |
+
axis=0,
|
| 1484 |
+
).astype(np.float32),
|
| 1485 |
+
"candidate_rollout_occluder_contact_map": np.stack(
|
| 1486 |
+
[payload["candidate_rollout_occluder_contact_map"] for payload in rollout_targets_by_mode],
|
| 1487 |
+
axis=0,
|
| 1488 |
+
).astype(np.float32),
|
| 1489 |
+
"candidate_rollout_grasp_affordance_map": np.stack(
|
| 1490 |
+
[payload["candidate_rollout_grasp_affordance_map"] for payload in rollout_targets_by_mode],
|
| 1491 |
+
axis=0,
|
| 1492 |
+
).astype(np.float32),
|
| 1493 |
+
"proposal_target_rollout_support_mode": np.stack(
|
| 1494 |
+
[payload["candidate_rollout_support_mode"] for payload in rollout_targets_by_mode],
|
| 1495 |
+
axis=0,
|
| 1496 |
+
).astype(np.int64),
|
| 1497 |
+
"proposal_target_rollout_corridor_feasible": np.stack(
|
| 1498 |
+
[payload["candidate_rollout_corridor_feasible"] for payload in rollout_targets_by_mode],
|
| 1499 |
+
axis=0,
|
| 1500 |
+
).astype(np.float32),
|
| 1501 |
+
"proposal_target_rollout_persistence_horizon": np.stack(
|
| 1502 |
+
[payload["candidate_rollout_persistence_horizon"] for payload in rollout_targets_by_mode],
|
| 1503 |
+
axis=0,
|
| 1504 |
+
).astype(np.float32),
|
| 1505 |
+
"proposal_target_rollout_disturbance_cost": np.stack(
|
| 1506 |
+
[payload["candidate_rollout_disturbance_cost"] for payload in rollout_targets_by_mode],
|
| 1507 |
+
axis=0,
|
| 1508 |
+
).astype(np.float32),
|
| 1509 |
+
"proposal_target_rollout_belief_map": np.stack(
|
| 1510 |
+
[payload["candidate_rollout_belief_map"] for payload in rollout_targets_by_mode],
|
| 1511 |
+
axis=0,
|
| 1512 |
+
).astype(np.float32),
|
| 1513 |
+
"proposal_target_rollout_visibility_map": np.stack(
|
| 1514 |
+
[payload["candidate_rollout_visibility_map"] for payload in rollout_targets_by_mode],
|
| 1515 |
+
axis=0,
|
| 1516 |
+
).astype(np.float32),
|
| 1517 |
+
"proposal_target_rollout_clearance_map": np.stack(
|
| 1518 |
+
[payload["candidate_rollout_clearance_map"] for payload in rollout_targets_by_mode],
|
| 1519 |
+
axis=0,
|
| 1520 |
+
).astype(np.float32),
|
| 1521 |
+
"proposal_target_rollout_support_stability": np.stack(
|
| 1522 |
+
[payload["candidate_rollout_support_stability"] for payload in rollout_targets_by_mode],
|
| 1523 |
+
axis=0,
|
| 1524 |
+
).astype(np.float32),
|
| 1525 |
+
"proposal_target_rollout_reocclusion_target": np.stack(
|
| 1526 |
+
[payload["candidate_rollout_reocclusion_target"] for payload in rollout_targets_by_mode],
|
| 1527 |
+
axis=0,
|
| 1528 |
+
).astype(np.float32),
|
| 1529 |
+
"proposal_target_rollout_occluder_contact_map": np.stack(
|
| 1530 |
+
[payload["candidate_rollout_occluder_contact_map"] for payload in rollout_targets_by_mode],
|
| 1531 |
+
axis=0,
|
| 1532 |
+
).astype(np.float32),
|
| 1533 |
+
"proposal_target_rollout_grasp_affordance_map": np.stack(
|
| 1534 |
+
[payload["candidate_rollout_grasp_affordance_map"] for payload in rollout_targets_by_mode],
|
| 1535 |
+
axis=0,
|
| 1536 |
+
).astype(np.float32),
|
| 1537 |
+
"episode_seed": int(episode_seed),
|
| 1538 |
+
"decision_step": int(decision_step),
|
| 1539 |
+
"selected_mode": selected_mode,
|
| 1540 |
+
**state_targets,
|
| 1541 |
+
}
|
| 1542 |
+
samples.append(sample)
|
| 1543 |
+
_execute_mode(obs_env, selected_mode)
|
| 1544 |
+
obs = obs_env.get_obs(obs_env.get_info())
|
| 1545 |
+
history.append(_init_history_entry(obs_bundle, proprio, CANONICAL_CHUNKS[selected_mode]))
|
| 1546 |
+
if _success_from_state(obs_env):
|
| 1547 |
+
episode_success = True
|
| 1548 |
+
break
|
| 1549 |
+
episode_records.append(
|
| 1550 |
+
{
|
| 1551 |
+
"episode_seed": int(episode_seed),
|
| 1552 |
+
"success": episode_success,
|
| 1553 |
+
"steps": len(history),
|
| 1554 |
+
}
|
| 1555 |
+
)
|
| 1556 |
+
print(
|
| 1557 |
+
json.dumps(
|
| 1558 |
+
{
|
| 1559 |
+
"phase": "collect_episode_complete",
|
| 1560 |
+
"split": split_name,
|
| 1561 |
+
"episode_seed": int(episode_seed),
|
| 1562 |
+
"success": episode_success,
|
| 1563 |
+
"steps": len(history),
|
| 1564 |
+
"samples_collected": len(samples),
|
| 1565 |
+
}
|
| 1566 |
+
),
|
| 1567 |
+
flush=True,
|
| 1568 |
+
)
|
| 1569 |
+
finally:
|
| 1570 |
+
obs_env.close()
|
| 1571 |
+
sim_env.close()
|
| 1572 |
+
|
| 1573 |
+
payload = {
|
| 1574 |
+
"split_name": split_name,
|
| 1575 |
+
"resolution": spec.resolution,
|
| 1576 |
+
"history_steps": spec.history_steps,
|
| 1577 |
+
"samples": samples,
|
| 1578 |
+
"episode_records": episode_records,
|
| 1579 |
+
}
|
| 1580 |
+
output_path.parent.mkdir(parents=True, exist_ok=True)
|
| 1581 |
+
torch.save(payload, output_path)
|
| 1582 |
+
return payload
|
| 1583 |
+
|
| 1584 |
+
|
| 1585 |
+
def _load_split(path: Path) -> dict[str, Any]:
|
| 1586 |
+
return torch.load(path, map_location="cpu", weights_only=False)
|
| 1587 |
+
|
| 1588 |
+
|
| 1589 |
+
def _print_split_supervision_summary(split_name: str, samples: Sequence[dict[str, Any]]) -> None:
|
| 1590 |
+
mode_counter = collections.Counter(str(sample.get("selected_mode", "unknown")) for sample in samples)
|
| 1591 |
+
support_counter = collections.Counter(int(sample.get("support_mode", -1)) for sample in samples if "support_mode" in sample)
|
| 1592 |
+
confidence_values = [float(sample.get("state_confidence_target", 0.0)) for sample in samples if "state_confidence_target" in sample]
|
| 1593 |
+
payload = {
|
| 1594 |
+
"phase": "split_supervision_summary",
|
| 1595 |
+
"split": split_name,
|
| 1596 |
+
"samples": len(samples),
|
| 1597 |
+
"selected_modes": dict(mode_counter),
|
| 1598 |
+
"support_modes": dict(support_counter),
|
| 1599 |
+
"mean_state_confidence_target": float(np.mean(confidence_values)) if confidence_values else None,
|
| 1600 |
+
}
|
| 1601 |
+
print(json.dumps(payload, sort_keys=True), flush=True)
|
| 1602 |
+
|
| 1603 |
+
|
| 1604 |
+
def _batch_from_obs(obs_bundle: dict[str, np.ndarray], proprio: np.ndarray, history: Sequence[dict[str, Any]], device: torch.device) -> dict[str, Any]:
|
| 1605 |
+
return {
|
| 1606 |
+
"images": torch.from_numpy(obs_bundle["images"]).permute(0, 3, 1, 2).unsqueeze(0).float().div(255.0).to(device),
|
| 1607 |
+
"depths": torch.from_numpy(obs_bundle["depths"]).unsqueeze(0).float().to(device),
|
| 1608 |
+
"depth_valid": torch.from_numpy(obs_bundle["depth_valid"]).unsqueeze(0).float().to(device),
|
| 1609 |
+
"camera_intrinsics": torch.from_numpy(obs_bundle["camera_intrinsics"]).unsqueeze(0).float().to(device),
|
| 1610 |
+
"camera_extrinsics": torch.from_numpy(obs_bundle["camera_extrinsics"]).unsqueeze(0).float().to(device),
|
| 1611 |
+
"history_images": torch.from_numpy(
|
| 1612 |
+
_history_stack(history, "images", pad_shape=obs_bundle["images"].shape, dtype=np.uint8, history_steps=HISTORY_STEPS)
|
| 1613 |
+
).permute(0, 1, 4, 2, 3).unsqueeze(0).float().div(255.0).to(device),
|
| 1614 |
+
"history_depths": torch.from_numpy(
|
| 1615 |
+
_history_stack(history, "depths", pad_shape=obs_bundle["depths"].shape, dtype=np.float32, history_steps=HISTORY_STEPS)
|
| 1616 |
+
).unsqueeze(0).float().to(device),
|
| 1617 |
+
"history_depth_valid": torch.from_numpy(
|
| 1618 |
+
_history_stack(history, "depth_valid", pad_shape=obs_bundle["depth_valid"].shape, dtype=np.float32, history_steps=HISTORY_STEPS)
|
| 1619 |
+
).unsqueeze(0).float().to(device),
|
| 1620 |
+
"history_camera_intrinsics": torch.from_numpy(
|
| 1621 |
+
_history_stack(history, "camera_intrinsics", pad_shape=obs_bundle["camera_intrinsics"].shape, dtype=np.float32, history_steps=HISTORY_STEPS)
|
| 1622 |
+
).unsqueeze(0).float().to(device),
|
| 1623 |
+
"history_camera_extrinsics": torch.from_numpy(
|
| 1624 |
+
_history_stack(history, "camera_extrinsics", pad_shape=obs_bundle["camera_extrinsics"].shape, dtype=np.float32, history_steps=HISTORY_STEPS)
|
| 1625 |
+
).unsqueeze(0).float().to(device),
|
| 1626 |
+
"history_proprio": torch.from_numpy(
|
| 1627 |
+
_history_stack(history, "proprio", pad_shape=(PROPRIO_DIM,), dtype=np.float32, history_steps=HISTORY_STEPS)
|
| 1628 |
+
).unsqueeze(0).float().to(device),
|
| 1629 |
+
"history_actions": torch.from_numpy(
|
| 1630 |
+
_history_stack(history, "action", pad_shape=(14,), dtype=np.float32, history_steps=HISTORY_STEPS)
|
| 1631 |
+
).unsqueeze(0).float().to(device),
|
| 1632 |
+
"proprio": torch.from_numpy(proprio).unsqueeze(0).float().to(device),
|
| 1633 |
+
"texts": [TEXT_PROMPT],
|
| 1634 |
+
"task_names": [TASK_NAME],
|
| 1635 |
+
"task_ids": torch.as_tensor([TASK_ID], dtype=torch.long, device=device),
|
| 1636 |
+
}
|
| 1637 |
+
|
| 1638 |
+
|
| 1639 |
+
def _eval_mode_name(
|
| 1640 |
+
model_output: dict[str, Any],
|
| 1641 |
+
checkpoint_mode: str,
|
| 1642 |
+
) -> tuple[str, bool, bool]:
|
| 1643 |
+
if checkpoint_mode == "adapter_active_ft" and "proposal_mode_names" in model_output and "best_candidate_indices" in model_output:
|
| 1644 |
+
active_mask = bool(_np(model_output.get("adapter_active_mask", np.asarray([False]))).reshape(-1)[0])
|
| 1645 |
+
if not active_mask:
|
| 1646 |
+
mode_name = _classify_mode_from_chunk(_np(model_output["action_mean"])[0])
|
| 1647 |
+
return mode_name, False, False
|
| 1648 |
+
best_index = int(_np(model_output["best_candidate_indices"])[0])
|
| 1649 |
+
proposal_mode_names = model_output["proposal_mode_names"][0]
|
| 1650 |
+
mode_name = str(proposal_mode_names[best_index]) if best_index < len(proposal_mode_names) else _classify_mode_from_chunk(
|
| 1651 |
+
_np(model_output["action_mean"])[0]
|
| 1652 |
+
)
|
| 1653 |
+
non_base = bool(best_index > 0)
|
| 1654 |
+
return mode_name, active_mask, non_base
|
| 1655 |
+
mode_name = _classify_mode_from_chunk(_np(model_output["action_mean"])[0])
|
| 1656 |
+
return mode_name, False, False
|
| 1657 |
+
|
| 1658 |
+
|
| 1659 |
+
def _evaluate_checkpoint(
|
| 1660 |
+
*,
|
| 1661 |
+
checkpoint_path: Path,
|
| 1662 |
+
adapter_mode: str,
|
| 1663 |
+
result_mode_name: str,
|
| 1664 |
+
seeds: Sequence[int],
|
| 1665 |
+
report_path: Path,
|
| 1666 |
+
train_spec: dict[str, Any] | None,
|
| 1667 |
+
dataset_seed: int,
|
| 1668 |
+
planner_overrides: dict[str, float] | None = None,
|
| 1669 |
+
) -> dict[str, Any]:
|
| 1670 |
+
model, checkpoint = _load_checkpoint(
|
| 1671 |
+
checkpoint_path,
|
| 1672 |
+
adapter_mode=adapter_mode if adapter_mode != "trunk_only" else None,
|
| 1673 |
+
planner_overrides=planner_overrides,
|
| 1674 |
+
)
|
| 1675 |
+
device = next(model.parameters()).device
|
| 1676 |
+
obs_env = PickClutterRevealEnv(
|
| 1677 |
+
obs_mode="rgb+depth+segmentation",
|
| 1678 |
+
control_mode="pd_ee_delta_pos",
|
| 1679 |
+
render_mode="rgb_array",
|
| 1680 |
+
)
|
| 1681 |
+
sim_env = PickClutterRevealEnv(
|
| 1682 |
+
obs_mode="none",
|
| 1683 |
+
control_mode="pd_ee_delta_pos",
|
| 1684 |
+
render_mode="rgb_array",
|
| 1685 |
+
)
|
| 1686 |
+
successes: list[int] = []
|
| 1687 |
+
episode_records: list[dict[str, Any]] = []
|
| 1688 |
+
reveal_steps: list[int] = []
|
| 1689 |
+
retrieve_steps: list[int] = []
|
| 1690 |
+
disturbance_values: list[float] = []
|
| 1691 |
+
intervention_events = 0
|
| 1692 |
+
non_base_events = 0
|
| 1693 |
+
total_decisions = 0
|
| 1694 |
+
try:
|
| 1695 |
+
for episode_seed in seeds:
|
| 1696 |
+
obs, _ = obs_env.reset(seed=int(episode_seed))
|
| 1697 |
+
sim_env.reset(seed=int(episode_seed))
|
| 1698 |
+
_sync_env_state(obs_env, sim_env)
|
| 1699 |
+
history: deque[dict[str, Any]] = deque(maxlen=HISTORY_STEPS)
|
| 1700 |
+
success = False
|
| 1701 |
+
first_reveal_step: int | None = None
|
| 1702 |
+
first_retrieve_step: int | None = None
|
| 1703 |
+
episode_disturbance: list[float] = []
|
| 1704 |
+
for decision_step in range(MAX_MACRO_STEPS):
|
| 1705 |
+
obs_bundle = _extract_sensor_bundle(obs)
|
| 1706 |
+
proprio = _build_proprio(obs_env)
|
| 1707 |
+
batch = _batch_from_obs(obs_bundle, proprio, list(history), device)
|
| 1708 |
+
with torch.no_grad():
|
| 1709 |
+
if adapter_mode == "trunk_only":
|
| 1710 |
+
outputs = model(**batch)
|
| 1711 |
+
else:
|
| 1712 |
+
outputs = model(
|
| 1713 |
+
**batch,
|
| 1714 |
+
adapter_mode=adapter_mode,
|
| 1715 |
+
use_transition_model=True,
|
| 1716 |
+
use_task_conditioning=True,
|
| 1717 |
+
)
|
| 1718 |
+
selected_mode, active_mask, non_base = _eval_mode_name(outputs, result_mode_name)
|
| 1719 |
+
start_positions = _all_positions(obs_env)
|
| 1720 |
+
_sync_env_state(obs_env, sim_env)
|
| 1721 |
+
_execute_mode(sim_env, selected_mode)
|
| 1722 |
+
end_metrics = _candidate_metrics(sim_env, start_positions=start_positions, current_obs_bundle=None)
|
| 1723 |
+
_sync_env_state(sim_env, obs_env)
|
| 1724 |
+
obs = obs_env.get_obs(obs_env.get_info())
|
| 1725 |
+
history.append(_init_history_entry(obs_bundle, proprio, CANONICAL_CHUNKS.get(selected_mode, CANONICAL_CHUNKS["base_action"])))
|
| 1726 |
+
total_decisions += 1
|
| 1727 |
+
intervention_events += int(active_mask)
|
| 1728 |
+
non_base_events += int(non_base)
|
| 1729 |
+
episode_disturbance.append(end_metrics["disturbance"])
|
| 1730 |
+
if selected_mode != "retrieve" and selected_mode not in {"base_action", "maintain_gap"} and first_reveal_step is None:
|
| 1731 |
+
first_reveal_step = decision_step + 1
|
| 1732 |
+
if selected_mode == "retrieve" and first_retrieve_step is None:
|
| 1733 |
+
first_retrieve_step = decision_step + 1
|
| 1734 |
+
if _success_from_state(obs_env):
|
| 1735 |
+
success = True
|
| 1736 |
+
break
|
| 1737 |
+
successes.append(int(success))
|
| 1738 |
+
if first_reveal_step is not None:
|
| 1739 |
+
reveal_steps.append(first_reveal_step)
|
| 1740 |
+
if first_retrieve_step is not None:
|
| 1741 |
+
retrieve_steps.append(first_retrieve_step)
|
| 1742 |
+
disturbance_values.append(float(np.mean(episode_disturbance)) if episode_disturbance else 0.0)
|
| 1743 |
+
episode_records.append(
|
| 1744 |
+
{
|
| 1745 |
+
"episode_seed": int(episode_seed),
|
| 1746 |
+
"success": success,
|
| 1747 |
+
"steps": len(history),
|
| 1748 |
+
"first_reveal_step": first_reveal_step,
|
| 1749 |
+
"first_retrieve_step": first_retrieve_step,
|
| 1750 |
+
"episode_disturbance": float(np.mean(episode_disturbance)) if episode_disturbance else 0.0,
|
| 1751 |
+
}
|
| 1752 |
+
)
|
| 1753 |
+
print(
|
| 1754 |
+
json.dumps(
|
| 1755 |
+
{
|
| 1756 |
+
"phase": "eval_episode_complete",
|
| 1757 |
+
"adapter_mode": result_mode_name,
|
| 1758 |
+
"episode_seed": int(episode_seed),
|
| 1759 |
+
"success": success,
|
| 1760 |
+
"steps": len(history),
|
| 1761 |
+
}
|
| 1762 |
+
),
|
| 1763 |
+
flush=True,
|
| 1764 |
+
)
|
| 1765 |
+
finally:
|
| 1766 |
+
obs_env.close()
|
| 1767 |
+
sim_env.close()
|
| 1768 |
+
|
| 1769 |
+
eval_protocol = build_public_eval_protocol(
|
| 1770 |
+
track_id="occlusion_track",
|
| 1771 |
+
eval_mode=result_mode_name,
|
| 1772 |
+
seed=int(dataset_seed),
|
| 1773 |
+
episodes=len(seeds),
|
| 1774 |
+
resolution=224,
|
| 1775 |
+
cameras=CAMERA_NAMES,
|
| 1776 |
+
)
|
| 1777 |
+
payload = {
|
| 1778 |
+
"track_id": "occlusion_track",
|
| 1779 |
+
"suite": "maniskill3",
|
| 1780 |
+
"benchmark_task": "PickClutterYCB-v1",
|
| 1781 |
+
"role": "target",
|
| 1782 |
+
"adapter_mode": result_mode_name,
|
| 1783 |
+
"episodes": len(seeds),
|
| 1784 |
+
"successes": successes,
|
| 1785 |
+
"success_rate": float(np.mean(successes)) if successes else 0.0,
|
| 1786 |
+
"intervention_rate": float(intervention_events / max(1, total_decisions)),
|
| 1787 |
+
"non_base_selection_rate": float(non_base_events / max(1, total_decisions)),
|
| 1788 |
+
"steps_to_first_reveal_or_access": float(np.mean(reveal_steps)) if reveal_steps else float(MAX_MACRO_STEPS),
|
| 1789 |
+
"steps_to_retrieve": float(np.mean(retrieve_steps)) if retrieve_steps else float(MAX_MACRO_STEPS),
|
| 1790 |
+
"disturbance_proxy": float(np.mean(disturbance_values)) if disturbance_values else 0.0,
|
| 1791 |
+
"episode_records": episode_records,
|
| 1792 |
+
"eval_protocol": eval_protocol,
|
| 1793 |
+
}
|
| 1794 |
+
if train_spec is not None:
|
| 1795 |
+
payload["train_spec"] = train_spec
|
| 1796 |
+
report_path.parent.mkdir(parents=True, exist_ok=True)
|
| 1797 |
+
report_path.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
|
| 1798 |
+
return payload
|
| 1799 |
+
|
| 1800 |
+
|
| 1801 |
+
def _summarize_smoke(results: Sequence[dict[str, Any]], output_dir: Path) -> dict[str, Any]:
|
| 1802 |
+
summary = summarize_public_benchmark_package(list(results), allow_partial=True)
|
| 1803 |
+
output_dir.mkdir(parents=True, exist_ok=True)
|
| 1804 |
+
json_path = output_dir / "public_benchmark_package_summary.json"
|
| 1805 |
+
md_path = output_dir / "public_benchmark_package_summary.md"
|
| 1806 |
+
json_path.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8")
|
| 1807 |
+
lines = [
|
| 1808 |
+
"# ManiSkill PickClutter Smoke Summary",
|
| 1809 |
+
"",
|
| 1810 |
+
f"- available_tracks: {summary['available_tracks']}",
|
| 1811 |
+
f"- target_macro_average_delta: {summary['target_macro_average_delta']:.3f}",
|
| 1812 |
+
f"- headline_pass: {summary['headline_pass']}",
|
| 1813 |
+
f"- sign_of_life_pass: {summary['sign_of_life_pass']}",
|
| 1814 |
+
"",
|
| 1815 |
+
]
|
| 1816 |
+
for track_id, payload in summary["tracks"].items():
|
| 1817 |
+
lines.append(f"## {track_id}")
|
| 1818 |
+
lines.append(f"- delta_active_vs_trunk: {payload.get('delta_active_vs_trunk', 0.0):.3f}")
|
| 1819 |
+
lines.append(f"- delta_noop_vs_trunk: {payload.get('delta_noop_vs_trunk', 0.0):.3f}")
|
| 1820 |
+
lines.append(f"- signs_of_life: {payload.get('signs_of_life', False)}")
|
| 1821 |
+
for mode, mode_payload in payload["modes"].items():
|
| 1822 |
+
lines.append(f"- {mode}: mean_success={mode_payload['mean_success']:.3f}")
|
| 1823 |
+
lines.append("")
|
| 1824 |
+
md_path.write_text("\n".join(lines).rstrip() + "\n", encoding="utf-8")
|
| 1825 |
+
return summary
|
| 1826 |
+
|
| 1827 |
+
|
| 1828 |
+
def _default_paths() -> SmokePaths:
|
| 1829 |
+
return SmokePaths()
|
| 1830 |
+
|
| 1831 |
+
|
| 1832 |
+
def _dataset_artifact_path(data_dir: Path, basename: str, *, dataset_seed: int) -> Path:
|
| 1833 |
+
if int(dataset_seed) == DEFAULT_SEED:
|
| 1834 |
+
return data_dir / basename
|
| 1835 |
+
artifact = Path(basename)
|
| 1836 |
+
return data_dir / f"{artifact.stem}_seed{int(dataset_seed)}{artifact.suffix}"
|
| 1837 |
+
|
| 1838 |
+
|
| 1839 |
+
def _parse_args() -> argparse.Namespace:
|
| 1840 |
+
parser = argparse.ArgumentParser(description="Minimum-sign-of-life ManiSkill PickClutter smoke run.")
|
| 1841 |
+
parser.add_argument("--stage", choices=("collect", "train", "eval", "all"), default="all")
|
| 1842 |
+
parser.add_argument("--data-dir", type=Path, default=_default_paths().data_dir)
|
| 1843 |
+
parser.add_argument("--output-dir", type=Path, default=_default_paths().output_dir)
|
| 1844 |
+
parser.add_argument("--report-dir", type=Path, default=_default_paths().report_dir)
|
| 1845 |
+
parser.add_argument("--seed", type=int, default=None, help="Deprecated alias for train/eval dataset seed.")
|
| 1846 |
+
parser.add_argument("--train-seed", type=int, default=None)
|
| 1847 |
+
parser.add_argument("--dataset-seed", type=int, default=None)
|
| 1848 |
+
parser.add_argument("--eval-split", choices=("val", "eval"), default="eval")
|
| 1849 |
+
parser.add_argument("--adapter-confidence-threshold", type=float, default=None)
|
| 1850 |
+
parser.add_argument("--retrieve-access-threshold", type=float, default=None)
|
| 1851 |
+
parser.add_argument("--retrieve-persistence-threshold", type=float, default=None)
|
| 1852 |
+
parser.add_argument("--retrieve-support-threshold", type=float, default=None)
|
| 1853 |
+
parser.add_argument("--retrieve-reocclusion-threshold", type=float, default=None)
|
| 1854 |
+
parser.add_argument("--planner-mode-preference-bonus", type=float, default=None)
|
| 1855 |
+
parser.add_argument("--planner-premature-retrieve-penalty", type=float, default=None)
|
| 1856 |
+
parser.add_argument("--planner-premature-insert-penalty", type=float, default=None)
|
| 1857 |
+
parser.add_argument("--planner-premature-occlusion-sweep-penalty", type=float, default=None)
|
| 1858 |
+
parser.add_argument("--planner-premature-maintain-penalty", type=float, default=None)
|
| 1859 |
+
parser.add_argument("--planner-retrieve-stage-access-threshold", type=float, default=None)
|
| 1860 |
+
parser.add_argument("--planner-retrieve-stage-reveal-threshold", type=float, default=None)
|
| 1861 |
+
parser.add_argument("--planner-retrieve-stage-persistence-threshold", type=float, default=None)
|
| 1862 |
+
parser.add_argument("--planner-retrieve-stage-support-threshold", type=float, default=None)
|
| 1863 |
+
parser.add_argument("--planner-insert-stage-access-threshold", type=float, default=None)
|
| 1864 |
+
parser.add_argument("--planner-insert-stage-visibility-threshold", type=float, default=None)
|
| 1865 |
+
parser.add_argument("--planner-insert-stage-support-threshold", type=float, default=None)
|
| 1866 |
+
parser.add_argument("--planner-occlusion-maintain-gap-min-access", type=float, default=None)
|
| 1867 |
+
parser.add_argument("--planner-occlusion-maintain-gap-min-visibility", type=float, default=None)
|
| 1868 |
+
return parser.parse_args()
|
| 1869 |
+
|
| 1870 |
+
|
| 1871 |
+
def _planner_overrides_from_args(args: argparse.Namespace) -> dict[str, float]:
|
| 1872 |
+
override_pairs = (
|
| 1873 |
+
("adapter_confidence_threshold", args.adapter_confidence_threshold),
|
| 1874 |
+
("retrieve_access_threshold", args.retrieve_access_threshold),
|
| 1875 |
+
("retrieve_persistence_threshold", args.retrieve_persistence_threshold),
|
| 1876 |
+
("retrieve_support_threshold", args.retrieve_support_threshold),
|
| 1877 |
+
("retrieve_reocclusion_threshold", args.retrieve_reocclusion_threshold),
|
| 1878 |
+
("mode_preference_bonus", args.planner_mode_preference_bonus),
|
| 1879 |
+
("premature_retrieve_penalty", args.planner_premature_retrieve_penalty),
|
| 1880 |
+
("premature_insert_penalty", args.planner_premature_insert_penalty),
|
| 1881 |
+
("premature_occlusion_sweep_penalty", args.planner_premature_occlusion_sweep_penalty),
|
| 1882 |
+
("premature_maintain_penalty", args.planner_premature_maintain_penalty),
|
| 1883 |
+
("retrieve_stage_access_threshold", args.planner_retrieve_stage_access_threshold),
|
| 1884 |
+
("retrieve_stage_reveal_threshold", args.planner_retrieve_stage_reveal_threshold),
|
| 1885 |
+
("retrieve_stage_persistence_threshold", args.planner_retrieve_stage_persistence_threshold),
|
| 1886 |
+
("retrieve_stage_support_threshold", args.planner_retrieve_stage_support_threshold),
|
| 1887 |
+
("insert_stage_access_threshold", args.planner_insert_stage_access_threshold),
|
| 1888 |
+
("insert_stage_visibility_threshold", args.planner_insert_stage_visibility_threshold),
|
| 1889 |
+
("insert_stage_support_threshold", args.planner_insert_stage_support_threshold),
|
| 1890 |
+
("occlusion_maintain_gap_min_access", args.planner_occlusion_maintain_gap_min_access),
|
| 1891 |
+
("occlusion_maintain_gap_min_visibility", args.planner_occlusion_maintain_gap_min_visibility),
|
| 1892 |
+
)
|
| 1893 |
+
return {key: value for key, value in override_pairs if value is not None}
|
| 1894 |
+
|
| 1895 |
+
|
| 1896 |
+
def main() -> None:
|
| 1897 |
+
args = _parse_args()
|
| 1898 |
+
planner_overrides = _planner_overrides_from_args(args)
|
| 1899 |
+
base_seed = DEFAULT_SEED if args.seed is None else int(args.seed)
|
| 1900 |
+
train_seed = int(args.train_seed) if args.train_seed is not None else base_seed
|
| 1901 |
+
dataset_seed = int(args.dataset_seed) if args.dataset_seed is not None else base_seed
|
| 1902 |
+
spec = SmokeSpec(dataset_seed=dataset_seed, train_seed=train_seed)
|
| 1903 |
+
splits = _build_episode_splits(spec)
|
| 1904 |
+
split_path = _dataset_artifact_path(args.data_dir, "episode_splits.json", dataset_seed=spec.dataset_seed)
|
| 1905 |
+
train_path = _dataset_artifact_path(args.data_dir, "train.pt", dataset_seed=spec.dataset_seed)
|
| 1906 |
+
val_path = _dataset_artifact_path(args.data_dir, "val.pt", dataset_seed=spec.dataset_seed)
|
| 1907 |
+
|
| 1908 |
+
if args.stage in {"collect", "all"}:
|
| 1909 |
+
_save_episode_splits(split_path, splits)
|
| 1910 |
+
if not train_path.exists():
|
| 1911 |
+
print(json.dumps({"phase": "collect_train_start", "episodes": len(splits["train"])}), flush=True)
|
| 1912 |
+
_collect_split(split_name="train", seeds=splits["train"], spec=spec, output_path=train_path)
|
| 1913 |
+
if not val_path.exists():
|
| 1914 |
+
print(json.dumps({"phase": "collect_val_start", "episodes": len(splits["val"])}), flush=True)
|
| 1915 |
+
_collect_split(split_name="val", seeds=splits["val"], spec=spec, output_path=val_path)
|
| 1916 |
+
|
| 1917 |
+
if args.stage == "collect":
|
| 1918 |
+
return
|
| 1919 |
+
|
| 1920 |
+
train_bundle = _load_split(train_path)
|
| 1921 |
+
val_bundle = _load_split(val_path)
|
| 1922 |
+
train_samples = train_bundle["samples"]
|
| 1923 |
+
val_samples = val_bundle["samples"]
|
| 1924 |
+
_print_split_supervision_summary("train", train_samples)
|
| 1925 |
+
_print_split_supervision_summary("val", val_samples)
|
| 1926 |
+
|
| 1927 |
+
trunk_checkpoint = args.output_dir / f"trunk_only_ft_seed{spec.train_seed}" / "checkpoint_best.pt"
|
| 1928 |
+
adapter_checkpoint = args.output_dir / f"adapter_active_ft_seed{spec.train_seed}" / "checkpoint_best.pt"
|
| 1929 |
+
trunk_train_spec: dict[str, Any] | None = None
|
| 1930 |
+
adapter_train_spec: dict[str, Any] | None = None
|
| 1931 |
+
|
| 1932 |
+
if args.stage in {"train", "all"}:
|
| 1933 |
+
if not trunk_checkpoint.exists():
|
| 1934 |
+
print(json.dumps({"phase": "train_variant_start", "variant": "trunk_only_ft"}), flush=True)
|
| 1935 |
+
trunk_checkpoint, trunk_train_spec = _train_variant(
|
| 1936 |
+
variant="trunk_only_ft",
|
| 1937 |
+
train_samples=train_samples,
|
| 1938 |
+
val_samples=val_samples,
|
| 1939 |
+
spec=spec,
|
| 1940 |
+
output_dir=args.output_dir / f"trunk_only_ft_seed{spec.train_seed}",
|
| 1941 |
+
)
|
| 1942 |
+
else:
|
| 1943 |
+
trunk_payload = torch.load(trunk_checkpoint, map_location="cpu", weights_only=False)
|
| 1944 |
+
trunk_train_spec = trunk_payload.get("train_spec")
|
| 1945 |
+
if not adapter_checkpoint.exists():
|
| 1946 |
+
print(json.dumps({"phase": "train_variant_start", "variant": "adapter_active_ft"}), flush=True)
|
| 1947 |
+
adapter_checkpoint, adapter_train_spec = _train_variant(
|
| 1948 |
+
variant="adapter_active_ft",
|
| 1949 |
+
train_samples=train_samples,
|
| 1950 |
+
val_samples=val_samples,
|
| 1951 |
+
spec=spec,
|
| 1952 |
+
output_dir=args.output_dir / f"adapter_active_ft_seed{spec.train_seed}",
|
| 1953 |
+
)
|
| 1954 |
+
else:
|
| 1955 |
+
adapter_payload = torch.load(adapter_checkpoint, map_location="cpu", weights_only=False)
|
| 1956 |
+
adapter_train_spec = adapter_payload.get("train_spec")
|
| 1957 |
+
|
| 1958 |
+
if args.stage == "train":
|
| 1959 |
+
return
|
| 1960 |
+
|
| 1961 |
+
if trunk_train_spec is None and trunk_checkpoint.exists():
|
| 1962 |
+
trunk_payload = torch.load(trunk_checkpoint, map_location="cpu", weights_only=False)
|
| 1963 |
+
trunk_train_spec = trunk_payload.get("train_spec")
|
| 1964 |
+
if adapter_train_spec is None and adapter_checkpoint.exists():
|
| 1965 |
+
adapter_payload = torch.load(adapter_checkpoint, map_location="cpu", weights_only=False)
|
| 1966 |
+
adapter_train_spec = adapter_payload.get("train_spec")
|
| 1967 |
+
|
| 1968 |
+
eval_seeds = splits[args.eval_split]
|
| 1969 |
+
print(json.dumps({"phase": "eval_start", "episodes": len(eval_seeds)}), flush=True)
|
| 1970 |
+
trunk_result = _evaluate_checkpoint(
|
| 1971 |
+
checkpoint_path=trunk_checkpoint,
|
| 1972 |
+
adapter_mode="trunk_only",
|
| 1973 |
+
result_mode_name="trunk_only_ft",
|
| 1974 |
+
seeds=eval_seeds,
|
| 1975 |
+
report_path=args.report_dir / f"trunk_only_ft_seed{spec.train_seed}.json",
|
| 1976 |
+
train_spec=trunk_train_spec,
|
| 1977 |
+
dataset_seed=spec.dataset_seed,
|
| 1978 |
+
planner_overrides=planner_overrides,
|
| 1979 |
+
)
|
| 1980 |
+
noop_result = _evaluate_checkpoint(
|
| 1981 |
+
checkpoint_path=adapter_checkpoint,
|
| 1982 |
+
adapter_mode="adapter_noop",
|
| 1983 |
+
result_mode_name="adapter_noop",
|
| 1984 |
+
seeds=eval_seeds,
|
| 1985 |
+
report_path=args.report_dir / f"adapter_noop_seed{spec.train_seed}.json",
|
| 1986 |
+
train_spec=adapter_train_spec,
|
| 1987 |
+
dataset_seed=spec.dataset_seed,
|
| 1988 |
+
planner_overrides=planner_overrides,
|
| 1989 |
+
)
|
| 1990 |
+
active_result = _evaluate_checkpoint(
|
| 1991 |
+
checkpoint_path=adapter_checkpoint,
|
| 1992 |
+
adapter_mode="adapter_active",
|
| 1993 |
+
result_mode_name="adapter_active_ft",
|
| 1994 |
+
seeds=eval_seeds,
|
| 1995 |
+
report_path=args.report_dir / f"adapter_active_ft_seed{spec.train_seed}.json",
|
| 1996 |
+
train_spec=adapter_train_spec,
|
| 1997 |
+
dataset_seed=spec.dataset_seed,
|
| 1998 |
+
planner_overrides=planner_overrides,
|
| 1999 |
+
)
|
| 2000 |
+
summary = _summarize_smoke([trunk_result, noop_result, active_result], args.report_dir)
|
| 2001 |
+
print(json.dumps({"phase": "complete", "summary": summary}, indent=2), flush=True)
|
| 2002 |
+
|
| 2003 |
+
|
| 2004 |
+
if __name__ == "__main__":
|
| 2005 |
+
main()
|
code/VLAarchtests4/VLAarchtests2_code/VLAarchtests/code/reveal_vla_bimanual/eval/run_public_benchmark_package.py
ADDED
|
@@ -0,0 +1,369 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
import argparse
|
| 4 |
+
import json
|
| 5 |
+
from pathlib import Path
|
| 6 |
+
from typing import Any
|
| 7 |
+
|
| 8 |
+
import numpy as np
|
| 9 |
+
|
| 10 |
+
from eval.public_benchmark_package import (
|
| 11 |
+
ANCHOR_ROLE,
|
| 12 |
+
DEFAULT_ANCHOR_TOLERANCE,
|
| 13 |
+
DEFAULT_SIGN_OF_LIFE_GAIN,
|
| 14 |
+
DEFAULT_SIGN_OF_LIFE_INTERVENTION,
|
| 15 |
+
DEFAULT_SIGN_OF_LIFE_NON_BASE,
|
| 16 |
+
TARGET_ROLE,
|
| 17 |
+
build_public_eval_protocol,
|
| 18 |
+
build_target_training_spec,
|
| 19 |
+
default_public_benchmark_manifest,
|
| 20 |
+
expected_eval_modes,
|
| 21 |
+
public_benchmark_tracks,
|
| 22 |
+
public_protocol_identity_signature,
|
| 23 |
+
public_track_by_id,
|
| 24 |
+
training_fairness_signature,
|
| 25 |
+
write_default_public_benchmark_manifest,
|
| 26 |
+
)
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
def _load_json(path: str | Path) -> dict[str, Any]:
|
| 30 |
+
with Path(path).open("r", encoding="utf-8") as handle:
|
| 31 |
+
payload = json.load(handle)
|
| 32 |
+
if not isinstance(payload, dict):
|
| 33 |
+
raise TypeError(f"Expected a JSON object in {path!s}, got {type(payload)!r}.")
|
| 34 |
+
return payload
|
| 35 |
+
|
| 36 |
+
|
| 37 |
+
def _normalize_success_samples(payload: dict[str, Any]) -> np.ndarray:
|
| 38 |
+
if "successes" in payload:
|
| 39 |
+
raw = np.asarray(payload["successes"], dtype=np.float32).reshape(-1)
|
| 40 |
+
return raw
|
| 41 |
+
if "success_rate" in payload:
|
| 42 |
+
return np.asarray([float(payload["success_rate"])], dtype=np.float32)
|
| 43 |
+
raise KeyError("Each result payload must include either `successes` or `success_rate`.")
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
def _mean_optional(records: list[dict[str, Any]], key: str) -> float | None:
|
| 47 |
+
values = [float(record[key]) for record in records if key in record]
|
| 48 |
+
if not values:
|
| 49 |
+
return None
|
| 50 |
+
return float(np.mean(values))
|
| 51 |
+
|
| 52 |
+
|
| 53 |
+
def _bootstrap_delta_ci(
|
| 54 |
+
lhs: np.ndarray,
|
| 55 |
+
rhs: np.ndarray,
|
| 56 |
+
*,
|
| 57 |
+
bootstrap_samples: int,
|
| 58 |
+
bootstrap_seed: int,
|
| 59 |
+
) -> tuple[float, float]:
|
| 60 |
+
if lhs.size == 0 or rhs.size == 0:
|
| 61 |
+
return 0.0, 0.0
|
| 62 |
+
rng = np.random.default_rng(int(bootstrap_seed))
|
| 63 |
+
deltas = np.empty(int(bootstrap_samples), dtype=np.float32)
|
| 64 |
+
for index in range(int(bootstrap_samples)):
|
| 65 |
+
lhs_sample = lhs[rng.integers(0, lhs.shape[0], size=lhs.shape[0])]
|
| 66 |
+
rhs_sample = rhs[rng.integers(0, rhs.shape[0], size=rhs.shape[0])]
|
| 67 |
+
deltas[index] = float(lhs_sample.mean() - rhs_sample.mean())
|
| 68 |
+
low, high = np.percentile(deltas, [2.5, 97.5])
|
| 69 |
+
return float(low), float(high)
|
| 70 |
+
|
| 71 |
+
|
| 72 |
+
def _normalize_record(payload: dict[str, Any]) -> dict[str, Any]:
|
| 73 |
+
if "track_id" not in payload:
|
| 74 |
+
raise KeyError("Missing required field `track_id`.")
|
| 75 |
+
if "adapter_mode" not in payload:
|
| 76 |
+
raise KeyError("Missing required field `adapter_mode`.")
|
| 77 |
+
track = public_track_by_id(str(payload["track_id"]))
|
| 78 |
+
success_samples = _normalize_success_samples(payload)
|
| 79 |
+
success_rate = float(payload.get("success_rate", float(success_samples.mean())))
|
| 80 |
+
episodes = int(payload.get("episodes", success_samples.shape[0]))
|
| 81 |
+
record = dict(payload)
|
| 82 |
+
record["track_id"] = track.track_id
|
| 83 |
+
record["suite"] = payload.get("suite", track.suite)
|
| 84 |
+
record["benchmark_task"] = payload.get("benchmark_task", track.benchmark_task)
|
| 85 |
+
record["role"] = payload.get("role", track.role)
|
| 86 |
+
record["adapter_mode"] = str(payload["adapter_mode"])
|
| 87 |
+
record["successes"] = success_samples.tolist()
|
| 88 |
+
record["success_rate"] = success_rate
|
| 89 |
+
record["episodes"] = episodes
|
| 90 |
+
return record
|
| 91 |
+
|
| 92 |
+
|
| 93 |
+
def _validate_protocols(records: list[dict[str, Any]]) -> None:
|
| 94 |
+
by_track: dict[str, list[dict[str, Any]]] = {}
|
| 95 |
+
for record in records:
|
| 96 |
+
by_track.setdefault(record["track_id"], []).append(record)
|
| 97 |
+
for track_id, grouped in by_track.items():
|
| 98 |
+
signatures = []
|
| 99 |
+
for record in grouped:
|
| 100 |
+
protocol = record.get("eval_protocol")
|
| 101 |
+
if protocol is None:
|
| 102 |
+
raise ValueError(
|
| 103 |
+
f"Missing eval_protocol for track {track_id!r}, mode {record['adapter_mode']!r}."
|
| 104 |
+
)
|
| 105 |
+
signatures.append(public_protocol_identity_signature(protocol))
|
| 106 |
+
if any(signature != signatures[0] for signature in signatures[1:]):
|
| 107 |
+
raise ValueError(f"Protocol identity mismatch detected for track {track_id!r}.")
|
| 108 |
+
|
| 109 |
+
|
| 110 |
+
def _validate_training_fairness(records: list[dict[str, Any]]) -> None:
|
| 111 |
+
grouped: dict[tuple[str, str], list[dict[str, Any]]] = {}
|
| 112 |
+
for record in records:
|
| 113 |
+
grouped.setdefault((record["track_id"], record["adapter_mode"]), []).append(record)
|
| 114 |
+
for track in public_benchmark_tracks(TARGET_ROLE):
|
| 115 |
+
trunk_records = grouped.get((track.track_id, "trunk_only_ft"), [])
|
| 116 |
+
active_records = grouped.get((track.track_id, "adapter_active_ft"), [])
|
| 117 |
+
if not trunk_records or not active_records:
|
| 118 |
+
continue
|
| 119 |
+
if len(trunk_records) != len(active_records):
|
| 120 |
+
raise ValueError(
|
| 121 |
+
f"Training fairness mismatch for {track.track_id!r}: different run counts "
|
| 122 |
+
f"between trunk_only_ft ({len(trunk_records)}) and adapter_active_ft ({len(active_records)})."
|
| 123 |
+
)
|
| 124 |
+
if any(record.get("train_spec") is None for record in trunk_records + active_records):
|
| 125 |
+
raise ValueError(
|
| 126 |
+
f"Training fairness mismatch for {track.track_id!r}: missing train_spec on a target-track result."
|
| 127 |
+
)
|
| 128 |
+
trunk_by_seed = {
|
| 129 |
+
int(record["train_spec"]["seed"]): training_fairness_signature(record["train_spec"])
|
| 130 |
+
for record in trunk_records
|
| 131 |
+
}
|
| 132 |
+
active_by_seed = {
|
| 133 |
+
int(record["train_spec"]["seed"]): training_fairness_signature(record["train_spec"])
|
| 134 |
+
for record in active_records
|
| 135 |
+
}
|
| 136 |
+
if set(trunk_by_seed) != set(active_by_seed):
|
| 137 |
+
raise ValueError(f"Training fairness mismatch for {track.track_id!r}: seed sets differ.")
|
| 138 |
+
for seed, trunk_signature in trunk_by_seed.items():
|
| 139 |
+
if trunk_signature != active_by_seed[seed]:
|
| 140 |
+
raise ValueError(
|
| 141 |
+
f"Training fairness mismatch for {track.track_id!r} at seed {seed}: "
|
| 142 |
+
"trunk_only_ft and adapter_active_ft do not share the same data/init signature."
|
| 143 |
+
)
|
| 144 |
+
|
| 145 |
+
|
| 146 |
+
def _aggregate_mode(records: list[dict[str, Any]]) -> dict[str, Any]:
|
| 147 |
+
success_rates = np.asarray([float(record["success_rate"]) for record in records], dtype=np.float32)
|
| 148 |
+
success_samples = np.concatenate(
|
| 149 |
+
[np.asarray(record["successes"], dtype=np.float32).reshape(-1) for record in records],
|
| 150 |
+
axis=0,
|
| 151 |
+
)
|
| 152 |
+
payload: dict[str, Any] = {
|
| 153 |
+
"num_runs": len(records),
|
| 154 |
+
"mean_success": float(success_rates.mean()) if success_rates.size else 0.0,
|
| 155 |
+
"success_samples": success_samples.tolist(),
|
| 156 |
+
}
|
| 157 |
+
for key in (
|
| 158 |
+
"intervention_rate",
|
| 159 |
+
"non_base_selection_rate",
|
| 160 |
+
"steps_to_first_reveal_or_access",
|
| 161 |
+
"steps_to_retrieve",
|
| 162 |
+
"disturbance_proxy",
|
| 163 |
+
):
|
| 164 |
+
mean_value = _mean_optional(records, key)
|
| 165 |
+
if mean_value is not None:
|
| 166 |
+
payload[key] = mean_value
|
| 167 |
+
return payload
|
| 168 |
+
|
| 169 |
+
|
| 170 |
+
def summarize_public_benchmark_package(
|
| 171 |
+
result_payloads: list[dict[str, Any]],
|
| 172 |
+
*,
|
| 173 |
+
bootstrap_samples: int = 2000,
|
| 174 |
+
bootstrap_seed: int = 0,
|
| 175 |
+
allow_partial: bool = False,
|
| 176 |
+
) -> dict[str, Any]:
|
| 177 |
+
records = [_normalize_record(payload) for payload in result_payloads]
|
| 178 |
+
_validate_protocols(records)
|
| 179 |
+
_validate_training_fairness(records)
|
| 180 |
+
|
| 181 |
+
grouped: dict[tuple[str, str], list[dict[str, Any]]] = {}
|
| 182 |
+
for record in records:
|
| 183 |
+
grouped.setdefault((record["track_id"], record["adapter_mode"]), []).append(record)
|
| 184 |
+
|
| 185 |
+
track_summaries: dict[str, Any] = {}
|
| 186 |
+
target_deltas: list[float] = []
|
| 187 |
+
anchor_pass = True
|
| 188 |
+
sign_of_life_tracks: list[str] = []
|
| 189 |
+
ci_above_zero_tracks: list[str] = []
|
| 190 |
+
available_tracks: list[str] = []
|
| 191 |
+
|
| 192 |
+
for track in public_benchmark_tracks():
|
| 193 |
+
track_modes = expected_eval_modes(track.track_id)
|
| 194 |
+
mode_payloads: dict[str, Any] = {}
|
| 195 |
+
missing_modes: list[str] = []
|
| 196 |
+
for mode in track_modes:
|
| 197 |
+
mode_records = grouped.get((track.track_id, mode), [])
|
| 198 |
+
if not mode_records:
|
| 199 |
+
missing_modes.append(mode)
|
| 200 |
+
continue
|
| 201 |
+
mode_payloads[mode] = _aggregate_mode(mode_records)
|
| 202 |
+
if missing_modes:
|
| 203 |
+
if allow_partial:
|
| 204 |
+
continue
|
| 205 |
+
raise ValueError(f"Missing results for track {track.track_id!r}, mode(s) {missing_modes!r}.")
|
| 206 |
+
available_tracks.append(track.track_id)
|
| 207 |
+
|
| 208 |
+
track_summary: dict[str, Any] = {
|
| 209 |
+
"suite": track.suite,
|
| 210 |
+
"benchmark_task": track.benchmark_task,
|
| 211 |
+
"role": track.role,
|
| 212 |
+
"task_family": track.task_family,
|
| 213 |
+
"target_behavior": track.target_behavior,
|
| 214 |
+
"public_source": track.public_source,
|
| 215 |
+
"notes": track.notes,
|
| 216 |
+
"modes": {
|
| 217 |
+
mode: {
|
| 218 |
+
key: value
|
| 219 |
+
for key, value in payload.items()
|
| 220 |
+
if key != "success_samples"
|
| 221 |
+
}
|
| 222 |
+
for mode, payload in mode_payloads.items()
|
| 223 |
+
},
|
| 224 |
+
}
|
| 225 |
+
|
| 226 |
+
if track.role == TARGET_ROLE:
|
| 227 |
+
trunk = mode_payloads["trunk_only_ft"]
|
| 228 |
+
active = mode_payloads["adapter_active_ft"]
|
| 229 |
+
noop = mode_payloads["adapter_noop"]
|
| 230 |
+
delta_active = float(active["mean_success"] - trunk["mean_success"])
|
| 231 |
+
delta_noop = float(noop["mean_success"] - trunk["mean_success"])
|
| 232 |
+
target_deltas.append(delta_active)
|
| 233 |
+
ci_low, ci_high = _bootstrap_delta_ci(
|
| 234 |
+
np.asarray(active["success_samples"], dtype=np.float32),
|
| 235 |
+
np.asarray(trunk["success_samples"], dtype=np.float32),
|
| 236 |
+
bootstrap_samples=bootstrap_samples,
|
| 237 |
+
bootstrap_seed=bootstrap_seed + len(target_deltas),
|
| 238 |
+
)
|
| 239 |
+
sign_of_life = bool(
|
| 240 |
+
float(active.get("intervention_rate", 0.0)) >= DEFAULT_SIGN_OF_LIFE_INTERVENTION
|
| 241 |
+
and float(active.get("non_base_selection_rate", 0.0)) >= DEFAULT_SIGN_OF_LIFE_NON_BASE
|
| 242 |
+
and delta_active >= DEFAULT_SIGN_OF_LIFE_GAIN
|
| 243 |
+
)
|
| 244 |
+
if sign_of_life:
|
| 245 |
+
sign_of_life_tracks.append(track.track_id)
|
| 246 |
+
if ci_low > 0.0:
|
| 247 |
+
ci_above_zero_tracks.append(track.track_id)
|
| 248 |
+
track_summary.update(
|
| 249 |
+
{
|
| 250 |
+
"delta_active_vs_trunk": delta_active,
|
| 251 |
+
"delta_noop_vs_trunk": delta_noop,
|
| 252 |
+
"delta_active_vs_trunk_ci95": [ci_low, ci_high],
|
| 253 |
+
"signs_of_life": sign_of_life,
|
| 254 |
+
}
|
| 255 |
+
)
|
| 256 |
+
else:
|
| 257 |
+
trunk = mode_payloads["trunk_only"]
|
| 258 |
+
active = mode_payloads["adapter_active"]
|
| 259 |
+
noop = mode_payloads["adapter_noop"]
|
| 260 |
+
active_delta = float(active["mean_success"] - trunk["mean_success"])
|
| 261 |
+
noop_delta = float(noop["mean_success"] - trunk["mean_success"])
|
| 262 |
+
within_tolerance = bool(
|
| 263 |
+
abs(active_delta) <= DEFAULT_ANCHOR_TOLERANCE
|
| 264 |
+
and abs(noop_delta) <= DEFAULT_ANCHOR_TOLERANCE
|
| 265 |
+
)
|
| 266 |
+
anchor_pass = anchor_pass and within_tolerance
|
| 267 |
+
track_summary.update(
|
| 268 |
+
{
|
| 269 |
+
"delta_active_vs_trunk": active_delta,
|
| 270 |
+
"delta_noop_vs_trunk": noop_delta,
|
| 271 |
+
"anchor_within_tolerance": within_tolerance,
|
| 272 |
+
}
|
| 273 |
+
)
|
| 274 |
+
|
| 275 |
+
track_summaries[track.track_id] = track_summary
|
| 276 |
+
|
| 277 |
+
headline_pass = bool(
|
| 278 |
+
target_deltas
|
| 279 |
+
and all(delta > 0.0 for delta in target_deltas)
|
| 280 |
+
and len(ci_above_zero_tracks) >= 1
|
| 281 |
+
)
|
| 282 |
+
sign_of_life_pass = len(sign_of_life_tracks) >= 2
|
| 283 |
+
|
| 284 |
+
return {
|
| 285 |
+
"package_name": default_public_benchmark_manifest()["package_name"],
|
| 286 |
+
"tracks": track_summaries,
|
| 287 |
+
"available_tracks": available_tracks,
|
| 288 |
+
"target_macro_average_delta": float(np.mean(target_deltas)) if target_deltas else 0.0,
|
| 289 |
+
"headline_pass": headline_pass,
|
| 290 |
+
"sign_of_life_pass": sign_of_life_pass,
|
| 291 |
+
"sign_of_life_track_count": len(sign_of_life_tracks),
|
| 292 |
+
"sign_of_life_tracks": sign_of_life_tracks,
|
| 293 |
+
"ci_above_zero_tracks": ci_above_zero_tracks,
|
| 294 |
+
"anchor_pass": anchor_pass,
|
| 295 |
+
}
|
| 296 |
+
|
| 297 |
+
|
| 298 |
+
def _write_markdown(output_path: Path, summary: dict[str, Any]) -> None:
|
| 299 |
+
lines = [
|
| 300 |
+
"# Hybrid Public Benchmark Package Summary",
|
| 301 |
+
"",
|
| 302 |
+
f"- package_name: {summary['package_name']}",
|
| 303 |
+
f"- headline_pass: {summary['headline_pass']}",
|
| 304 |
+
f"- sign_of_life_pass: {summary['sign_of_life_pass']}",
|
| 305 |
+
f"- sign_of_life_track_count: {summary['sign_of_life_track_count']}",
|
| 306 |
+
f"- anchor_pass: {summary['anchor_pass']}",
|
| 307 |
+
f"- target_macro_average_delta: {summary['target_macro_average_delta']:.3f}",
|
| 308 |
+
"",
|
| 309 |
+
]
|
| 310 |
+
for track_id, payload in summary["tracks"].items():
|
| 311 |
+
lines.append(f"## {track_id}")
|
| 312 |
+
lines.append(f"- suite: {payload['suite']}")
|
| 313 |
+
lines.append(f"- benchmark_task: {payload['benchmark_task']}")
|
| 314 |
+
lines.append(f"- role: {payload['role']}")
|
| 315 |
+
for mode, mode_payload in payload["modes"].items():
|
| 316 |
+
lines.append(f"- {mode}: mean_success={mode_payload['mean_success']:.3f}, num_runs={mode_payload['num_runs']}")
|
| 317 |
+
if "delta_active_vs_trunk" in payload:
|
| 318 |
+
lines.append(f"- delta_active_vs_trunk: {payload['delta_active_vs_trunk']:.3f}")
|
| 319 |
+
if "delta_noop_vs_trunk" in payload:
|
| 320 |
+
lines.append(f"- delta_noop_vs_trunk: {payload['delta_noop_vs_trunk']:.3f}")
|
| 321 |
+
if "delta_active_vs_trunk_ci95" in payload:
|
| 322 |
+
low, high = payload["delta_active_vs_trunk_ci95"]
|
| 323 |
+
lines.append(f"- delta_active_vs_trunk_ci95: [{low:.3f}, {high:.3f}]")
|
| 324 |
+
if "signs_of_life" in payload:
|
| 325 |
+
lines.append(f"- signs_of_life: {payload['signs_of_life']}")
|
| 326 |
+
if "anchor_within_tolerance" in payload:
|
| 327 |
+
lines.append(f"- anchor_within_tolerance: {payload['anchor_within_tolerance']}")
|
| 328 |
+
lines.append("")
|
| 329 |
+
output_path.parent.mkdir(parents=True, exist_ok=True)
|
| 330 |
+
output_path.write_text("\n".join(lines).rstrip() + "\n", encoding="utf-8")
|
| 331 |
+
|
| 332 |
+
|
| 333 |
+
def _parse_args() -> argparse.Namespace:
|
| 334 |
+
parser = argparse.ArgumentParser(description="Validate and summarize the public benchmark package results.")
|
| 335 |
+
parser.add_argument("--result", action="append", default=[], help="Path to a normalized benchmark result JSON.")
|
| 336 |
+
parser.add_argument("--output-dir", type=Path, default=Path("/workspace") / "reports" / "public_hybrid_benchmark_v1")
|
| 337 |
+
parser.add_argument("--bootstrap-samples", type=int, default=2000)
|
| 338 |
+
parser.add_argument("--bootstrap-seed", type=int, default=0)
|
| 339 |
+
parser.add_argument("--write-default-manifest", type=Path, default=None)
|
| 340 |
+
return parser.parse_args()
|
| 341 |
+
|
| 342 |
+
|
| 343 |
+
def main() -> None:
|
| 344 |
+
args = _parse_args()
|
| 345 |
+
if args.write_default_manifest is not None:
|
| 346 |
+
path = write_default_public_benchmark_manifest(args.write_default_manifest)
|
| 347 |
+
print(json.dumps({"wrote_manifest": str(path)}, indent=2))
|
| 348 |
+
if not args.result:
|
| 349 |
+
return
|
| 350 |
+
if not args.result:
|
| 351 |
+
raise SystemExit("No results provided. Pass one or more --result files or use --write-default-manifest.")
|
| 352 |
+
|
| 353 |
+
payloads = [_load_json(path) for path in args.result]
|
| 354 |
+
summary = summarize_public_benchmark_package(
|
| 355 |
+
payloads,
|
| 356 |
+
bootstrap_samples=args.bootstrap_samples,
|
| 357 |
+
bootstrap_seed=args.bootstrap_seed,
|
| 358 |
+
)
|
| 359 |
+
|
| 360 |
+
args.output_dir.mkdir(parents=True, exist_ok=True)
|
| 361 |
+
json_path = args.output_dir / "public_benchmark_package_summary.json"
|
| 362 |
+
md_path = args.output_dir / "public_benchmark_package_summary.md"
|
| 363 |
+
json_path.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8")
|
| 364 |
+
_write_markdown(md_path, summary)
|
| 365 |
+
print(json.dumps({"summary_json": str(json_path), "summary_md": str(md_path)}, indent=2))
|
| 366 |
+
|
| 367 |
+
|
| 368 |
+
if __name__ == "__main__":
|
| 369 |
+
main()
|
code/VLAarchtests4/VLAarchtests2_code/VLAarchtests/code/reveal_vla_bimanual/eval/run_rlbench_hybrid_smoke.py
ADDED
|
@@ -0,0 +1,269 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
import argparse
|
| 4 |
+
import json
|
| 5 |
+
import os
|
| 6 |
+
import time
|
| 7 |
+
from dataclasses import asdict, dataclass
|
| 8 |
+
from pathlib import Path
|
| 9 |
+
from typing import Any
|
| 10 |
+
|
| 11 |
+
import numpy as np
|
| 12 |
+
|
| 13 |
+
from eval.public_benchmark_package import (
|
| 14 |
+
build_public_eval_protocol,
|
| 15 |
+
build_target_training_spec,
|
| 16 |
+
public_track_by_id,
|
| 17 |
+
)
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
WORKSPACE_ROOT = Path("/workspace")
|
| 21 |
+
DEFAULT_REPORT_DIR = WORKSPACE_ROOT / "reports" / "public_hybrid_benchmark_smoke_v1" / "rlbench2"
|
| 22 |
+
DEFAULT_COPPELIASIM_ROOT = WORKSPACE_ROOT / "assets" / "coppeliasim_v4_1_0"
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
def _configure_runtime_env() -> None:
|
| 26 |
+
coppeliasim_root = Path(os.environ.get("COPPELIASIM_ROOT", str(DEFAULT_COPPELIASIM_ROOT)))
|
| 27 |
+
os.environ.setdefault("COPPELIASIM_ROOT", str(coppeliasim_root))
|
| 28 |
+
os.environ["LD_LIBRARY_PATH"] = f"{coppeliasim_root}:{os.environ.get('LD_LIBRARY_PATH', '')}".rstrip(":")
|
| 29 |
+
os.environ.setdefault("QT_QPA_PLATFORM_PLUGIN_PATH", str(coppeliasim_root))
|
| 30 |
+
os.environ.setdefault("QT_PLUGIN_PATH", str(coppeliasim_root))
|
| 31 |
+
os.environ.setdefault("XDG_RUNTIME_DIR", "/tmp/runtime-root")
|
| 32 |
+
Path(os.environ["XDG_RUNTIME_DIR"]).mkdir(parents=True, exist_ok=True)
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
_configure_runtime_env()
|
| 36 |
+
|
| 37 |
+
from rlbench.action_modes.action_mode import BimanualMoveArmThenGripper # noqa: E402
|
| 38 |
+
from rlbench.action_modes.arm_action_modes import BimanualEndEffectorPoseViaPlanning # noqa: E402
|
| 39 |
+
from rlbench.action_modes.gripper_action_modes import BimanualDiscrete # noqa: E402
|
| 40 |
+
from rlbench.environment import Environment # noqa: E402
|
| 41 |
+
from rlbench.observation_config import CameraConfig, ObservationConfig # noqa: E402
|
| 42 |
+
from rlbench.bimanual_tasks.bimanual_lift_tray import BimanualLiftTray # noqa: E402
|
| 43 |
+
from rlbench.bimanual_tasks.bimanual_put_bottle_in_fridge import BimanualPutBottleInFridge # noqa: E402
|
| 44 |
+
from rlbench.bimanual_tasks.bimanual_take_shoes_out_of_box import BimanualTakeShoesOutOfBox # noqa: E402
|
| 45 |
+
from rlbench.bimanual_tasks.bimanual_straighten_rope import BimanualStraightenRope # noqa: E402
|
| 46 |
+
from rlbench.bimanual_tasks.bimanual_sweep_to_dustpan import BimanualSweepToDustpan # noqa: E402
|
| 47 |
+
from rlbench.bimanual_tasks.bimanual_take_tray_out_of_oven import BimanualTakeTrayOutOfOven # noqa: E402
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
@dataclass(frozen=True)
|
| 51 |
+
class RLBenchTaskSpec:
|
| 52 |
+
track_id: str
|
| 53 |
+
task_cls: type
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
RLBENCH_TASKS: tuple[RLBenchTaskSpec, ...] = (
|
| 57 |
+
RLBenchTaskSpec("rlbench2_put_bottle_in_fridge", BimanualPutBottleInFridge),
|
| 58 |
+
RLBenchTaskSpec("rlbench2_take_out_tray", BimanualTakeTrayOutOfOven),
|
| 59 |
+
RLBenchTaskSpec("rlbench2_take_shoes_out_of_box", BimanualTakeShoesOutOfBox),
|
| 60 |
+
RLBenchTaskSpec("rlbench2_lift_tray", BimanualLiftTray),
|
| 61 |
+
RLBenchTaskSpec("rlbench2_straighten_rope", BimanualStraightenRope),
|
| 62 |
+
RLBenchTaskSpec("rlbench2_sweep_to_dustpan", BimanualSweepToDustpan),
|
| 63 |
+
)
|
| 64 |
+
|
| 65 |
+
|
| 66 |
+
def _camera_configs(resolution: int) -> dict[str, CameraConfig]:
|
| 67 |
+
disabled = CameraConfig(
|
| 68 |
+
rgb=False,
|
| 69 |
+
depth=False,
|
| 70 |
+
point_cloud=False,
|
| 71 |
+
mask=False,
|
| 72 |
+
image_size=(resolution, resolution),
|
| 73 |
+
)
|
| 74 |
+
return {
|
| 75 |
+
"over_shoulder_left": disabled,
|
| 76 |
+
"over_shoulder_right": disabled,
|
| 77 |
+
"overhead": disabled,
|
| 78 |
+
"wrist_right": disabled,
|
| 79 |
+
"wrist_left": disabled,
|
| 80 |
+
"front": disabled,
|
| 81 |
+
}
|
| 82 |
+
|
| 83 |
+
|
| 84 |
+
def _obs_config(resolution: int) -> ObservationConfig:
|
| 85 |
+
return ObservationConfig(
|
| 86 |
+
camera_configs=_camera_configs(resolution),
|
| 87 |
+
joint_velocities=False,
|
| 88 |
+
joint_positions=False,
|
| 89 |
+
joint_forces=False,
|
| 90 |
+
gripper_open=True,
|
| 91 |
+
gripper_pose=True,
|
| 92 |
+
gripper_matrix=False,
|
| 93 |
+
gripper_joint_positions=False,
|
| 94 |
+
gripper_touch_forces=False,
|
| 95 |
+
wrist_camera_matrix=False,
|
| 96 |
+
task_low_dim_state=True,
|
| 97 |
+
record_ignore_collisions=True,
|
| 98 |
+
robot_name="dual_panda",
|
| 99 |
+
)
|
| 100 |
+
|
| 101 |
+
|
| 102 |
+
def _nudge_action(obs: Any, z_delta: float) -> np.ndarray:
|
| 103 |
+
right_pose = np.asarray(obs.right.gripper_pose, dtype=np.float32).reshape(-1)[:7].copy()
|
| 104 |
+
left_pose = np.asarray(obs.left.gripper_pose, dtype=np.float32).reshape(-1)[:7].copy()
|
| 105 |
+
right_pose[2] += np.float32(z_delta)
|
| 106 |
+
left_pose[2] += np.float32(z_delta)
|
| 107 |
+
return np.concatenate(
|
| 108 |
+
[
|
| 109 |
+
right_pose,
|
| 110 |
+
np.asarray([1.0, 1.0], dtype=np.float32),
|
| 111 |
+
left_pose,
|
| 112 |
+
np.asarray([1.0, 1.0], dtype=np.float32),
|
| 113 |
+
],
|
| 114 |
+
axis=0,
|
| 115 |
+
)
|
| 116 |
+
|
| 117 |
+
|
| 118 |
+
def _mode_train_variant(adapter_mode: str) -> str:
|
| 119 |
+
return "adapter_active_ft" if adapter_mode == "adapter_noop" else adapter_mode
|
| 120 |
+
|
| 121 |
+
|
| 122 |
+
def _run_task(
|
| 123 |
+
spec: RLBenchTaskSpec,
|
| 124 |
+
*,
|
| 125 |
+
adapter_mode: str,
|
| 126 |
+
seed: int,
|
| 127 |
+
resolution: int,
|
| 128 |
+
z_delta: float,
|
| 129 |
+
) -> dict[str, Any]:
|
| 130 |
+
track = public_track_by_id(spec.track_id)
|
| 131 |
+
env = Environment(
|
| 132 |
+
action_mode=BimanualMoveArmThenGripper(
|
| 133 |
+
BimanualEndEffectorPoseViaPlanning(),
|
| 134 |
+
BimanualDiscrete(),
|
| 135 |
+
),
|
| 136 |
+
obs_config=_obs_config(resolution),
|
| 137 |
+
headless=True,
|
| 138 |
+
robot_setup="dual_panda",
|
| 139 |
+
)
|
| 140 |
+
started_at = time.time()
|
| 141 |
+
try:
|
| 142 |
+
env.launch()
|
| 143 |
+
task = env.get_task(spec.task_cls)
|
| 144 |
+
descriptions, obs = task.reset()
|
| 145 |
+
action = _nudge_action(obs, z_delta)
|
| 146 |
+
_, reward, terminate = task.step(action)
|
| 147 |
+
right_pose_norm = float(np.linalg.norm(np.asarray(obs.right.gripper_pose, dtype=np.float32)[:3]))
|
| 148 |
+
left_pose_norm = float(np.linalg.norm(np.asarray(obs.left.gripper_pose, dtype=np.float32)[:3]))
|
| 149 |
+
action_translation_norm = float(np.linalg.norm(action[[0, 1, 2, 9, 10, 11]]))
|
| 150 |
+
smoke_pass = bool(
|
| 151 |
+
np.isfinite(right_pose_norm)
|
| 152 |
+
and np.isfinite(left_pose_norm)
|
| 153 |
+
and np.isfinite(action_translation_norm)
|
| 154 |
+
and max(right_pose_norm, left_pose_norm, action_translation_norm) > 0.0
|
| 155 |
+
)
|
| 156 |
+
result = {
|
| 157 |
+
"track_id": track.track_id,
|
| 158 |
+
"suite": track.suite,
|
| 159 |
+
"benchmark_task": track.benchmark_task,
|
| 160 |
+
"adapter_mode": adapter_mode,
|
| 161 |
+
"seed": seed,
|
| 162 |
+
"episodes": 1,
|
| 163 |
+
"successes": [int(smoke_pass)],
|
| 164 |
+
"success_rate": float(smoke_pass),
|
| 165 |
+
"eval_protocol": build_public_eval_protocol(
|
| 166 |
+
track_id=track.track_id,
|
| 167 |
+
eval_mode=adapter_mode,
|
| 168 |
+
seed=seed,
|
| 169 |
+
episodes=1,
|
| 170 |
+
resolution=resolution,
|
| 171 |
+
),
|
| 172 |
+
"train_spec": build_target_training_spec(
|
| 173 |
+
track_id=track.track_id,
|
| 174 |
+
model_variant=_mode_train_variant(adapter_mode),
|
| 175 |
+
seed=seed,
|
| 176 |
+
),
|
| 177 |
+
"smoke_pass": smoke_pass,
|
| 178 |
+
"task_success": bool(float(reward) > 0.0),
|
| 179 |
+
"reward": float(reward),
|
| 180 |
+
"terminate": bool(terminate),
|
| 181 |
+
"description_count": len(descriptions),
|
| 182 |
+
"right_pose_norm": right_pose_norm,
|
| 183 |
+
"left_pose_norm": left_pose_norm,
|
| 184 |
+
"action_translation_norm": action_translation_norm,
|
| 185 |
+
"runtime_sec": float(time.time() - started_at),
|
| 186 |
+
}
|
| 187 |
+
finally:
|
| 188 |
+
env.shutdown()
|
| 189 |
+
return result
|
| 190 |
+
|
| 191 |
+
|
| 192 |
+
def _parse_args() -> argparse.Namespace:
|
| 193 |
+
parser = argparse.ArgumentParser(description="Run minimal RLBench2 smoke checks for the hybrid public benchmark.")
|
| 194 |
+
parser.add_argument("--output-dir", type=Path, default=DEFAULT_REPORT_DIR)
|
| 195 |
+
parser.add_argument("--adapter-mode", type=str, default="trunk_only_ft")
|
| 196 |
+
parser.add_argument("--seed", type=int, default=17)
|
| 197 |
+
parser.add_argument("--resolution", type=int, default=256)
|
| 198 |
+
parser.add_argument("--z-delta", type=float, default=0.01)
|
| 199 |
+
return parser.parse_args()
|
| 200 |
+
|
| 201 |
+
|
| 202 |
+
def main() -> None:
|
| 203 |
+
args = _parse_args()
|
| 204 |
+
args.output_dir.mkdir(parents=True, exist_ok=True)
|
| 205 |
+
|
| 206 |
+
results: list[dict[str, Any]] = []
|
| 207 |
+
failures: list[str] = []
|
| 208 |
+
for spec in RLBENCH_TASKS:
|
| 209 |
+
track = public_track_by_id(spec.track_id)
|
| 210 |
+
try:
|
| 211 |
+
result = _run_task(
|
| 212 |
+
spec,
|
| 213 |
+
adapter_mode=args.adapter_mode,
|
| 214 |
+
seed=args.seed,
|
| 215 |
+
resolution=args.resolution,
|
| 216 |
+
z_delta=args.z_delta,
|
| 217 |
+
)
|
| 218 |
+
except Exception as exc: # pragma: no cover - exercised in live smoke runs only
|
| 219 |
+
result = {
|
| 220 |
+
"track_id": track.track_id,
|
| 221 |
+
"suite": track.suite,
|
| 222 |
+
"benchmark_task": track.benchmark_task,
|
| 223 |
+
"adapter_mode": args.adapter_mode,
|
| 224 |
+
"seed": args.seed,
|
| 225 |
+
"episodes": 1,
|
| 226 |
+
"successes": [0],
|
| 227 |
+
"success_rate": 0.0,
|
| 228 |
+
"eval_protocol": build_public_eval_protocol(
|
| 229 |
+
track_id=track.track_id,
|
| 230 |
+
eval_mode=args.adapter_mode,
|
| 231 |
+
seed=args.seed,
|
| 232 |
+
episodes=1,
|
| 233 |
+
resolution=args.resolution,
|
| 234 |
+
),
|
| 235 |
+
"train_spec": build_target_training_spec(
|
| 236 |
+
track_id=track.track_id,
|
| 237 |
+
model_variant=_mode_train_variant(args.adapter_mode),
|
| 238 |
+
seed=args.seed,
|
| 239 |
+
),
|
| 240 |
+
"smoke_pass": False,
|
| 241 |
+
"error": repr(exc),
|
| 242 |
+
}
|
| 243 |
+
failures.append(track.track_id)
|
| 244 |
+
results.append(result)
|
| 245 |
+
output_path = args.output_dir / f"{track.track_id}__{args.adapter_mode}_seed{args.seed}.json"
|
| 246 |
+
output_path.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8")
|
| 247 |
+
|
| 248 |
+
summary = {
|
| 249 |
+
"suite": "rlbench2",
|
| 250 |
+
"adapter_mode": args.adapter_mode,
|
| 251 |
+
"seed": args.seed,
|
| 252 |
+
"tracks": [result["track_id"] for result in results],
|
| 253 |
+
"smoke_pass_count": sum(int(result.get("smoke_pass", False)) for result in results),
|
| 254 |
+
"task_success_count": sum(int(result.get("task_success", False)) for result in results),
|
| 255 |
+
"failures": failures,
|
| 256 |
+
"result_files": [
|
| 257 |
+
str(args.output_dir / f"{result['track_id']}__{args.adapter_mode}_seed{args.seed}.json")
|
| 258 |
+
for result in results
|
| 259 |
+
],
|
| 260 |
+
}
|
| 261 |
+
summary_path = args.output_dir / "rlbench_hybrid_smoke_summary.json"
|
| 262 |
+
summary_path.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8")
|
| 263 |
+
print(json.dumps({"summary_path": str(summary_path), "smoke_pass_count": summary["smoke_pass_count"]}, indent=2))
|
| 264 |
+
if failures:
|
| 265 |
+
raise SystemExit(1)
|
| 266 |
+
|
| 267 |
+
|
| 268 |
+
if __name__ == "__main__":
|
| 269 |
+
main()
|
third_party/AnyBimanual/third_party/PyRep/.gitignore
ADDED
|
@@ -0,0 +1,102 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Byte-compiled / optimized / DLL files
|
| 2 |
+
__pycache__/
|
| 3 |
+
*.py[cod]
|
| 4 |
+
*$py.class
|
| 5 |
+
|
| 6 |
+
./idea
|
| 7 |
+
.idea/
|
| 8 |
+
|
| 9 |
+
# Distribution / packaging
|
| 10 |
+
.Python
|
| 11 |
+
build/
|
| 12 |
+
develop-eggs/
|
| 13 |
+
dist/
|
| 14 |
+
downloads/
|
| 15 |
+
eggs/
|
| 16 |
+
.eggs/
|
| 17 |
+
lib/
|
| 18 |
+
lib64/
|
| 19 |
+
parts/
|
| 20 |
+
sdist/
|
| 21 |
+
var/
|
| 22 |
+
wheels/
|
| 23 |
+
*.egg-info/
|
| 24 |
+
.installed.cfg
|
| 25 |
+
*.egg
|
| 26 |
+
|
| 27 |
+
# PyInstaller
|
| 28 |
+
# Usually these files are written by a python script from a template
|
| 29 |
+
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
| 30 |
+
*.manifest
|
| 31 |
+
*.spec
|
| 32 |
+
|
| 33 |
+
# Installer logs
|
| 34 |
+
pip-log.txt
|
| 35 |
+
pip-delete-this-directory.txt
|
| 36 |
+
|
| 37 |
+
# Unit test / coverage reports
|
| 38 |
+
htmlcov/
|
| 39 |
+
.tox/
|
| 40 |
+
.coverage
|
| 41 |
+
.coverage.*
|
| 42 |
+
.cache
|
| 43 |
+
nosetests.xml
|
| 44 |
+
coverage.xml
|
| 45 |
+
*.cover
|
| 46 |
+
.hypothesis/
|
| 47 |
+
|
| 48 |
+
# Translations
|
| 49 |
+
*.mo
|
| 50 |
+
*.pot
|
| 51 |
+
|
| 52 |
+
# Django stuff:
|
| 53 |
+
*.log
|
| 54 |
+
local_settings.py
|
| 55 |
+
|
| 56 |
+
# Flask stuff:
|
| 57 |
+
instance/
|
| 58 |
+
.webassets-cache
|
| 59 |
+
|
| 60 |
+
# Scrapy stuff:
|
| 61 |
+
.scrapy
|
| 62 |
+
|
| 63 |
+
# Sphinx documentation
|
| 64 |
+
docs/_build/
|
| 65 |
+
|
| 66 |
+
# PyBuilder
|
| 67 |
+
target/
|
| 68 |
+
|
| 69 |
+
# Jupyter Notebook
|
| 70 |
+
.ipynb_checkpoints
|
| 71 |
+
|
| 72 |
+
# pyenv
|
| 73 |
+
.python-version
|
| 74 |
+
|
| 75 |
+
# celery beat schedule file
|
| 76 |
+
celerybeat-schedule
|
| 77 |
+
|
| 78 |
+
# SageMath parsed files
|
| 79 |
+
*.sage.py
|
| 80 |
+
|
| 81 |
+
# Environments
|
| 82 |
+
.env
|
| 83 |
+
.venv
|
| 84 |
+
env/
|
| 85 |
+
venv/
|
| 86 |
+
ENV/
|
| 87 |
+
|
| 88 |
+
# Spyder project settings
|
| 89 |
+
.spyderproject
|
| 90 |
+
.spyproject
|
| 91 |
+
|
| 92 |
+
# Rope project settings
|
| 93 |
+
.ropeproject
|
| 94 |
+
|
| 95 |
+
# mkdocs documentation
|
| 96 |
+
/site
|
| 97 |
+
|
| 98 |
+
# mypy
|
| 99 |
+
.mypy_cache/
|
| 100 |
+
|
| 101 |
+
pyrep/backend/_sim_cffi*
|
| 102 |
+
|
third_party/AnyBimanual/third_party/PyRep/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
MIT License
|
| 2 |
+
|
| 3 |
+
Copyright (c) 2019 Stephen James
|
| 4 |
+
|
| 5 |
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
| 6 |
+
of this software and associated documentation files (the "Software"), to deal
|
| 7 |
+
in the Software without restriction, including without limitation the rights
|
| 8 |
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
| 9 |
+
copies of the Software, and to permit persons to whom the Software is
|
| 10 |
+
furnished to do so, subject to the following conditions:
|
| 11 |
+
|
| 12 |
+
The above copyright notice and this permission notice shall be included in all
|
| 13 |
+
copies or substantial portions of the Software.
|
| 14 |
+
|
| 15 |
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
| 16 |
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
| 17 |
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
| 18 |
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
| 19 |
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
| 20 |
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
| 21 |
+
SOFTWARE.
|
third_party/AnyBimanual/third_party/PyRep/README.md
ADDED
|
@@ -0,0 +1,284 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# PyRep [](https://github.com/stepjam/PyRep/actions) [](https://discord.gg/eTMsa5Y)
|
| 2 |
+
|
| 3 |
+
|
| 4 |
+
__PyRep is a toolkit for robot learning research, built on top of [CoppeliaSim](http://www.coppeliarobotics.com/) (previously called V-REP).__
|
| 5 |
+
|
| 6 |
+
|
| 7 |
+
- [Install](#install)
|
| 8 |
+
- [Running Headless](#running-headless)
|
| 9 |
+
- [Getting Started](#getting-started)
|
| 10 |
+
- [Usage](#usage)
|
| 11 |
+
- [Supported Robots](#supported-robots)
|
| 12 |
+
- [Adding Robots](#adding-robots)
|
| 13 |
+
- [Contributing](#contributing)
|
| 14 |
+
- [Projects Using PyRep](#projects-using-pyrep)
|
| 15 |
+
- [What Happened to V-REP?](#what-happened-to-v-rep)
|
| 16 |
+
- [Citation](#citation)
|
| 17 |
+
|
| 18 |
+
|
| 19 |
+
## Install
|
| 20 |
+
|
| 21 |
+
PyRep requires version **4.1** of CoppeliaSim. Download:
|
| 22 |
+
- [Ubuntu 16.04](https://www.coppeliarobotics.com/files/CoppeliaSim_Edu_V4_1_0_Ubuntu16_04.tar.xz)
|
| 23 |
+
- [Ubuntu 18.04](https://www.coppeliarobotics.com/files/CoppeliaSim_Edu_V4_1_0_Ubuntu18_04.tar.xz)
|
| 24 |
+
- [Ubuntu 20.04](https://www.coppeliarobotics.com/files/CoppeliaSim_Edu_V4_1_0_Ubuntu20_04.tar.xz)
|
| 25 |
+
|
| 26 |
+
Once you have downloaded CoppeliaSim, you can pull PyRep from git:
|
| 27 |
+
|
| 28 |
+
```bash
|
| 29 |
+
git clone https://github.com/stepjam/PyRep.git
|
| 30 |
+
cd PyRep
|
| 31 |
+
```
|
| 32 |
+
|
| 33 |
+
Add the following to your *~/.bashrc* file: (__NOTE__: the 'EDIT ME' in the first line)
|
| 34 |
+
|
| 35 |
+
```bash
|
| 36 |
+
export COPPELIASIM_ROOT=EDIT/ME/PATH/TO/COPPELIASIM/INSTALL/DIR
|
| 37 |
+
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$COPPELIASIM_ROOT
|
| 38 |
+
export QT_QPA_PLATFORM_PLUGIN_PATH=$COPPELIASIM_ROOT
|
| 39 |
+
```
|
| 40 |
+
|
| 41 |
+
__Remember to source your bashrc (`source ~/.bashrc`) or
|
| 42 |
+
zshrc (`source ~/.zshrc`) after this.
|
| 43 |
+
|
| 44 |
+
Finally install the python library:
|
| 45 |
+
|
| 46 |
+
```bash
|
| 47 |
+
pip3 install -r requirements.txt
|
| 48 |
+
pip3 install .
|
| 49 |
+
```
|
| 50 |
+
|
| 51 |
+
You should be good to go!
|
| 52 |
+
Try running one of the examples in the *examples/* folder.
|
| 53 |
+
|
| 54 |
+
_Although you can use CoppeliaSim on any platform, communication via PyRep is currently only supported on Linux._
|
| 55 |
+
|
| 56 |
+
#### Troubleshooting
|
| 57 |
+
|
| 58 |
+
Below are some problems you may encounter during installation. If none of these solve your problem, please raise an issue.
|
| 59 |
+
|
| 60 |
+
- ModuleNotFoundError: No module named 'pyrep.backend._v_rep_cffi'
|
| 61 |
+
- If you are getting this error, then please check that you are not running the interpreter from the project root. If you are, then your Python interpreter will try to import those files rather the installed files.
|
| 62 |
+
- error: command 'x86_64-linux-gnu-gcc' failed
|
| 63 |
+
- You may be missing packages needed for building python extensions. Try: `sudo apt-get install python3-dev`, and then re-run the installation.
|
| 64 |
+
|
| 65 |
+
## Running Headless
|
| 66 |
+
|
| 67 |
+
You can run PyRep/CoppeliaSim headlessly with VirtualGL. VirtualGL is an open source toolkit that gives any Unix or Linux remote display software the ability to run OpenGL applications **with full 3D hardware acceleration**.
|
| 68 |
+
First insure that you have the nVidia proprietary driver installed. I.e. you should get an output when running `nvidia-smi`. Now run the following commands:
|
| 69 |
+
```bash
|
| 70 |
+
sudo apt-get install xorg libxcb-randr0-dev libxrender-dev libxkbcommon-dev libxkbcommon-x11-0 libavcodec-dev libavformat-dev libswscale-dev
|
| 71 |
+
sudo nvidia-xconfig -a --use-display-device=None --virtual=1280x1024
|
| 72 |
+
# Install VirtualGL
|
| 73 |
+
wget https://sourceforge.net/projects/virtualgl/files/2.5.2/virtualgl_2.5.2_amd64.deb/download -O virtualgl_2.5.2_amd64.deb
|
| 74 |
+
sudo dpkg -i virtualgl*.deb
|
| 75 |
+
rm virtualgl*.deb
|
| 76 |
+
```
|
| 77 |
+
You will now need to reboot, and then start the X server:
|
| 78 |
+
```bash
|
| 79 |
+
sudo reboot
|
| 80 |
+
nohup sudo X &
|
| 81 |
+
```
|
| 82 |
+
Now we are good to go! To render the application with the first GPU, you can do the following:
|
| 83 |
+
```bash
|
| 84 |
+
export DISPLAY=:0.0
|
| 85 |
+
python my_pyrep_app.py
|
| 86 |
+
```
|
| 87 |
+
To render with the second GPU, you will insetad set display as: `export DISPLAY=:0.1`, and so on.
|
| 88 |
+
|
| 89 |
+
**Acknowledgement**: Special thanks to Boyuan Chen (UC Berkeley) for bringing VirtualGL to my attention!
|
| 90 |
+
|
| 91 |
+
## Getting Started
|
| 92 |
+
|
| 93 |
+
1. First take a look at [Usage](#usage) and the examples in the *examples/* folder to see if PyRep might be able to accelerate your research.
|
| 94 |
+
2. Take a look at the CoppeliaSim [tutorials](http://www.coppeliarobotics.com/helpFiles/en/tutorials.htm).
|
| 95 |
+
|
| 96 |
+
## Usage
|
| 97 |
+
|
| 98 |
+
The best way to see how PyRep can help in your research is to look at the examples in the *examples/* folder!
|
| 99 |
+
|
| 100 |
+
#### Launching the simulation
|
| 101 |
+
|
| 102 |
+
```python
|
| 103 |
+
from pyrep import PyRep
|
| 104 |
+
|
| 105 |
+
pr = PyRep()
|
| 106 |
+
# Launch the application with a scene file in headless mode
|
| 107 |
+
pr.launch('scene.ttt', headless=True)
|
| 108 |
+
pr.start() # Start the simulation
|
| 109 |
+
|
| 110 |
+
# Do some stuff
|
| 111 |
+
|
| 112 |
+
pr.stop() # Stop the simulation
|
| 113 |
+
pr.shutdown() # Close the application
|
| 114 |
+
```
|
| 115 |
+
|
| 116 |
+
|
| 117 |
+
#### Modifying the Scene
|
| 118 |
+
|
| 119 |
+
```python
|
| 120 |
+
from pyrep.objects.shape import Shape
|
| 121 |
+
from pyrep.const import PrimitiveShape
|
| 122 |
+
|
| 123 |
+
object = Shape.create(type=PrimitiveShape.CYLINDER,
|
| 124 |
+
color=[r,g,b], size=[w, h, d],
|
| 125 |
+
position=[x, y, z])
|
| 126 |
+
object.set_color([r, g, b])
|
| 127 |
+
object.set_position([x, y, z])
|
| 128 |
+
```
|
| 129 |
+
|
| 130 |
+
#### Using Robots
|
| 131 |
+
|
| 132 |
+
Robots are designed to be modular; arms are treated separately to grippers.
|
| 133 |
+
|
| 134 |
+
Use the robot ttm files defined in robots/ttms. These have been altered slightly from the original ones shipped with CoppeliaSim to allow them to be used with motional planning out of the box.
|
| 135 |
+
The 'tip' of the robot may not be where you want it, so feel free to play around with this.
|
| 136 |
+
|
| 137 |
+
```python
|
| 138 |
+
from pyrep import PyRep
|
| 139 |
+
from pyrep.robots.arms.panda import Panda
|
| 140 |
+
from pyrep.robots.end_effectors.panda_gripper import PandaGripper
|
| 141 |
+
|
| 142 |
+
pr = PyRep()
|
| 143 |
+
# Launch the application with a scene file that contains a robot
|
| 144 |
+
pr.launch('scene_with_panda.ttt')
|
| 145 |
+
pr.start() # Start the simulation
|
| 146 |
+
|
| 147 |
+
arm = Panda() # Get the panda from the scene
|
| 148 |
+
gripper = PandaGripper() # Get the panda gripper from the scene
|
| 149 |
+
|
| 150 |
+
velocities = [.1, .2, .3, .4, .5, .6, .7]
|
| 151 |
+
arm.set_joint_target_velocities(velocities)
|
| 152 |
+
pr.step() # Step physics simulation
|
| 153 |
+
|
| 154 |
+
done = False
|
| 155 |
+
# Open the gripper halfway at a velocity of 0.04.
|
| 156 |
+
while not done:
|
| 157 |
+
done = gripper.actuate(0.5, velocity=0.04)
|
| 158 |
+
pr.step()
|
| 159 |
+
|
| 160 |
+
pr.stop() # Stop the simulation
|
| 161 |
+
pr.shutdown() # Close the application
|
| 162 |
+
```
|
| 163 |
+
|
| 164 |
+
We recommend constructing your robot in a dictionary or a small structure, e.g.
|
| 165 |
+
|
| 166 |
+
|
| 167 |
+
```python
|
| 168 |
+
class MyRobot(object):
|
| 169 |
+
def __init__(self, arm, gripper):
|
| 170 |
+
self.arm = arm
|
| 171 |
+
self.gripper = gripper
|
| 172 |
+
|
| 173 |
+
arm = Panda() # Get the panda from the scene
|
| 174 |
+
gripper = PandaGripper() # Get the panda gripper from the scene
|
| 175 |
+
|
| 176 |
+
# Create robot structure
|
| 177 |
+
my_robot_1 = MyRobot(arm, gripper)
|
| 178 |
+
# OR
|
| 179 |
+
my_robot_2 = {
|
| 180 |
+
'arm': arm,
|
| 181 |
+
'gripper': gripper
|
| 182 |
+
}
|
| 183 |
+
```
|
| 184 |
+
|
| 185 |
+
#### Running Multiple PyRep Instances
|
| 186 |
+
|
| 187 |
+
Each PyRep instance needs its own process. This can be achieved using Pythons [multiprocessing](https://docs.python.org/3.6/library/multiprocessing.html) module. Here is a simple example:
|
| 188 |
+
|
| 189 |
+
```python
|
| 190 |
+
from multiprocessing import Process
|
| 191 |
+
|
| 192 |
+
PROCESSES = 10
|
| 193 |
+
|
| 194 |
+
def run():
|
| 195 |
+
pr = PyRep()
|
| 196 |
+
pr.launch('my_scene.ttt', headless=True)
|
| 197 |
+
pr.start()
|
| 198 |
+
# Do stuff...
|
| 199 |
+
pr.stop()
|
| 200 |
+
pr.shutdown()
|
| 201 |
+
|
| 202 |
+
processes = [Process(target=run, args=()) for i in range(PROCESSES)]
|
| 203 |
+
[p.start() for p in processes]
|
| 204 |
+
[p.join() for p in processes]
|
| 205 |
+
```
|
| 206 |
+
|
| 207 |
+
## Supported Robots
|
| 208 |
+
|
| 209 |
+
Here is a list of robots currently supported by PyRep:
|
| 210 |
+
|
| 211 |
+
#### Arms
|
| 212 |
+
|
| 213 |
+
- Kinova Mico
|
| 214 |
+
- Kinova Jaco
|
| 215 |
+
- Rethink Baxter
|
| 216 |
+
- Rethink Sawyer
|
| 217 |
+
- Franka Emika Panda
|
| 218 |
+
- Kuka LBR iiwa 7 R800
|
| 219 |
+
- Kuka LBR iiwa 14 R820
|
| 220 |
+
- Universal Robots UR3
|
| 221 |
+
- Universal Robots UR5
|
| 222 |
+
- Universal Robots UR10
|
| 223 |
+
|
| 224 |
+
#### Grippers
|
| 225 |
+
|
| 226 |
+
- Kinova Mico Hand
|
| 227 |
+
- Kinova Jaco Hand
|
| 228 |
+
- Rethink Baxter Gripper
|
| 229 |
+
- Franka Emika Panda Gripper
|
| 230 |
+
|
| 231 |
+
#### Mobile Robots
|
| 232 |
+
|
| 233 |
+
- Kuka YouBot
|
| 234 |
+
- Turtle Bot
|
| 235 |
+
- Line Tracer
|
| 236 |
+
|
| 237 |
+
Feel free to send pull requests for new robots!
|
| 238 |
+
|
| 239 |
+
## Adding Robots
|
| 240 |
+
|
| 241 |
+
If the robot you want is not currently supported, then why not add it in!
|
| 242 |
+
|
| 243 |
+
[Here is a tutorial for adding robots.](tutorials/adding_robots.md)
|
| 244 |
+
|
| 245 |
+
## Contributing
|
| 246 |
+
|
| 247 |
+
We want to make PyRep the best tool for rapid robot learning research. If you would like to get involved, then please [get in contact](https://www.doc.ic.ac.uk/~slj12/)!
|
| 248 |
+
|
| 249 |
+
Pull requests welcome for bug fixes!
|
| 250 |
+
|
| 251 |
+
## Projects Using PyRep
|
| 252 |
+
|
| 253 |
+
If you use PyRep in your work, then get in contact and we can add you to the list!
|
| 254 |
+
|
| 255 |
+
- [RLBench: The Robot Learning Benchmark & Learning Environment, arxiv 2019](https://arxiv.org/abs/1909.12271)
|
| 256 |
+
- [Learning One-Shot Imitation from Humans without Humans, arxiv 2019](https://arxiv.org/abs/1911.01103)
|
| 257 |
+
- [Task-Embedded Control Networks for Few-Shot Imitation Learning, CoRL 2018](https://arxiv.org/abs/1810.03237)
|
| 258 |
+
- [Transferring End-to-End Visuomotor Control from Simulation to Real World for a Multi-Stage Task
|
| 259 |
+
, CoRL 2017](https://arxiv.org/abs/1707.02267)
|
| 260 |
+
|
| 261 |
+
## Acknowledgements
|
| 262 |
+
|
| 263 |
+
- Georges Nomicos (Imperial College London) for the addition of mobile platforms.
|
| 264 |
+
- Boyuan Chen (UC Berkeley) for bringing VirtualGL to my attention.
|
| 265 |
+
|
| 266 |
+
## What Happened to V-REP?
|
| 267 |
+
|
| 268 |
+
Coppelia Robotics discontinued development of __V-REP__. Instead, they now focus
|
| 269 |
+
their efforts on __CoppeliaSim__. CoppeliaSim is 100% compatible with V-REP.
|
| 270 |
+
See more information [here](http://coppeliarobotics.com/helpFiles/en/versionInfo.htm#coppeliaSim4.0.0).
|
| 271 |
+
|
| 272 |
+
PyRep is fully compatible with both V-REP and CoppeliaSim.
|
| 273 |
+
|
| 274 |
+
|
| 275 |
+
## Citation
|
| 276 |
+
|
| 277 |
+
```
|
| 278 |
+
@article{james2019pyrep,
|
| 279 |
+
title={PyRep: Bringing V-REP to Deep Robot Learning},
|
| 280 |
+
author={James, Stephen and Freese, Marc and Davison, Andrew J.},
|
| 281 |
+
journal={arXiv preprint arXiv:1906.11176},
|
| 282 |
+
year={2019}
|
| 283 |
+
}
|
| 284 |
+
```
|
third_party/AnyBimanual/third_party/PyRep/pyrep/__init__.py
ADDED
|
@@ -0,0 +1,5 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
__version__ = '4.1.0.3'
|
| 2 |
+
|
| 3 |
+
testing = False
|
| 4 |
+
|
| 5 |
+
from .pyrep import PyRep
|
third_party/AnyBimanual/third_party/PyRep/pyrep/const.py
ADDED
|
@@ -0,0 +1,110 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from pyrep.backend import sim
|
| 2 |
+
from enum import Enum
|
| 3 |
+
|
| 4 |
+
|
| 5 |
+
class PrimitiveShape(Enum):
|
| 6 |
+
CUBOID = 0
|
| 7 |
+
SPHERE = 1
|
| 8 |
+
CYLINDER = 2
|
| 9 |
+
CONE = 3
|
| 10 |
+
|
| 11 |
+
|
| 12 |
+
class ObjectType(Enum):
|
| 13 |
+
ALL = sim.sim_handle_all
|
| 14 |
+
SHAPE = sim.sim_object_shape_type
|
| 15 |
+
JOINT = sim.sim_object_joint_type
|
| 16 |
+
DUMMY = sim.sim_object_dummy_type
|
| 17 |
+
PROXIMITY_SENSOR = sim.sim_object_proximitysensor_type
|
| 18 |
+
GRAPH = sim.sim_object_graph_type
|
| 19 |
+
CAMERA = sim.sim_object_camera_type
|
| 20 |
+
PATH = sim.sim_object_path_type
|
| 21 |
+
VISION_SENSOR = sim.sim_object_visionsensor_type
|
| 22 |
+
VOLUME = sim.sim_object_volume_type
|
| 23 |
+
MILl = sim.sim_object_mill_type
|
| 24 |
+
FORCE_SENSOR = sim.sim_object_forcesensor_type
|
| 25 |
+
LIGHT = sim.sim_object_light_type
|
| 26 |
+
MIRROR = sim.sim_object_mirror_type
|
| 27 |
+
OCTREE = sim.sim_object_octree_type
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
class JointType(Enum):
|
| 31 |
+
REVOLUTE = sim.sim_joint_revolute_subtype
|
| 32 |
+
PRISMATIC = sim.sim_joint_prismatic_subtype
|
| 33 |
+
SPHERICAL = sim.sim_joint_spherical_subtype
|
| 34 |
+
|
| 35 |
+
|
| 36 |
+
class JointMode(Enum):
|
| 37 |
+
PASSIVE = sim.sim_jointmode_passive
|
| 38 |
+
IK = sim.sim_jointmode_ik
|
| 39 |
+
IK_DEPENDENT = sim.sim_jointmode_ikdependent
|
| 40 |
+
DEPENDENT = sim.sim_jointmode_dependent
|
| 41 |
+
FORCE = sim.sim_jointmode_force
|
| 42 |
+
|
| 43 |
+
|
| 44 |
+
class ConfigurationPathAlgorithms(Enum):
|
| 45 |
+
BiTRRT = 'BiTRRT'
|
| 46 |
+
BITstar = 'BITstar'
|
| 47 |
+
BKPIECE1 = 'BKPIECE1'
|
| 48 |
+
CForest = 'CForest'
|
| 49 |
+
EST = 'EST'
|
| 50 |
+
FMT = 'FMT'
|
| 51 |
+
KPIECE1 = 'KPIECE1'
|
| 52 |
+
LazyPRM = 'LazyPRM'
|
| 53 |
+
LazyPRMstar = 'LazyPRMstar'
|
| 54 |
+
LazyRRT = 'LazyRRT'
|
| 55 |
+
LBKPIECE1 = 'LBKPIECE1'
|
| 56 |
+
LBTRRT = 'LBTRRT'
|
| 57 |
+
PDST = 'PDST'
|
| 58 |
+
PRM = 'PRM'
|
| 59 |
+
PRMstar = 'PRMstar'
|
| 60 |
+
pRRT = 'pRRT'
|
| 61 |
+
pSBL = 'pSBL'
|
| 62 |
+
RRT = 'RRT'
|
| 63 |
+
RRTConnect = 'RRTConnect'
|
| 64 |
+
RRTstar = 'RRTstar'
|
| 65 |
+
SBL = 'SBL'
|
| 66 |
+
SPARS = 'SPARS'
|
| 67 |
+
SPARStwo = 'SPARStwo'
|
| 68 |
+
STRIDE = 'STRIDE'
|
| 69 |
+
TRRT = 'TRRT'
|
| 70 |
+
|
| 71 |
+
|
| 72 |
+
class TextureMappingMode(Enum):
|
| 73 |
+
PLANE = sim.sim_texturemap_plane
|
| 74 |
+
CYLINDER = sim.sim_texturemap_cylinder
|
| 75 |
+
SPHERE = sim.sim_texturemap_sphere
|
| 76 |
+
CUBE = sim.sim_texturemap_cube
|
| 77 |
+
|
| 78 |
+
|
| 79 |
+
class PerspectiveMode(Enum):
|
| 80 |
+
ORTHOGRAPHIC = 0
|
| 81 |
+
PERSPECTIVE = 1
|
| 82 |
+
|
| 83 |
+
|
| 84 |
+
class RenderMode(Enum):
|
| 85 |
+
OPENGL = sim.sim_rendermode_opengl
|
| 86 |
+
OPENGL_AUXILIARY = sim.sim_rendermode_auxchannels
|
| 87 |
+
OPENGL_COLOR_CODED = sim.sim_rendermode_colorcoded
|
| 88 |
+
POV_RAY = sim.sim_rendermode_povray
|
| 89 |
+
EXTERNAL = sim.sim_rendermode_extrenderer
|
| 90 |
+
EXTERNAL_WINDOWED = sim.sim_rendermode_extrendererwindowed
|
| 91 |
+
OPENGL3 = sim.sim_rendermode_opengl3
|
| 92 |
+
OPENGL3_WINDOWED = sim.sim_rendermode_opengl3windowed
|
| 93 |
+
|
| 94 |
+
|
| 95 |
+
class Verbosity(Enum):
|
| 96 |
+
NONE = 'none'
|
| 97 |
+
ERRORS = 'errors'
|
| 98 |
+
WARNINGS = 'warnings'
|
| 99 |
+
LOAD_INFOS = 'loadinfos'
|
| 100 |
+
SCRIPT_ERRORS = 'scripterrors'
|
| 101 |
+
SCRIPT_WARNINGS = 'scriptwarnings'
|
| 102 |
+
SCRIPT_INFOS = 'scriptinfos'
|
| 103 |
+
INFOS = 'infos'
|
| 104 |
+
DEBUG = 'debug'
|
| 105 |
+
TRACE = 'trace'
|
| 106 |
+
TRACE_LUA = 'tracelua'
|
| 107 |
+
TYRACE_ALL = 'traceall'
|
| 108 |
+
|
| 109 |
+
|
| 110 |
+
PYREP_SCRIPT_TYPE = sim.sim_scripttype_addonscript
|
third_party/AnyBimanual/third_party/PyRep/pyrep/errors.py
ADDED
|
@@ -0,0 +1,30 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
class PyRepError(Exception):
|
| 2 |
+
pass
|
| 3 |
+
|
| 4 |
+
|
| 5 |
+
class WrongObjectTypeError(Exception):
|
| 6 |
+
pass
|
| 7 |
+
|
| 8 |
+
|
| 9 |
+
class ObjectAlreadyRemovedError(Exception):
|
| 10 |
+
pass
|
| 11 |
+
|
| 12 |
+
|
| 13 |
+
class ObjectIsNotModelError(Exception):
|
| 14 |
+
pass
|
| 15 |
+
|
| 16 |
+
|
| 17 |
+
class ConfigurationError(Exception):
|
| 18 |
+
pass
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
class ConfigurationPathError(Exception):
|
| 22 |
+
pass
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
class GripperError(Exception):
|
| 26 |
+
pass
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
class IKError(Exception):
|
| 30 |
+
pass
|
third_party/AnyBimanual/third_party/PyRep/pyrep/misc/__init__.py
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
"""Compatibility placeholder for PyRep packaging."""
|
third_party/AnyBimanual/third_party/PyRep/pyrep/pyrep.py
ADDED
|
@@ -0,0 +1,332 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import numpy as np
|
| 2 |
+
from contextlib import contextmanager
|
| 3 |
+
from pyrep.backend import sim, utils
|
| 4 |
+
from pyrep.const import Verbosity
|
| 5 |
+
from pyrep.objects.object import Object
|
| 6 |
+
from pyrep.objects.shape import Shape
|
| 7 |
+
from pyrep.textures.texture import Texture
|
| 8 |
+
from pyrep.errors import PyRepError
|
| 9 |
+
from pyrep.backend import sim
|
| 10 |
+
import os
|
| 11 |
+
import sys
|
| 12 |
+
import time
|
| 13 |
+
import threading
|
| 14 |
+
from typing import Tuple, List
|
| 15 |
+
import warnings
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
class PyRep(object):
|
| 19 |
+
"""Used for interfacing with the CoppeliaSim simulation.
|
| 20 |
+
|
| 21 |
+
Can be used for starting, stopping, and stepping the simulation. As well
|
| 22 |
+
as getting, and creating scene objects and robots.
|
| 23 |
+
"""
|
| 24 |
+
|
| 25 |
+
def __init__(self):
|
| 26 |
+
self.running = False
|
| 27 |
+
self._process = None
|
| 28 |
+
self._robot_to_count = {}
|
| 29 |
+
self.connected = False
|
| 30 |
+
|
| 31 |
+
self._ui_thread = None
|
| 32 |
+
self._responsive_ui_thread = None
|
| 33 |
+
|
| 34 |
+
self._init_thread_id = None
|
| 35 |
+
self._shutting_down = False
|
| 36 |
+
|
| 37 |
+
self._handles_to_objects = {}
|
| 38 |
+
|
| 39 |
+
if 'COPPELIASIM_ROOT' not in os.environ:
|
| 40 |
+
raise PyRepError(
|
| 41 |
+
'COPPELIASIM_ROOT not defined. See installation instructions.')
|
| 42 |
+
self._vrep_root = os.environ['COPPELIASIM_ROOT']
|
| 43 |
+
if not os.path.exists(self._vrep_root):
|
| 44 |
+
raise PyRepError(
|
| 45 |
+
'COPPELIASIM_ROOT was not a correct path. '
|
| 46 |
+
'See installation instructions')
|
| 47 |
+
|
| 48 |
+
def _run_ui_thread(self, scene_file: str, headless: bool,
|
| 49 |
+
verbosity: Verbosity) -> None:
|
| 50 |
+
# Need this otherwise extensions will not be loaded
|
| 51 |
+
os.chdir(self._vrep_root)
|
| 52 |
+
options = sim.sim_gui_headless if headless else sim.sim_gui_all
|
| 53 |
+
sim.simSetStringParameter(
|
| 54 |
+
sim.sim_stringparam_verbosity, verbosity.value)
|
| 55 |
+
sim.simExtLaunchUIThread(
|
| 56 |
+
options=options, scene=scene_file, pyrep_root=self._vrep_root)
|
| 57 |
+
|
| 58 |
+
def _run_responsive_ui_thread(self) -> None:
|
| 59 |
+
while True:
|
| 60 |
+
if not self.running:
|
| 61 |
+
with utils.step_lock:
|
| 62 |
+
if self._shutting_down or sim.simExtGetExitRequest():
|
| 63 |
+
break
|
| 64 |
+
sim.simExtStep(False)
|
| 65 |
+
time.sleep(0.01)
|
| 66 |
+
# If the exit request was from the UI, then call shutdown, otherwise
|
| 67 |
+
# shutdown caused this thread to terminate.
|
| 68 |
+
if not self._shutting_down:
|
| 69 |
+
self.shutdown()
|
| 70 |
+
|
| 71 |
+
def launch(self, scene_file: str = "", headless: bool = False,
|
| 72 |
+
responsive_ui: bool = False, blocking: bool = False,
|
| 73 |
+
verbosity: Verbosity = Verbosity.NONE) -> None:
|
| 74 |
+
"""Launches CoppeliaSim.
|
| 75 |
+
|
| 76 |
+
Launches the UI thread, waits until the UI thread has finished, this
|
| 77 |
+
results in the current thread becoming the simulation thread.
|
| 78 |
+
|
| 79 |
+
:param scene_file: The scene file to load. Empty string for empty scene.
|
| 80 |
+
:param headless: Run CoppeliaSim in simulation mode.
|
| 81 |
+
:param responsive_ui: If True, then a separate thread will be created to
|
| 82 |
+
asynchronously step the UI of CoppeliaSim. Note, that will reduce
|
| 83 |
+
the responsiveness of the simulation thread.
|
| 84 |
+
:param blocking: Causes CoppeliaSim to launch as if running the default
|
| 85 |
+
c++ client application. This is causes the function to block.
|
| 86 |
+
For most users, this will be set to False.
|
| 87 |
+
:param verbosity: The verbosity level for CoppeliaSim.
|
| 88 |
+
Usually Verbosity.NONE or Verbosity.LOAD_INFOS.
|
| 89 |
+
"""
|
| 90 |
+
abs_scene_file = os.path.abspath(scene_file)
|
| 91 |
+
if len(scene_file) > 0 and not os.path.isfile(abs_scene_file):
|
| 92 |
+
raise PyRepError('Scene file does not exist: %s' % scene_file)
|
| 93 |
+
cwd = os.getcwd()
|
| 94 |
+
self._ui_thread = threading.Thread(
|
| 95 |
+
target=self._run_ui_thread,
|
| 96 |
+
args=(abs_scene_file, headless, verbosity))
|
| 97 |
+
self._ui_thread.daemon = True
|
| 98 |
+
self._ui_thread.start()
|
| 99 |
+
|
| 100 |
+
while not sim.simExtCanInitSimThread():
|
| 101 |
+
time.sleep(0.1)
|
| 102 |
+
|
| 103 |
+
sim.simExtSimThreadInit()
|
| 104 |
+
time.sleep(0.2) # Stops CoppeliaSim crashing if restarted too quickly.
|
| 105 |
+
|
| 106 |
+
if blocking:
|
| 107 |
+
while not sim.simExtGetExitRequest():
|
| 108 |
+
sim.simExtStep()
|
| 109 |
+
self.shutdown()
|
| 110 |
+
elif responsive_ui:
|
| 111 |
+
self._responsive_ui_thread = threading.Thread(
|
| 112 |
+
target=self._run_responsive_ui_thread)
|
| 113 |
+
self._responsive_ui_thread.daemon = True
|
| 114 |
+
try:
|
| 115 |
+
self._responsive_ui_thread.start()
|
| 116 |
+
except (KeyboardInterrupt, SystemExit):
|
| 117 |
+
if not self._shutting_down:
|
| 118 |
+
self.shutdown()
|
| 119 |
+
sys.exit()
|
| 120 |
+
self.step()
|
| 121 |
+
else:
|
| 122 |
+
self.step()
|
| 123 |
+
os.chdir(cwd) # Go back to the previous cwd
|
| 124 |
+
|
| 125 |
+
def script_call(self, function_name_at_script_name: str,
|
| 126 |
+
script_handle_or_type: int,
|
| 127 |
+
ints=(), floats=(), strings=(), bytes='') -> (
|
| 128 |
+
Tuple[List[int], List[float], List[str], str]):
|
| 129 |
+
"""Calls a script function (from a plugin, the main client application,
|
| 130 |
+
or from another script). This represents a callback inside of a script.
|
| 131 |
+
|
| 132 |
+
:param function_name_at_script_name: A string representing the function
|
| 133 |
+
name and script name, e.g. myFunctionName@theScriptName. When the
|
| 134 |
+
script is not associated with an object, then just specify the
|
| 135 |
+
function name.
|
| 136 |
+
:param script_handle_or_type: The handle of the script, otherwise the
|
| 137 |
+
type of the script.
|
| 138 |
+
:param ints: The input ints to the script.
|
| 139 |
+
:param floats: The input floats to the script.
|
| 140 |
+
:param strings: The input strings to the script.
|
| 141 |
+
:param bytes: The input bytes to the script (as a string).
|
| 142 |
+
:return: Any number of return values from the called Lua function.
|
| 143 |
+
"""
|
| 144 |
+
return utils.script_call(
|
| 145 |
+
function_name_at_script_name, script_handle_or_type, ints, floats,
|
| 146 |
+
strings, bytes)
|
| 147 |
+
|
| 148 |
+
def shutdown(self) -> None:
|
| 149 |
+
"""Shuts down the CoppeliaSim simulation.
|
| 150 |
+
"""
|
| 151 |
+
if self._ui_thread is None:
|
| 152 |
+
raise PyRepError(
|
| 153 |
+
'CoppeliaSim has not been launched. Call launch first.')
|
| 154 |
+
if self._ui_thread is not None:
|
| 155 |
+
self._shutting_down = True
|
| 156 |
+
self.stop()
|
| 157 |
+
self.step_ui()
|
| 158 |
+
sim.simExtPostExitRequest()
|
| 159 |
+
sim.simExtSimThreadDestroy()
|
| 160 |
+
self._ui_thread.join()
|
| 161 |
+
if self._responsive_ui_thread is not None:
|
| 162 |
+
self._responsive_ui_thread.join()
|
| 163 |
+
# CoppeliaSim crashes if new instance opened too quickly after shutdown.
|
| 164 |
+
# TODO: A small sleep stops this for now.
|
| 165 |
+
time.sleep(0.1)
|
| 166 |
+
self._ui_thread = None
|
| 167 |
+
self._shutting_down = False
|
| 168 |
+
|
| 169 |
+
def start(self) -> None:
|
| 170 |
+
"""Starts the physics simulation if it is not already running.
|
| 171 |
+
"""
|
| 172 |
+
if self._ui_thread is None:
|
| 173 |
+
raise PyRepError(
|
| 174 |
+
'CoppeliaSim has not been launched. Call launch first.')
|
| 175 |
+
if not self.running:
|
| 176 |
+
sim.simStartSimulation()
|
| 177 |
+
self.running = True
|
| 178 |
+
|
| 179 |
+
def stop(self) -> None:
|
| 180 |
+
"""Stops the physics simulation if it is running.
|
| 181 |
+
"""
|
| 182 |
+
if self._ui_thread is None:
|
| 183 |
+
raise PyRepError(
|
| 184 |
+
'CoppeliaSim has not been launched. Call launch first.')
|
| 185 |
+
if self.running:
|
| 186 |
+
sim.simStopSimulation()
|
| 187 |
+
self.running = False
|
| 188 |
+
# Need this so the UI updates
|
| 189 |
+
[self.step() for _ in range(5)] # type: ignore
|
| 190 |
+
|
| 191 |
+
def step(self) -> None:
|
| 192 |
+
"""Execute the next simulation step.
|
| 193 |
+
|
| 194 |
+
If the physics simulation is not running, then this will only update
|
| 195 |
+
the UI.
|
| 196 |
+
"""
|
| 197 |
+
with utils.step_lock:
|
| 198 |
+
sim.simExtStep()
|
| 199 |
+
|
| 200 |
+
def step_ui(self) -> None:
|
| 201 |
+
"""Update the UI.
|
| 202 |
+
|
| 203 |
+
This will not execute the next simulation step, even if the physics
|
| 204 |
+
simulation is running.
|
| 205 |
+
This is only applicable when PyRep was launched without a responsive UI.
|
| 206 |
+
"""
|
| 207 |
+
with utils.step_lock:
|
| 208 |
+
sim.simExtStep(False)
|
| 209 |
+
|
| 210 |
+
def set_simulation_timestep(self, dt: float) -> None:
|
| 211 |
+
"""Sets the simulation time step. Default is 0.05.
|
| 212 |
+
|
| 213 |
+
:param dt: The time step value in seconds.
|
| 214 |
+
"""
|
| 215 |
+
sim.simSetFloatParameter(sim.sim_floatparam_simulation_time_step, dt)
|
| 216 |
+
if not np.allclose(self.get_simulation_timestep(), dt):
|
| 217 |
+
warnings.warn('Could not change simulation timestep. You may need '
|
| 218 |
+
'to change it to "custom dt" using simulation '
|
| 219 |
+
'settings dialog.')
|
| 220 |
+
|
| 221 |
+
def get_simulation_timestep(self) -> float:
|
| 222 |
+
"""Gets the simulation time step.
|
| 223 |
+
|
| 224 |
+
:return: The time step value in seconds.
|
| 225 |
+
"""
|
| 226 |
+
return sim.simGetSimulationTimeStep()
|
| 227 |
+
|
| 228 |
+
def set_configuration_tree(self, config_tree: bytes) -> None:
|
| 229 |
+
"""Restores configuration information previously retrieved.
|
| 230 |
+
|
| 231 |
+
Configuration information (object relative positions/orientations,
|
| 232 |
+
joint/path values) can be retrieved with
|
| 233 |
+
:py:meth:`Object.get_configuration_tree`. Dynamically simulated
|
| 234 |
+
objects will implicitly be reset before the command is applied
|
| 235 |
+
(i.e. similar to calling :py:meth:`Object.reset_dynamic_object` just
|
| 236 |
+
before).
|
| 237 |
+
|
| 238 |
+
:param config_tree: The configuration tree to restore.
|
| 239 |
+
"""
|
| 240 |
+
sim.simSetConfigurationTree(config_tree)
|
| 241 |
+
|
| 242 |
+
def group_objects(self, objects: List[Shape]) -> Shape:
|
| 243 |
+
"""Groups several shapes into a compound shape (or simple shape).
|
| 244 |
+
|
| 245 |
+
:param objects: The list of shapes to group.
|
| 246 |
+
:return: A single grouped shape.
|
| 247 |
+
"""
|
| 248 |
+
handles = [o.get_handle() for o in objects]
|
| 249 |
+
handle = sim.simGroupShapes(handles)
|
| 250 |
+
return Shape(handle)
|
| 251 |
+
|
| 252 |
+
def merge_objects(self, objects: List[Shape]) -> Shape:
|
| 253 |
+
"""Merges several shapes into a compound shape (or simple shape).
|
| 254 |
+
|
| 255 |
+
:param objects: The list of shapes to group.
|
| 256 |
+
:return: A single merged shape.
|
| 257 |
+
"""
|
| 258 |
+
handles = [o.get_handle() for o in objects]
|
| 259 |
+
# FIXME: sim.simGroupShapes(merge=True) won't return correct handle,
|
| 260 |
+
# so we use name to find correct handle of the merged shape.
|
| 261 |
+
name = objects[-1].get_name()
|
| 262 |
+
sim.simGroupShapes(handles, merge=True)
|
| 263 |
+
return Shape(name)
|
| 264 |
+
|
| 265 |
+
def export_scene(self, filename: str) -> None:
|
| 266 |
+
"""Saves the current scene.
|
| 267 |
+
|
| 268 |
+
:param filename: scene filename. The filename extension is required
|
| 269 |
+
("ttt").
|
| 270 |
+
"""
|
| 271 |
+
sim.simSaveScene(filename)
|
| 272 |
+
|
| 273 |
+
def import_model(self, filename: str) -> Object:
|
| 274 |
+
""" Loads a previously saved model.
|
| 275 |
+
|
| 276 |
+
:param filename: model filename. The filename extension is required
|
| 277 |
+
("ttm"). An optional "@copy" can be appended to the filename, in
|
| 278 |
+
which case the model's objects will be named/renamed as if an
|
| 279 |
+
associated script was attached to the model.
|
| 280 |
+
:return: The imported model.
|
| 281 |
+
"""
|
| 282 |
+
handle = sim.simLoadModel(filename)
|
| 283 |
+
return utils.to_type(handle)
|
| 284 |
+
|
| 285 |
+
def create_texture(self, filename: str, interpolate=True, decal_mode=False,
|
| 286 |
+
repeat_along_u=False, repeat_along_v=False
|
| 287 |
+
) -> Tuple[Shape, Texture]:
|
| 288 |
+
"""Creates a planar shape that is textured.
|
| 289 |
+
|
| 290 |
+
:param filename: Path to the texture to load.
|
| 291 |
+
:param interpolate: Adjacent texture pixels are not interpolated.
|
| 292 |
+
:param decal_mode: Texture is applied as a decal (its appearance
|
| 293 |
+
won't be influenced by light conditions).
|
| 294 |
+
:param repeat_along_u: Texture will be repeated along the U direction.
|
| 295 |
+
:param repeat_along_v: Texture will be repeated along the V direction.
|
| 296 |
+
:return: A tuple containing the textured plane and the texture.
|
| 297 |
+
"""
|
| 298 |
+
options = 0
|
| 299 |
+
if not interpolate:
|
| 300 |
+
options |= 1
|
| 301 |
+
if decal_mode:
|
| 302 |
+
options |= 2
|
| 303 |
+
if repeat_along_u:
|
| 304 |
+
options |= 3
|
| 305 |
+
if repeat_along_v:
|
| 306 |
+
options |= 4
|
| 307 |
+
handle = sim.simCreateTexture(filename, options)
|
| 308 |
+
s = Shape(handle)
|
| 309 |
+
return s, s.get_texture()
|
| 310 |
+
|
| 311 |
+
def get_objects_in_tree(self, root_object=None, *args, **kwargs
|
| 312 |
+
) -> List[Object]:
|
| 313 |
+
"""Retrieves the objects in a given hierarchy tree.
|
| 314 |
+
|
| 315 |
+
:param root_object: The root object in the tree. Pass None to retrieve
|
| 316 |
+
all objects in the configuration tree. :py:class:`Object` or `int`.
|
| 317 |
+
:param object_type: The object type to retrieve.
|
| 318 |
+
One of :py:class:`.ObjectType`.
|
| 319 |
+
:param exclude_base: Exclude the tree base from the returned list.
|
| 320 |
+
:param first_generation_only: Include in the returned list only the
|
| 321 |
+
object's first children. Otherwise, entire hierarchy is returned.
|
| 322 |
+
:return: A list of objects in the hierarchy tree.
|
| 323 |
+
"""
|
| 324 |
+
return Object._get_objects_in_tree(root_object, *args, **kwargs)
|
| 325 |
+
|
| 326 |
+
def get_collection_handle_by_name(self, collection_name: str) -> int:
|
| 327 |
+
"""Retrieves the integer handle for a given collection.
|
| 328 |
+
|
| 329 |
+
:param collection_name: Name of the collection to retrieve the integer handle for
|
| 330 |
+
:return: An integer handle for the collection
|
| 331 |
+
"""
|
| 332 |
+
return sim.simGetCollectionHandle(collection_name)
|
third_party/AnyBimanual/third_party/PyRep/pyrep/robots/__init__.py
ADDED
|
File without changes
|
third_party/AnyBimanual/third_party/PyRep/pyrep/robots/end_effectors/xarm_gripper.py
ADDED
|
@@ -0,0 +1,13 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from pyrep.robots.end_effectors.gripper import Gripper
|
| 2 |
+
|
| 3 |
+
|
| 4 |
+
class XArmGripper(Gripper):
|
| 5 |
+
|
| 6 |
+
def __init__(self, count: int = 0):
|
| 7 |
+
super().__init__(count, 'xarm_gripper',
|
| 8 |
+
['xarm_left_inner_knuckle_joint',
|
| 9 |
+
'xarm_right_inner_knuckle_joint',
|
| 10 |
+
'xarm_drive_joint',
|
| 11 |
+
'xarm_left_finger_joint',
|
| 12 |
+
'xarm_right_outer_knuckle_joint',
|
| 13 |
+
'xarm_right_finger_joint'])
|
third_party/AnyBimanual/third_party/PyRep/pyrep/robots/mobiles/holonomic_base.py
ADDED
|
@@ -0,0 +1,215 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from pyrep.robots.mobiles.mobile_base import MobileBase
|
| 2 |
+
from pyrep.robots.configuration_paths.holonomic_configuration_path import (
|
| 3 |
+
HolonomicConfigurationPath)
|
| 4 |
+
from pyrep.backend import utils
|
| 5 |
+
from pyrep.const import PYREP_SCRIPT_TYPE
|
| 6 |
+
from pyrep.const import ConfigurationPathAlgorithms as Algos
|
| 7 |
+
from pyrep.errors import ConfigurationPathError
|
| 8 |
+
from typing import List
|
| 9 |
+
from pyrep.objects.joint import Joint
|
| 10 |
+
from math import pi, sqrt
|
| 11 |
+
|
| 12 |
+
|
| 13 |
+
class HolonomicBase(MobileBase):
|
| 14 |
+
|
| 15 |
+
def __init__(self,
|
| 16 |
+
count: int,
|
| 17 |
+
num_wheels: int,
|
| 18 |
+
distance_from_target: float,
|
| 19 |
+
name: str,
|
| 20 |
+
max_velocity: float = 4,
|
| 21 |
+
max_velocity_rotation: float = 6,
|
| 22 |
+
max_acceleration: float = 0.035):
|
| 23 |
+
"""Init.
|
| 24 |
+
|
| 25 |
+
:param count: used for multiple copies of robots.
|
| 26 |
+
:param num_wheels: number of actuated wheels.
|
| 27 |
+
:param distance_from_target: offset from target.
|
| 28 |
+
:param name: string with robot name (same as base in vrep model).
|
| 29 |
+
:param max_velocity: bounds x,y velocity for motion planning.
|
| 30 |
+
:param max_velocity_rotation: bounds yaw velocity for motion planning.
|
| 31 |
+
:param max_acceleration: bounds acceleration for motion planning.
|
| 32 |
+
"""
|
| 33 |
+
|
| 34 |
+
super().__init__(count, num_wheels, name)
|
| 35 |
+
|
| 36 |
+
suffix = '' if count == 0 else '#%d' % (count - 1)
|
| 37 |
+
|
| 38 |
+
self.paramP = 20
|
| 39 |
+
self.paramO = 10
|
| 40 |
+
# self.paramO = 30
|
| 41 |
+
self.previous_forw_back_vel = 0
|
| 42 |
+
self.previous_left_right_vel = 0
|
| 43 |
+
self.previous_rot_vel = 0
|
| 44 |
+
self.accelF = max_acceleration
|
| 45 |
+
self.maxV = max_velocity
|
| 46 |
+
self.max_v_rot = max_velocity_rotation
|
| 47 |
+
self.dist1 = distance_from_target
|
| 48 |
+
|
| 49 |
+
joint_slipping_names = [
|
| 50 |
+
'%s_slipping_m_joint%s%s' % (name, str(i + 1), suffix) for i in
|
| 51 |
+
range(self.num_wheels)]
|
| 52 |
+
self.joints_slipping = [Joint(jsname)
|
| 53 |
+
for jsname in joint_slipping_names]
|
| 54 |
+
|
| 55 |
+
def set_base_angular_velocites(self, velocity: List[float]):
|
| 56 |
+
"""Calls required functions to achieve desired omnidirectional effect.
|
| 57 |
+
|
| 58 |
+
:param velocity: A List with forwardBackward, leftRight and rotation
|
| 59 |
+
velocity (in radian/s)
|
| 60 |
+
"""
|
| 61 |
+
# self._reset_wheel()
|
| 62 |
+
fBVel = velocity[0]
|
| 63 |
+
lRVel = velocity[1]
|
| 64 |
+
rVel = velocity[2]
|
| 65 |
+
self.set_joint_target_velocities(
|
| 66 |
+
[-fBVel - lRVel - rVel, -fBVel + lRVel - rVel,
|
| 67 |
+
-fBVel - lRVel + rVel, -fBVel + lRVel + rVel])
|
| 68 |
+
|
| 69 |
+
def get_linear_path(self, position: List[float],
|
| 70 |
+
angle=0) -> HolonomicConfigurationPath:
|
| 71 |
+
"""Initialize linear path and check for collision along it.
|
| 72 |
+
|
| 73 |
+
Must specify either rotation in euler or quaternions, but not both!
|
| 74 |
+
|
| 75 |
+
:param position: The x, y position of the target.
|
| 76 |
+
:param angle: The z orientation of the target (in radians).
|
| 77 |
+
:raises: ConfigurationPathError if no path could be created.
|
| 78 |
+
|
| 79 |
+
:return: A linear path in the 2d space.
|
| 80 |
+
"""
|
| 81 |
+
position_base = self.get_position()
|
| 82 |
+
angle_base = self.get_orientation()[-1]
|
| 83 |
+
|
| 84 |
+
self.target_base.set_position(
|
| 85 |
+
[position[0], position[1], self.target_z])
|
| 86 |
+
self.target_base.set_orientation([0, 0, angle])
|
| 87 |
+
|
| 88 |
+
handle_base = self.get_handle()
|
| 89 |
+
handle_target_base = self.target_base.get_handle()
|
| 90 |
+
_, ret_floats, _, _ = utils.script_call(
|
| 91 |
+
'getBoxAdjustedMatrixAndFacingAngle@PyRep', PYREP_SCRIPT_TYPE,
|
| 92 |
+
ints=[handle_base, handle_target_base])
|
| 93 |
+
|
| 94 |
+
m = ret_floats[:-1]
|
| 95 |
+
angle = ret_floats[-1]
|
| 96 |
+
self.intermediate_target_base.set_position(
|
| 97 |
+
[m[3] - m[0] * self.dist1, m[7] - m[4] * self.dist1,
|
| 98 |
+
self.target_z])
|
| 99 |
+
self.intermediate_target_base.set_orientation([0, 0, angle])
|
| 100 |
+
self.target_base.set_orientation([0, 0, angle])
|
| 101 |
+
|
| 102 |
+
path = [[position_base[0], position_base[1], angle_base],
|
| 103 |
+
[position[0], position[1], angle]]
|
| 104 |
+
|
| 105 |
+
if self._check_collision_linear_path(path):
|
| 106 |
+
raise ConfigurationPathError(
|
| 107 |
+
'Could not create path. '
|
| 108 |
+
'An object was detected on the linear path.')
|
| 109 |
+
|
| 110 |
+
return HolonomicConfigurationPath(self, path)
|
| 111 |
+
|
| 112 |
+
def get_nonlinear_path(self, position: List[float],
|
| 113 |
+
angle=0,
|
| 114 |
+
boundaries=2,
|
| 115 |
+
path_pts=600,
|
| 116 |
+
ignore_collisions=False,
|
| 117 |
+
algorithm=Algos.RRTConnect
|
| 118 |
+
) -> HolonomicConfigurationPath:
|
| 119 |
+
"""Gets a non-linear (planned) configuration path given a target pose.
|
| 120 |
+
|
| 121 |
+
:param position: The x, y, z position of the target.
|
| 122 |
+
:param angle: The z orientation of the target (in radians).
|
| 123 |
+
:param boundaries: A float defining the path search in x and y direction
|
| 124 |
+
[[-boundaries,boundaries],[-boundaries,boundaries]].
|
| 125 |
+
:param path_pts: The number of sampled points returned from the
|
| 126 |
+
computed path
|
| 127 |
+
:param ignore_collisions: If collision checking should be disabled.
|
| 128 |
+
:param algorithm: Algorithm used to compute path
|
| 129 |
+
:raises: ConfigurationPathError if no path could be created.
|
| 130 |
+
|
| 131 |
+
:return: A non-linear path (x,y,angle) in the xy configuration space.
|
| 132 |
+
"""
|
| 133 |
+
|
| 134 |
+
path = self._get_nonlinear_path_points(
|
| 135 |
+
position, angle, boundaries, path_pts, ignore_collisions, algorithm)
|
| 136 |
+
|
| 137 |
+
return HolonomicConfigurationPath(self, path)
|
| 138 |
+
|
| 139 |
+
def get_base_actuation(self):
|
| 140 |
+
"""Proportional controller.
|
| 141 |
+
|
| 142 |
+
:return: A list with left and right joint velocity,
|
| 143 |
+
and a bool representing target is reached.
|
| 144 |
+
"""
|
| 145 |
+
|
| 146 |
+
handleBase = self.get_handle()
|
| 147 |
+
handle_inter_target_base = self.intermediate_target_base.get_handle()
|
| 148 |
+
pos_v = self.target_base.get_position(relative_to=self)
|
| 149 |
+
or_v = self.target_base.get_orientation(relative_to=self)
|
| 150 |
+
|
| 151 |
+
pos_inter = self.intermediate_target_base.get_position(
|
| 152 |
+
relative_to=self)
|
| 153 |
+
or_inter = self.intermediate_target_base.get_orientation(
|
| 154 |
+
relative_to=self)
|
| 155 |
+
|
| 156 |
+
if (sqrt((pos_v[0]) ** 2 + (
|
| 157 |
+
pos_v[1]) ** 2) - self.dist1) < 0.001 and or_v[-1] < 0.1 * pi / 180:
|
| 158 |
+
return [self.previous_forw_back_vel, self.previous_left_right_vel,
|
| 159 |
+
self.previous_rot_vel], True
|
| 160 |
+
|
| 161 |
+
forw_back_vel = pos_inter[1] * self.paramP
|
| 162 |
+
left_right_vel = pos_inter[0] * self.paramP
|
| 163 |
+
rot_vel = - or_inter[2] * self.paramO
|
| 164 |
+
|
| 165 |
+
v = sqrt(forw_back_vel * forw_back_vel +
|
| 166 |
+
left_right_vel * left_right_vel)
|
| 167 |
+
if v > self.maxV:
|
| 168 |
+
forw_back_vel = forw_back_vel * self.maxV / v
|
| 169 |
+
left_right_vel = left_right_vel * self.maxV / v
|
| 170 |
+
|
| 171 |
+
if (abs(rot_vel) > self.max_v_rot):
|
| 172 |
+
rot_vel = self.max_v_rot * rot_vel / abs(rot_vel)
|
| 173 |
+
|
| 174 |
+
df = forw_back_vel - self.previous_forw_back_vel
|
| 175 |
+
ds = left_right_vel - self.previous_left_right_vel
|
| 176 |
+
dr = rot_vel - self.previous_rot_vel
|
| 177 |
+
|
| 178 |
+
if abs(df) > self.maxV * self.accelF:
|
| 179 |
+
df = abs(df) * (self.maxV * self.accelF) / df
|
| 180 |
+
|
| 181 |
+
if abs(ds) > self.maxV * self.accelF:
|
| 182 |
+
ds = abs(ds) * (self.maxV * self.accelF) / ds
|
| 183 |
+
|
| 184 |
+
if abs(dr) > self.max_v_rot * self.accelF:
|
| 185 |
+
dr = abs(dr) * (self.max_v_rot * self.accelF) / dr
|
| 186 |
+
|
| 187 |
+
forw_back_vel = self.previous_forw_back_vel + df
|
| 188 |
+
left_right_vel = self.previous_left_right_vel + ds
|
| 189 |
+
rot_vel = self.previous_rot_vel + dr
|
| 190 |
+
|
| 191 |
+
self.previous_forw_back_vel = forw_back_vel
|
| 192 |
+
self.previous_left_right_vel = left_right_vel
|
| 193 |
+
self.previous_rot_vel = rot_vel
|
| 194 |
+
|
| 195 |
+
return [forw_back_vel, left_right_vel, rot_vel], False
|
| 196 |
+
|
| 197 |
+
def _reset_wheel(self):
|
| 198 |
+
"""Required to achieve desired omnidirectional wheel effect.
|
| 199 |
+
"""
|
| 200 |
+
[j.reset_dynamic_object() for j in self.wheels]
|
| 201 |
+
|
| 202 |
+
p = [[-pi / 4, 0, 0], [pi / 4, 0, pi], [-pi / 4, 0, 0], [pi / 4, 0, pi]]
|
| 203 |
+
|
| 204 |
+
for i in range(self.num_wheels):
|
| 205 |
+
self.joints_slipping[i].set_position([0, 0, 0],
|
| 206 |
+
relative_to=self.joints[i],
|
| 207 |
+
reset_dynamics=False)
|
| 208 |
+
self.joints_slipping[i].set_orientation(p[i],
|
| 209 |
+
relative_to=self.joints[i],
|
| 210 |
+
reset_dynamics=False)
|
| 211 |
+
self.wheels[i].set_position([0, 0, 0], relative_to=self.joints[i],
|
| 212 |
+
reset_dynamics=False)
|
| 213 |
+
self.wheels[i].set_orientation([0, 0, 0],
|
| 214 |
+
relative_to=self.joints[i],
|
| 215 |
+
reset_dynamics=False)
|
third_party/AnyBimanual/third_party/PyRep/pyrep/robots/mobiles/line_tracer.py
ADDED
|
@@ -0,0 +1,6 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from pyrep.robots.mobiles.mobile_base import MobileBase
|
| 2 |
+
|
| 3 |
+
|
| 4 |
+
class LineTracer(MobileBase):
|
| 5 |
+
def __init__(self, count: int = 0):
|
| 6 |
+
super().__init__(count, 2, 'LineTracer')
|
third_party/AnyBimanual/third_party/PyRep/pyrep/robots/mobiles/locobot.py
ADDED
|
@@ -0,0 +1,6 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from pyrep.robots.mobiles.nonholonomic_base import NonHolonomicBase
|
| 2 |
+
|
| 3 |
+
|
| 4 |
+
class LoCoBot(NonHolonomicBase):
|
| 5 |
+
def __init__(self, count: int = 0):
|
| 6 |
+
super().__init__(count, 2, 'LoCoBot')
|
third_party/AnyBimanual/third_party/PyRep/pyrep/robots/mobiles/mobile_base.py
ADDED
|
@@ -0,0 +1,202 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from pyrep.backend import sim, utils
|
| 2 |
+
from pyrep.objects.dummy import Dummy
|
| 3 |
+
from pyrep.objects.shape import Shape
|
| 4 |
+
from pyrep.robots.robot_component import RobotComponent
|
| 5 |
+
from pyrep.const import ConfigurationPathAlgorithms as Algos
|
| 6 |
+
from pyrep.errors import ConfigurationPathError
|
| 7 |
+
from pyrep.const import PYREP_SCRIPT_TYPE
|
| 8 |
+
from typing import List, Union
|
| 9 |
+
from math import sqrt
|
| 10 |
+
import numpy as np
|
| 11 |
+
|
| 12 |
+
|
| 13 |
+
class MobileBase(RobotComponent):
|
| 14 |
+
"""Base class representing a robot mobile base with path planning support.
|
| 15 |
+
"""
|
| 16 |
+
|
| 17 |
+
def __init__(self, count: int, num_wheels: int, name: str):
|
| 18 |
+
"""Count is used for when we have multiple copies of mobile bases.
|
| 19 |
+
|
| 20 |
+
:param count: used for multiple copies of robots
|
| 21 |
+
:param num_wheels: number of actuated wheels
|
| 22 |
+
:param name: string with robot name (same as base in vrep model).
|
| 23 |
+
"""
|
| 24 |
+
|
| 25 |
+
joint_names = ['%s_m_joint%s' % (name, str(i + 1)) for i in
|
| 26 |
+
range(num_wheels)]
|
| 27 |
+
super().__init__(count, name, joint_names)
|
| 28 |
+
|
| 29 |
+
self.num_wheels = num_wheels
|
| 30 |
+
suffix = '' if count == 0 else '#%d' % (count - 1)
|
| 31 |
+
|
| 32 |
+
wheel_names = ['%s_wheel%s%s' % (name, str(i + 1), suffix) for i in
|
| 33 |
+
range(self.num_wheels)]
|
| 34 |
+
self.wheels = [Shape(name) for name in wheel_names]
|
| 35 |
+
|
| 36 |
+
# Motion planning handles
|
| 37 |
+
self.intermediate_target_base = Dummy(
|
| 38 |
+
'%s_intermediate_target_base%s' % (name, suffix))
|
| 39 |
+
self.target_base = Dummy('%s_target_base%s' % (name, suffix))
|
| 40 |
+
|
| 41 |
+
self._collision_collection = sim.simGetCollectionHandle(
|
| 42 |
+
'%s_base%s' % (name, suffix))
|
| 43 |
+
|
| 44 |
+
# Robot parameters and handle
|
| 45 |
+
self.z_pos = self.get_position()[2]
|
| 46 |
+
self.target_z = self.target_base.get_position()[-1]
|
| 47 |
+
self.wheel_radius = self.wheels[0].get_bounding_box()[5] # Z
|
| 48 |
+
self.wheel_distance = np.linalg.norm(
|
| 49 |
+
np.array(self.wheels[0].get_position()) -
|
| 50 |
+
np.array(self.wheels[1].get_position()))
|
| 51 |
+
|
| 52 |
+
# Make sure dummies are orphan if loaded with ttm
|
| 53 |
+
self.intermediate_target_base.set_parent(None)
|
| 54 |
+
self.target_base.set_parent(None)
|
| 55 |
+
|
| 56 |
+
def get_2d_pose(self) -> np.ndarray:
|
| 57 |
+
"""Gets the 2D (top-down) pose of the robot [x, y, yaw].
|
| 58 |
+
|
| 59 |
+
:return: A List containing the x, y, yaw (in radians).
|
| 60 |
+
"""
|
| 61 |
+
return np.r_[self.get_position()[:2], self.get_orientation()[-1:]]
|
| 62 |
+
|
| 63 |
+
def set_2d_pose(self, pose: Union[List[float], np.ndarray]) -> None:
|
| 64 |
+
"""Sets the 2D (top-down) pose of the robot [x, y, yaw]
|
| 65 |
+
|
| 66 |
+
:param pose: A List containing the x, y, yaw (in radians).
|
| 67 |
+
"""
|
| 68 |
+
x, y, yaw = pose
|
| 69 |
+
self.set_position([x, y, self.z_pos])
|
| 70 |
+
self.set_orientation([0, 0, yaw])
|
| 71 |
+
|
| 72 |
+
def assess_collision(self):
|
| 73 |
+
"""Silent detection of the robot base with all other entities present in the scene.
|
| 74 |
+
|
| 75 |
+
:return: True if collision is detected
|
| 76 |
+
"""
|
| 77 |
+
return sim.simCheckCollision(self._collision_collection,
|
| 78 |
+
sim.sim_handle_all) == 1
|
| 79 |
+
|
| 80 |
+
def set_cartesian_position(self, position: List[float]):
|
| 81 |
+
"""Set a delta target position (x,y) and rotation position
|
| 82 |
+
|
| 83 |
+
:param position: length 3 list containing x and y position, and angle position
|
| 84 |
+
|
| 85 |
+
NOTE: not supported for nonholonomic mobile bases.
|
| 86 |
+
"""
|
| 87 |
+
vel = [0, 0, 0]
|
| 88 |
+
vel[-1] = position[-1]
|
| 89 |
+
for i in range(2):
|
| 90 |
+
vel[i] = position[i] / (0.05 * self.wheel_radius / 2) # "0.05 is dt"
|
| 91 |
+
|
| 92 |
+
self.set_base_angular_velocites(vel)
|
| 93 |
+
|
| 94 |
+
def set_base_angular_velocites(self, velocity: List[float]):
|
| 95 |
+
"""This function has no effect for two_wheels robot. More control is required for omnidirectional robot.
|
| 96 |
+
|
| 97 |
+
:param velocity: for two wheels robot: each wheel velocity, for omnidirectional robot forwardBackward, leftRight and rotation velocity
|
| 98 |
+
"""
|
| 99 |
+
raise NotImplementedError()
|
| 100 |
+
|
| 101 |
+
def _get_nonlinear_path_points(self, position: List[float],
|
| 102 |
+
angle=0,
|
| 103 |
+
boundaries=2,
|
| 104 |
+
path_pts=600,
|
| 105 |
+
ignore_collisions=False,
|
| 106 |
+
algorithm=Algos.RRTConnect) -> List[List[float]]:
|
| 107 |
+
"""Gets a non-linear (planned) configuration path given a target pose.
|
| 108 |
+
|
| 109 |
+
:param position: The x, y, z position of the target.
|
| 110 |
+
:param angle: The z orientation of the target (in radians).
|
| 111 |
+
:param boundaries: A float defining the path search in x and y direction
|
| 112 |
+
[[-boundaries,boundaries],[-boundaries,boundaries]].
|
| 113 |
+
:param path_pts: number of sampled points returned from the computed path
|
| 114 |
+
:param ignore_collisions: If collision checking should be disabled.
|
| 115 |
+
:param algorithm: Algorithm used to compute path
|
| 116 |
+
:raises: ConfigurationPathError if no path could be created.
|
| 117 |
+
|
| 118 |
+
:return: A non-linear path (x,y,angle) in the xy configuration space.
|
| 119 |
+
"""
|
| 120 |
+
|
| 121 |
+
# Base dummy required to be parent of the robot tree
|
| 122 |
+
# self.base_ref.set_parent(None)
|
| 123 |
+
# self.set_parent(self.base_ref)
|
| 124 |
+
|
| 125 |
+
# Missing the dist1 for intermediate target
|
| 126 |
+
|
| 127 |
+
self.target_base.set_position([position[0], position[1], self.target_z])
|
| 128 |
+
self.target_base.set_orientation([0, 0, angle])
|
| 129 |
+
|
| 130 |
+
handle_base = self.get_handle()
|
| 131 |
+
handle_target_base = self.target_base.get_handle()
|
| 132 |
+
|
| 133 |
+
_, ret_floats, _, _ = utils.script_call(
|
| 134 |
+
'getNonlinearPathMobile@PyRep', PYREP_SCRIPT_TYPE,
|
| 135 |
+
ints=[handle_base, handle_target_base,
|
| 136 |
+
self._collision_collection,
|
| 137 |
+
int(ignore_collisions), path_pts], floats=[boundaries],
|
| 138 |
+
strings=[algorithm.value])
|
| 139 |
+
|
| 140 |
+
# self.set_parent(None)
|
| 141 |
+
# self.base_ref.set_parent(self)
|
| 142 |
+
|
| 143 |
+
if len(ret_floats) == 0:
|
| 144 |
+
raise ConfigurationPathError('Could not create path.')
|
| 145 |
+
|
| 146 |
+
path = []
|
| 147 |
+
for i in range(0, len(ret_floats) // 3):
|
| 148 |
+
inst = ret_floats[3 * i:3 * i + 3]
|
| 149 |
+
if i > 0:
|
| 150 |
+
dist_change = sqrt((inst[0] - prev_inst[0]) ** 2 + (
|
| 151 |
+
inst[1] - prev_inst[1]) ** 2)
|
| 152 |
+
else:
|
| 153 |
+
dist_change = 0
|
| 154 |
+
inst.append(dist_change)
|
| 155 |
+
|
| 156 |
+
path.append(inst)
|
| 157 |
+
|
| 158 |
+
prev_inst = inst
|
| 159 |
+
|
| 160 |
+
return path
|
| 161 |
+
|
| 162 |
+
def _check_collision_linear_path(self,path):
|
| 163 |
+
"""Check for collision on a linear path from start to goal
|
| 164 |
+
|
| 165 |
+
:param path: A list containing start and goal as [x,y,yaw]
|
| 166 |
+
:return: A bool, True if collision was detected
|
| 167 |
+
"""
|
| 168 |
+
start = path[0]
|
| 169 |
+
end = path[1]
|
| 170 |
+
|
| 171 |
+
m = (end[1] - start[1])/(end[0] - start[0])
|
| 172 |
+
b = start[1] - m * start[0]
|
| 173 |
+
x_range = [start[0],end[0]]
|
| 174 |
+
x_span = start[0] - end[0]
|
| 175 |
+
|
| 176 |
+
incr = round(abs(x_span)/50, 3)
|
| 177 |
+
if x_range[1] < x_range[0]:
|
| 178 |
+
incr = - incr
|
| 179 |
+
|
| 180 |
+
x = x_range[0]
|
| 181 |
+
for k in range(50):
|
| 182 |
+
x += incr
|
| 183 |
+
y = m * x + b
|
| 184 |
+
self.set_2d_pose([x,y,start[-1] if k < 46 else end[-1]])
|
| 185 |
+
status_collision = self.assess_collision()
|
| 186 |
+
if status_collision == True:
|
| 187 |
+
break
|
| 188 |
+
|
| 189 |
+
return status_collision
|
| 190 |
+
|
| 191 |
+
def get_base_actuation(self):
|
| 192 |
+
"""Controller for mobile bases.
|
| 193 |
+
"""
|
| 194 |
+
raise NotImplementedError()
|
| 195 |
+
|
| 196 |
+
def copy(self) -> RobotComponent:
|
| 197 |
+
self.intermediate_target_base.set_parent(self)
|
| 198 |
+
self.target_base.set_parent(self)
|
| 199 |
+
c = super().copy()
|
| 200 |
+
self.intermediate_target_base.set_parent(None)
|
| 201 |
+
self.target_base.set_parent(None)
|
| 202 |
+
return c
|
third_party/AnyBimanual/third_party/PyRep/pyrep/robots/mobiles/nonholonomic_base.py
ADDED
|
@@ -0,0 +1,140 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from math import atan2
|
| 2 |
+
from math import cos
|
| 3 |
+
from math import sin
|
| 4 |
+
from math import sqrt
|
| 5 |
+
from typing import List
|
| 6 |
+
|
| 7 |
+
from pyrep.const import ConfigurationPathAlgorithms as Algos
|
| 8 |
+
from pyrep.errors import ConfigurationPathError
|
| 9 |
+
from pyrep.robots.configuration_paths.nonholonomic_configuration_path import (
|
| 10 |
+
NonHolonomicConfigurationPath)
|
| 11 |
+
from pyrep.robots.mobiles.mobile_base import MobileBase
|
| 12 |
+
|
| 13 |
+
|
| 14 |
+
class NonHolonomicBase(MobileBase):
|
| 15 |
+
"""Currently only differential drive robots.
|
| 16 |
+
Can be refactored to include other types of non-holonomic bases in future.
|
| 17 |
+
"""
|
| 18 |
+
|
| 19 |
+
def __init__(self,
|
| 20 |
+
count: int,
|
| 21 |
+
num_wheels: int,
|
| 22 |
+
name: str):
|
| 23 |
+
|
| 24 |
+
super().__init__(count, num_wheels, name)
|
| 25 |
+
|
| 26 |
+
self.cummulative_error = 0
|
| 27 |
+
self.prev_error = 0
|
| 28 |
+
|
| 29 |
+
# PID controller values.
|
| 30 |
+
# TODO: expose to user through constructor.
|
| 31 |
+
self.Kp = 1.0
|
| 32 |
+
self.Ki = 0.01
|
| 33 |
+
self.Kd = 0.1
|
| 34 |
+
self.desired_velocity = 0.05
|
| 35 |
+
|
| 36 |
+
def get_base_velocities(self) -> List[float]:
|
| 37 |
+
"""Gets linear and angular velocities of the mobile robot base
|
| 38 |
+
calculated from kinematics equations. Left joint should have index 1,
|
| 39 |
+
right joint should have index.
|
| 40 |
+
|
| 41 |
+
:return: A list with linear and angular velocity of the robot base.
|
| 42 |
+
"""
|
| 43 |
+
wv = self.get_joint_velocities()
|
| 44 |
+
|
| 45 |
+
lv = sum(wv) * self.wheel_radius / 2.0
|
| 46 |
+
v_diff = wv[1] - wv[0]
|
| 47 |
+
av = self.wheel_radius * v_diff / self.wheel_distance
|
| 48 |
+
|
| 49 |
+
return [lv, av]
|
| 50 |
+
|
| 51 |
+
def get_linear_path(self, position: List[float],
|
| 52 |
+
angle=0) -> NonHolonomicConfigurationPath:
|
| 53 |
+
"""Initialize linear path and check for collision along it.
|
| 54 |
+
|
| 55 |
+
Must specify either rotation in euler or quaternions, but not both!
|
| 56 |
+
|
| 57 |
+
:param position: The x, y position of the target.
|
| 58 |
+
:param angle: The z orientation of the target (in radians).
|
| 59 |
+
:raises: ConfigurationPathError if no path could be created.
|
| 60 |
+
|
| 61 |
+
:return: A linear path in the 2d space.
|
| 62 |
+
"""
|
| 63 |
+
position_base = self.get_position()
|
| 64 |
+
angle_base = self.get_orientation()[-1]
|
| 65 |
+
|
| 66 |
+
self.target_base.set_position(
|
| 67 |
+
[position[0], position[1], self.target_z])
|
| 68 |
+
self.target_base.set_orientation([0, 0, angle])
|
| 69 |
+
self.intermediate_target_base.set_position(
|
| 70 |
+
[position[0], position[1], self.target_z])
|
| 71 |
+
self.intermediate_target_base.set_orientation([0, 0, angle])
|
| 72 |
+
|
| 73 |
+
path = [[position_base[0], position_base[1], angle_base],
|
| 74 |
+
[position[0], position[1], angle]]
|
| 75 |
+
|
| 76 |
+
if self._check_collision_linear_path(path):
|
| 77 |
+
raise ConfigurationPathError(
|
| 78 |
+
'Could not create path. '
|
| 79 |
+
'An object was detected on the linear path.')
|
| 80 |
+
|
| 81 |
+
return NonHolonomicConfigurationPath(self, path)
|
| 82 |
+
|
| 83 |
+
def get_nonlinear_path(self, position: List[float],
|
| 84 |
+
angle=0,
|
| 85 |
+
boundaries=2,
|
| 86 |
+
path_pts=600,
|
| 87 |
+
ignore_collisions=False,
|
| 88 |
+
algorithm=Algos.RRTConnect
|
| 89 |
+
) -> NonHolonomicConfigurationPath:
|
| 90 |
+
"""Gets a non-linear (planned) configuration path given a target pose.
|
| 91 |
+
|
| 92 |
+
:param position: The x, y, z position of the target.
|
| 93 |
+
:param angle: The z orientation of the target (in radians).
|
| 94 |
+
:param boundaries: A float defining the path search in x and y direction
|
| 95 |
+
[[-boundaries,boundaries],[-boundaries,boundaries]].
|
| 96 |
+
:param path_pts: number of sampled points returned from the computed path
|
| 97 |
+
:param ignore_collisions: If collision checking should be disabled.
|
| 98 |
+
:param algorithm: Algorithm used to compute path
|
| 99 |
+
:raises: ConfigurationPathError if no path could be created.
|
| 100 |
+
|
| 101 |
+
:return: A non-linear path (x,y,angle) in the xy configuration space.
|
| 102 |
+
"""
|
| 103 |
+
|
| 104 |
+
path = self._get_nonlinear_path_points(
|
| 105 |
+
position, angle, boundaries, path_pts, ignore_collisions, algorithm)
|
| 106 |
+
|
| 107 |
+
return NonHolonomicConfigurationPath(self, path)
|
| 108 |
+
|
| 109 |
+
def get_base_actuation(self):
|
| 110 |
+
"""A controller using PID.
|
| 111 |
+
|
| 112 |
+
:return: A list with left and right joint velocity, and bool if target is reached.
|
| 113 |
+
"""
|
| 114 |
+
|
| 115 |
+
d_x, d_y, _ = self.intermediate_target_base.get_position(
|
| 116 |
+
relative_to=self)
|
| 117 |
+
|
| 118 |
+
d_x_final, d_y_final, _ = self.target_base.get_position(
|
| 119 |
+
relative_to=self)
|
| 120 |
+
|
| 121 |
+
if sqrt((d_x_final) ** 2 + (d_y_final) ** 2) < 0.1:
|
| 122 |
+
return [0., 0.], True
|
| 123 |
+
|
| 124 |
+
alpha = atan2(d_y, d_x)
|
| 125 |
+
e = atan2(sin(alpha), cos(alpha))
|
| 126 |
+
e_P = e
|
| 127 |
+
e_I = self.cummulative_error + e
|
| 128 |
+
e_D = e - self.prev_error
|
| 129 |
+
w = self.Kp * e_P + self.Ki * e_I + self.Kd * e_D
|
| 130 |
+
w = atan2(sin(w), cos(w))
|
| 131 |
+
|
| 132 |
+
self.cummulative_error = self.cummulative_error + e
|
| 133 |
+
self.prev_error = e
|
| 134 |
+
|
| 135 |
+
vr = ((2. * self.desired_velocity + w * self.wheel_distance) /
|
| 136 |
+
(2. * self.wheel_radius))
|
| 137 |
+
vl = ((2. * self.desired_velocity - w * self.wheel_distance) /
|
| 138 |
+
(2. * self.wheel_radius))
|
| 139 |
+
|
| 140 |
+
return [vl, vr], False
|
third_party/AnyBimanual/third_party/PyRep/pyrep/robots/mobiles/turtlebot.py
ADDED
|
@@ -0,0 +1,6 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from pyrep.robots.mobiles.nonholonomic_base import NonHolonomicBase
|
| 2 |
+
|
| 3 |
+
|
| 4 |
+
class TurtleBot(NonHolonomicBase):
|
| 5 |
+
def __init__(self, count: int = 0):
|
| 6 |
+
super().__init__(count, 2, 'turtlebot')
|
third_party/AnyBimanual/third_party/PyRep/pyrep/robots/mobiles/youbot.py
ADDED
|
@@ -0,0 +1,7 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from pyrep.robots.mobiles.holonomic_base import HolonomicBase
|
| 2 |
+
|
| 3 |
+
|
| 4 |
+
class YouBot(HolonomicBase):
|
| 5 |
+
def __init__(self, count: int = 0, distance_from_target: float = 0):
|
| 6 |
+
super().__init__(
|
| 7 |
+
count, 4, distance_from_target, 'youBot', 4, 6, 0.035)
|
third_party/AnyBimanual/third_party/PyRep/pyrep/robots/robot_component.py
ADDED
|
@@ -0,0 +1,290 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from typing import List, Tuple
|
| 2 |
+
|
| 3 |
+
from pyrep.objects.shape import Shape
|
| 4 |
+
|
| 5 |
+
from pyrep.backend import sim, utils
|
| 6 |
+
from pyrep.const import JointType
|
| 7 |
+
from pyrep.objects.object import Object
|
| 8 |
+
from pyrep.objects.joint import Joint
|
| 9 |
+
from pyrep.const import ObjectType, JointMode
|
| 10 |
+
|
| 11 |
+
|
| 12 |
+
class RobotComponent(Object):
|
| 13 |
+
"""Collection of joints representing arms, end effectors, mobile bases, etc.
|
| 14 |
+
"""
|
| 15 |
+
|
| 16 |
+
def __init__(self, count: int, name: str, joint_names: List[str],
|
| 17 |
+
base_name: str = None):
|
| 18 |
+
self.name = name
|
| 19 |
+
suffix = '' if count == 0 else '#%d' % (count - 1)
|
| 20 |
+
super().__init__(
|
| 21 |
+
name + suffix if base_name is None else base_name + suffix)
|
| 22 |
+
self._num_joints = len(joint_names)
|
| 23 |
+
|
| 24 |
+
# Joint handles
|
| 25 |
+
self.joints = [Joint(jname + suffix)
|
| 26 |
+
for jname in joint_names]
|
| 27 |
+
self._joint_handles = [j.get_handle() for j in self.joints]
|
| 28 |
+
|
| 29 |
+
def copy(self) -> 'RobotComponent':
|
| 30 |
+
"""Copy and pastes the arm in the scene.
|
| 31 |
+
|
| 32 |
+
The arm is copied together with all its associated calculation
|
| 33 |
+
objects and associated scripts.
|
| 34 |
+
|
| 35 |
+
:return: The new pasted arm.
|
| 36 |
+
"""
|
| 37 |
+
# Copy whole model
|
| 38 |
+
handle = sim.simCopyPasteObjects([self._handle], 1)[0]
|
| 39 |
+
name = sim.simGetObjectName(handle)
|
| 40 |
+
# Find the number of this arm
|
| 41 |
+
num = name[name.rfind('#') + 1:]
|
| 42 |
+
if len(num) > 0:
|
| 43 |
+
num = int(num) + 1
|
| 44 |
+
else:
|
| 45 |
+
num = 0
|
| 46 |
+
# FIXME: Pass valid name and joint_names.
|
| 47 |
+
return self.__class__(num) # type: ignore
|
| 48 |
+
|
| 49 |
+
def _get_requested_type(self) -> ObjectType:
|
| 50 |
+
"""Gets the type of the object.
|
| 51 |
+
|
| 52 |
+
:return: Type of the object.
|
| 53 |
+
"""
|
| 54 |
+
return ObjectType(sim.simGetObjectType(self.get_handle()))
|
| 55 |
+
|
| 56 |
+
def get_joint_count(self) -> int:
|
| 57 |
+
"""Gets the number of joints in this component.
|
| 58 |
+
|
| 59 |
+
:return: The number of joints.
|
| 60 |
+
"""
|
| 61 |
+
return self._num_joints
|
| 62 |
+
|
| 63 |
+
def get_joint_types(self) -> List[JointType]:
|
| 64 |
+
"""Retrieves the type of the joints in this component.
|
| 65 |
+
|
| 66 |
+
:return: A list containing the types of the joints.
|
| 67 |
+
"""
|
| 68 |
+
return [j.get_joint_type() for j in self.joints]
|
| 69 |
+
|
| 70 |
+
def get_joint_positions(self) -> List[float]:
|
| 71 |
+
"""Retrieves the intrinsic position of the joints.
|
| 72 |
+
|
| 73 |
+
See :py:meth:`Joint.get_joint_position` for more information.
|
| 74 |
+
|
| 75 |
+
:return: A list of intrinsic position of the joints.
|
| 76 |
+
"""
|
| 77 |
+
return [j.get_joint_position() for j in self.joints]
|
| 78 |
+
|
| 79 |
+
def set_joint_positions(self, positions: List[float],
|
| 80 |
+
disable_dynamics: bool = False) -> None:
|
| 81 |
+
"""Sets the intrinsic position of the joints.
|
| 82 |
+
|
| 83 |
+
See :py:meth:`Joint.set_joint_position` for more information.
|
| 84 |
+
|
| 85 |
+
:param disable_dynamics: If True, then the position can be set even
|
| 86 |
+
when the joint mode is in Force mode. It will disable dynamics,
|
| 87 |
+
move the joint, and then re-enable dynamics.
|
| 88 |
+
|
| 89 |
+
:param positions: A list of positions of the joints (angular or linear
|
| 90 |
+
values depending on the joint type).
|
| 91 |
+
"""
|
| 92 |
+
self._assert_len(positions)
|
| 93 |
+
if not disable_dynamics:
|
| 94 |
+
[sim.simSetJointPosition(jh, p) # type: ignore
|
| 95 |
+
for jh, p in zip(self._joint_handles, positions)]
|
| 96 |
+
return
|
| 97 |
+
|
| 98 |
+
is_model = self.is_model()
|
| 99 |
+
if not is_model:
|
| 100 |
+
self.set_model(True)
|
| 101 |
+
|
| 102 |
+
prior = sim.simGetModelProperty(self.get_handle())
|
| 103 |
+
p = prior | sim.sim_modelproperty_not_dynamic
|
| 104 |
+
# Disable the dynamics
|
| 105 |
+
sim.simSetModelProperty(self._handle, p)
|
| 106 |
+
with utils.step_lock:
|
| 107 |
+
sim.simExtStep(True) # Have to step for changes to take effect
|
| 108 |
+
|
| 109 |
+
[sim.simSetJointPosition(jh, p) # type: ignore
|
| 110 |
+
for jh, p in zip(self._joint_handles, positions)]
|
| 111 |
+
[j.set_joint_target_position(p) # type: ignore
|
| 112 |
+
for j, p in zip(self.joints, positions)]
|
| 113 |
+
with utils.step_lock:
|
| 114 |
+
sim.simExtStep(True) # Have to step for changes to take effect
|
| 115 |
+
# Re-enable the dynamics
|
| 116 |
+
sim.simSetModelProperty(self._handle, prior)
|
| 117 |
+
self.set_model(is_model)
|
| 118 |
+
|
| 119 |
+
def get_joint_target_positions(self) -> List[float]:
|
| 120 |
+
"""Retrieves the target positions of the joints.
|
| 121 |
+
|
| 122 |
+
:return: A list of target position of the joints (angular or linear
|
| 123 |
+
values depending on the joint type).
|
| 124 |
+
"""
|
| 125 |
+
return [j.get_joint_target_position() for j in self.joints]
|
| 126 |
+
|
| 127 |
+
def set_joint_target_positions(self, positions: List[float]) -> None:
|
| 128 |
+
"""Sets the target positions of the joints.
|
| 129 |
+
|
| 130 |
+
See :py:meth:`Joint.set_joint_target_position` for more information.
|
| 131 |
+
|
| 132 |
+
:param positions: List of target position of the joints (angular or
|
| 133 |
+
linear values depending on the joint type).
|
| 134 |
+
"""
|
| 135 |
+
self._assert_len(positions)
|
| 136 |
+
[j.set_joint_target_position(p) # type: ignore
|
| 137 |
+
for j, p in zip(self.joints, positions)]
|
| 138 |
+
|
| 139 |
+
def get_joint_target_velocities(self) -> List[float]:
|
| 140 |
+
"""Retrieves the intrinsic target velocities of the joints.
|
| 141 |
+
|
| 142 |
+
:return: List of the target velocity of the joints (linear or angular
|
| 143 |
+
velocity depending on the joint-type).
|
| 144 |
+
"""
|
| 145 |
+
return [j.get_joint_target_velocity() # type: ignore
|
| 146 |
+
for j in self.joints]
|
| 147 |
+
|
| 148 |
+
def set_joint_target_velocities(self, velocities: List[float]) -> None:
|
| 149 |
+
"""Sets the intrinsic target velocities of the joints.
|
| 150 |
+
|
| 151 |
+
:param velocities: List of the target velocity of the joints (linear
|
| 152 |
+
or angular velocities depending on the joint-type).
|
| 153 |
+
"""
|
| 154 |
+
self._assert_len(velocities)
|
| 155 |
+
[j.set_joint_target_velocity(v) # type: ignore
|
| 156 |
+
for j, v in zip(self.joints, velocities)]
|
| 157 |
+
|
| 158 |
+
def get_joint_forces(self) -> List[float]:
|
| 159 |
+
"""Retrieves the forces or torques of the joints.
|
| 160 |
+
|
| 161 |
+
See :py:meth:`Joint.get_joint_force` for more information.
|
| 162 |
+
|
| 163 |
+
:return: A list of the forces or the torques applied to the joints
|
| 164 |
+
along/about their z-axis.
|
| 165 |
+
"""
|
| 166 |
+
return [j.get_joint_force() for j in self.joints]
|
| 167 |
+
|
| 168 |
+
def set_joint_forces(self, forces: List[float]) -> None:
|
| 169 |
+
"""Sets the maximum force or torque that the joints can exert.
|
| 170 |
+
|
| 171 |
+
See :py:meth:`Joint.set_joint_force` for more information.
|
| 172 |
+
|
| 173 |
+
:param forces: The maximum force or torque that the joints can exert.
|
| 174 |
+
These cannot be negative values.
|
| 175 |
+
"""
|
| 176 |
+
self._assert_len(forces)
|
| 177 |
+
[j.set_joint_force(f) # type: ignore
|
| 178 |
+
for j, f in zip(self.joints, forces)]
|
| 179 |
+
|
| 180 |
+
def get_joint_velocities(self) -> List[float]:
|
| 181 |
+
"""Get the current joint velocities.
|
| 182 |
+
|
| 183 |
+
:return: List containing the velocities of the joints (linear or
|
| 184 |
+
angular velocities depending on the joint-type).
|
| 185 |
+
"""
|
| 186 |
+
return [j.get_joint_velocity() for j in self.joints]
|
| 187 |
+
|
| 188 |
+
def get_joint_intervals(self) -> Tuple[List[bool], List[List[float]]]:
|
| 189 |
+
"""Retrieves the interval parameters of the joints.
|
| 190 |
+
|
| 191 |
+
See :py:meth:`Joint.get_joint_interval` for more information.
|
| 192 |
+
|
| 193 |
+
:return: A tuple containing a list of bools indicates whether the joint
|
| 194 |
+
is cyclic (the joint varies between -pi and +pi in a cyclic manner),
|
| 195 |
+
and a 2D list containing the interval of the joints.
|
| 196 |
+
"""
|
| 197 |
+
cyclics, intervals = [], []
|
| 198 |
+
for j in self.joints:
|
| 199 |
+
c, i = j.get_joint_interval()
|
| 200 |
+
cyclics.append(c)
|
| 201 |
+
intervals.append(i)
|
| 202 |
+
return cyclics, intervals
|
| 203 |
+
|
| 204 |
+
def set_joint_intervals(self, cyclic: List[bool],
|
| 205 |
+
intervals: List[List[float]]) -> None:
|
| 206 |
+
"""Sets the interval parameters of the joints (i.e. range values).
|
| 207 |
+
|
| 208 |
+
See :py:meth:`Joint.set_joint_interval` for more information.
|
| 209 |
+
|
| 210 |
+
:param cyclic: List of bools indicates whether the joint is cyclic.
|
| 211 |
+
Only revolute joints with a pitch of 0 can be cyclic.
|
| 212 |
+
:param intervals: 2D list containing the intervals of the joints.
|
| 213 |
+
"""
|
| 214 |
+
self._assert_len(cyclic)
|
| 215 |
+
self._assert_len(intervals)
|
| 216 |
+
[j.set_joint_interval(c, i) # type: ignore
|
| 217 |
+
for j, c, i in zip( self.joints, cyclic, intervals)]
|
| 218 |
+
|
| 219 |
+
def get_joint_upper_velocity_limits(self) -> List[float]:
|
| 220 |
+
"""Gets upper velocity limits of the joints.
|
| 221 |
+
|
| 222 |
+
:return: List of the upper velocity limits.
|
| 223 |
+
"""
|
| 224 |
+
return [j.get_joint_upper_velocity_limit() for j in self.joints]
|
| 225 |
+
|
| 226 |
+
def set_control_loop_enabled(self, value: bool) -> None:
|
| 227 |
+
"""Sets whether the control loop is enable for all joints.
|
| 228 |
+
|
| 229 |
+
:param value: The new value for the control loop state.
|
| 230 |
+
"""
|
| 231 |
+
[j.set_control_loop_enabled(value) # type: ignore
|
| 232 |
+
for j in self.joints]
|
| 233 |
+
|
| 234 |
+
def set_motor_locked_at_zero_velocity(self, value: bool) -> None:
|
| 235 |
+
"""Sets if motor is locked when target velocity is zero for all joints.
|
| 236 |
+
|
| 237 |
+
When enabled in velocity mode and its target velocity is zero, then the
|
| 238 |
+
joint is locked in place.
|
| 239 |
+
|
| 240 |
+
:param value: If the motors should be locked at zero velocity.
|
| 241 |
+
"""
|
| 242 |
+
[j.set_motor_locked_at_zero_velocity(value) # type: ignore
|
| 243 |
+
for j in self.joints]
|
| 244 |
+
|
| 245 |
+
def set_joint_mode(self, value: JointMode) -> None:
|
| 246 |
+
"""Sets the operation mode of the joint group.
|
| 247 |
+
|
| 248 |
+
:param value: The new joint mode value.
|
| 249 |
+
"""
|
| 250 |
+
[j.set_joint_mode(value) # type: ignore
|
| 251 |
+
for j in self.joints]
|
| 252 |
+
|
| 253 |
+
def get_joint_modes(self) -> List[JointMode]:
|
| 254 |
+
"""Gets the operation mode of the joint group.
|
| 255 |
+
|
| 256 |
+
:return: A list of joint modes.
|
| 257 |
+
"""
|
| 258 |
+
return [j.get_joint_mode() for j in self.joints]
|
| 259 |
+
|
| 260 |
+
def get_visuals(self) -> List[Object]:
|
| 261 |
+
"""Gets a list of the visual elements of this component.
|
| 262 |
+
|
| 263 |
+
Can be useful for methods such as domain randomization.
|
| 264 |
+
Should ideally be overridden for each robot.
|
| 265 |
+
|
| 266 |
+
:return: A list of visual shapes.
|
| 267 |
+
"""
|
| 268 |
+
tree = self.get_objects_in_tree(ObjectType.SHAPE, exclude_base=False)
|
| 269 |
+
return [obj for obj in tree if 'visual' in obj.get_name()]
|
| 270 |
+
|
| 271 |
+
def _assert_len(self, inputs: list) -> None:
|
| 272 |
+
if len(self.joints) != len(inputs):
|
| 273 |
+
raise RuntimeError(
|
| 274 |
+
'Tried to set values for %d joints, but joint group consists '
|
| 275 |
+
'of %d joints.' % (len(inputs), len(self.joints)))
|
| 276 |
+
|
| 277 |
+
def colorize(self, color: List[float]):
|
| 278 |
+
for v in self.get_visuals():
|
| 279 |
+
if v.get_name().startswith(self.name):
|
| 280 |
+
v.set_color(color)
|
| 281 |
+
|
| 282 |
+
def hide(self):
|
| 283 |
+
for v in self.get_visuals():
|
| 284 |
+
if v.get_name().startswith(self.name):
|
| 285 |
+
v.set_renderable(False)
|
| 286 |
+
|
| 287 |
+
def show(self):
|
| 288 |
+
for v in self.get_visuals():
|
| 289 |
+
if v.get_name().startswith(self.name):
|
| 290 |
+
v.set_renderable(True)
|
third_party/AnyBimanual/third_party/PyRep/pyrep/sensors/__init__.py
ADDED
|
File without changes
|
third_party/AnyBimanual/third_party/PyRep/pyrep/sensors/accelerometer.py
ADDED
|
@@ -0,0 +1,30 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from typing import List
|
| 2 |
+
|
| 3 |
+
from pyrep.backend import sim
|
| 4 |
+
from pyrep.const import ObjectType
|
| 5 |
+
from pyrep.objects.force_sensor import ForceSensor
|
| 6 |
+
from pyrep.objects.object import Object
|
| 7 |
+
from pyrep.objects.shape import Shape
|
| 8 |
+
|
| 9 |
+
|
| 10 |
+
class Accelerometer(Object):
|
| 11 |
+
"""An object able to measure accelerations that are applied to it.
|
| 12 |
+
"""
|
| 13 |
+
|
| 14 |
+
def __init__(self, name):
|
| 15 |
+
super().__init__(name)
|
| 16 |
+
self._mass_object = Shape('%s_mass' % (self.get_name()))
|
| 17 |
+
self._sensor = ForceSensor('%s_force_sensor' % (self.get_name()))
|
| 18 |
+
|
| 19 |
+
def _get_requested_type(self) -> ObjectType:
|
| 20 |
+
return ObjectType(sim.simGetObjectType(self.get_handle()))
|
| 21 |
+
|
| 22 |
+
def read(self) -> List[float]:
|
| 23 |
+
"""Reads the acceleration applied to accelerometer.
|
| 24 |
+
|
| 25 |
+
:return: A list containing applied accelerations along
|
| 26 |
+
the sensor's x, y and z-axes
|
| 27 |
+
"""
|
| 28 |
+
forces, _ = self._sensor.read()
|
| 29 |
+
accel = [force / self._mass_object.get_mass() for force in forces]
|
| 30 |
+
return accel
|
third_party/AnyBimanual/third_party/PyRep/pyrep/sensors/gyroscope.py
ADDED
|
@@ -0,0 +1,40 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from pyrep.backend import sim
|
| 2 |
+
from pyrep.const import ObjectType
|
| 3 |
+
from pyrep.objects.object import Object
|
| 4 |
+
|
| 5 |
+
|
| 6 |
+
class Gyroscope(Object):
|
| 7 |
+
"""An object able to measure angular velocities that are applied to it.
|
| 8 |
+
"""
|
| 9 |
+
def __init__(self, name):
|
| 10 |
+
super().__init__(name)
|
| 11 |
+
self._ref = '%s_reference' % (self.get_name())
|
| 12 |
+
|
| 13 |
+
self._last_time = sim.simGetSimulationTime()
|
| 14 |
+
self._old_transformation_matrix = self.get_matrix()[:3, :4].reshape(
|
| 15 |
+
(12, )).tolist()
|
| 16 |
+
|
| 17 |
+
def _get_requested_type(self) -> ObjectType:
|
| 18 |
+
return ObjectType(sim.simGetObjectType(self.get_handle()))
|
| 19 |
+
|
| 20 |
+
def read(self):
|
| 21 |
+
"""Reads the angular velocities applied to gyroscope.
|
| 22 |
+
|
| 23 |
+
:return: A list containing applied angular velocities along
|
| 24 |
+
the sensor's x, y and z-axes.
|
| 25 |
+
"""
|
| 26 |
+
current_time = sim.simGetSimulationTime()
|
| 27 |
+
dt = current_time - self._last_time
|
| 28 |
+
|
| 29 |
+
inv_old_matrix = sim.simInvertMatrix(self._old_transformation_matrix)
|
| 30 |
+
transformation_matrix = self.get_matrix()[:3, :4].reshape(
|
| 31 |
+
(12, )).tolist()
|
| 32 |
+
mat = sim.simMultiplyMatrices(inv_old_matrix, transformation_matrix)
|
| 33 |
+
euler_angles = sim.simGetEulerAnglesFromMatrix(mat)
|
| 34 |
+
|
| 35 |
+
self._last_time = current_time
|
| 36 |
+
self._old_transformation_matrix = transformation_matrix
|
| 37 |
+
|
| 38 |
+
if dt != 0:
|
| 39 |
+
return [euler_angle / dt for euler_angle in euler_angles]
|
| 40 |
+
return [0.0, 0.0, 0.0]
|
third_party/AnyBimanual/third_party/PyRep/pyrep/sensors/spherical_vision_sensor.py
ADDED
|
@@ -0,0 +1,296 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import numpy as np
|
| 2 |
+
from typing import List, Sequence
|
| 3 |
+
|
| 4 |
+
from pyrep.backend import sim, utils
|
| 5 |
+
from pyrep.objects.object import Object
|
| 6 |
+
from pyrep.objects.vision_sensor import VisionSensor
|
| 7 |
+
from pyrep.const import PYREP_SCRIPT_TYPE, RenderMode, ObjectType
|
| 8 |
+
|
| 9 |
+
MIN_DIVISOR = 1e-12
|
| 10 |
+
|
| 11 |
+
|
| 12 |
+
class SphericalVisionSensor(Object):
|
| 13 |
+
"""An object able capture 360 degree omni-directional images.
|
| 14 |
+
"""
|
| 15 |
+
|
| 16 |
+
def __init__(self, name):
|
| 17 |
+
super().__init__(name)
|
| 18 |
+
|
| 19 |
+
# final omni-directional sensors
|
| 20 |
+
self._sensor_depth = VisionSensor('%s_sensorDepth' % (self.get_name()))
|
| 21 |
+
self._sensor_rgb = VisionSensor('%s_sensorRGB' % (self.get_name()))
|
| 22 |
+
|
| 23 |
+
# directed sub-sensors
|
| 24 |
+
names = ['front', 'top', 'back', 'bottom', 'left', 'right']
|
| 25 |
+
self._six_sensors = [VisionSensor('%s_%s' % (self.get_name(), n)) for n in names]
|
| 26 |
+
self._front = self._six_sensors[0]
|
| 27 |
+
|
| 28 |
+
# directed sub-sensor handles list
|
| 29 |
+
self._six_sensor_handles = [vs.get_handle() for vs in self._six_sensors]
|
| 30 |
+
|
| 31 |
+
# set all to explicit handling
|
| 32 |
+
self._set_all_to_explicit_handling()
|
| 33 |
+
|
| 34 |
+
# assert sensor properties are matching
|
| 35 |
+
self._assert_matching_resolutions()
|
| 36 |
+
self._assert_matching_render_modes()
|
| 37 |
+
self._assert_matching_entities_to_render()
|
| 38 |
+
self._assert_matching_window_sizes()
|
| 39 |
+
self._assert_matching_near_clipping_planes()
|
| 40 |
+
self._assert_matching_far_clipping_planes()
|
| 41 |
+
|
| 42 |
+
# omni resolution
|
| 43 |
+
self._omni_resolution = self._sensor_rgb.get_resolution()
|
| 44 |
+
|
| 45 |
+
# near and far clipping plane
|
| 46 |
+
self._near_clipping_plane = self.get_near_clipping_plane()
|
| 47 |
+
self._far_clipping_plane = self.get_far_clipping_plane()
|
| 48 |
+
self._clipping_plane_diff = self._far_clipping_plane - self._near_clipping_plane
|
| 49 |
+
|
| 50 |
+
# projective cam region image
|
| 51 |
+
self._planar_to_radial_depth_scalars = self._get_planar_to_radial_depth_scalars()
|
| 52 |
+
|
| 53 |
+
# Private #
|
| 54 |
+
# --------#
|
| 55 |
+
|
| 56 |
+
def _set_all_to_explicit_handling(self):
|
| 57 |
+
self._sensor_depth.set_explicit_handling(1)
|
| 58 |
+
self._sensor_rgb.set_explicit_handling(1)
|
| 59 |
+
for sensor in self._six_sensors:
|
| 60 |
+
sensor.set_explicit_handling(1)
|
| 61 |
+
|
| 62 |
+
def _assert_matching_resolutions(self):
|
| 63 |
+
assert self._sensor_depth.get_resolution() == self._sensor_rgb.get_resolution()
|
| 64 |
+
front_sensor_res = self._front.get_resolution()
|
| 65 |
+
for sensor in self._six_sensors[1:]:
|
| 66 |
+
assert sensor.get_resolution() == front_sensor_res
|
| 67 |
+
|
| 68 |
+
def _assert_matching_render_modes(self):
|
| 69 |
+
assert self._sensor_depth.get_render_mode() == self._sensor_rgb.get_render_mode()
|
| 70 |
+
front_sensor_render_mode = self._front.get_render_mode()
|
| 71 |
+
for sensor in self._six_sensors[1:]:
|
| 72 |
+
assert sensor.get_render_mode() == front_sensor_render_mode
|
| 73 |
+
|
| 74 |
+
def _assert_matching_entities_to_render(self):
|
| 75 |
+
assert self._sensor_depth.get_entity_to_render() == self._sensor_rgb.get_entity_to_render()
|
| 76 |
+
front_sensor_entity_to_render = self._front.get_entity_to_render()
|
| 77 |
+
for sensor in self._six_sensors[1:]:
|
| 78 |
+
assert sensor.get_entity_to_render() == front_sensor_entity_to_render
|
| 79 |
+
|
| 80 |
+
def _assert_matching_window_sizes(self):
|
| 81 |
+
assert self._sensor_depth.get_windowed_size() == self._sensor_rgb.get_windowed_size()
|
| 82 |
+
|
| 83 |
+
def _assert_matching_near_clipping_planes(self):
|
| 84 |
+
front_sensor_near = self._front.get_near_clipping_plane()
|
| 85 |
+
for sensor in self._six_sensors[1:]:
|
| 86 |
+
assert sensor.get_near_clipping_plane() == front_sensor_near
|
| 87 |
+
|
| 88 |
+
def _assert_matching_far_clipping_planes(self):
|
| 89 |
+
front_sensor_far = self._front.get_far_clipping_plane()
|
| 90 |
+
for sensor in self._six_sensors[1:]:
|
| 91 |
+
assert sensor.get_far_clipping_plane() == front_sensor_far
|
| 92 |
+
|
| 93 |
+
def _get_requested_type(self) -> ObjectType:
|
| 94 |
+
return ObjectType(sim.simGetObjectType(self.get_handle()))
|
| 95 |
+
|
| 96 |
+
@staticmethod
|
| 97 |
+
def _create_uniform_pixel_coords_image(image_dims):
|
| 98 |
+
pixel_x_coords = np.reshape(np.tile(np.arange(image_dims[1]), [image_dims[0]]),
|
| 99 |
+
(image_dims[0], image_dims[1], 1)).astype(np.float32)
|
| 100 |
+
pixel_y_coords_ = np.reshape(np.tile(np.arange(image_dims[0]), [image_dims[1]]),
|
| 101 |
+
(image_dims[1], image_dims[0], 1)).astype(np.float32)
|
| 102 |
+
pixel_y_coords = np.transpose(pixel_y_coords_, (1, 0, 2))
|
| 103 |
+
return np.concatenate((pixel_x_coords, pixel_y_coords), -1)
|
| 104 |
+
|
| 105 |
+
def _get_planar_to_radial_depth_scalars(self):
|
| 106 |
+
|
| 107 |
+
# reference: https://en.wikipedia.org/wiki/Spherical_coordinate_system#Cartesian_coordinates
|
| 108 |
+
|
| 109 |
+
# polar coord image
|
| 110 |
+
coord_img = self._create_uniform_pixel_coords_image([self._omni_resolution[1], self._omni_resolution[0]])
|
| 111 |
+
pixels_per_rad = self._omni_resolution[1] / np.pi
|
| 112 |
+
polar_angles = coord_img / pixels_per_rad
|
| 113 |
+
phi = polar_angles[..., 0:1]
|
| 114 |
+
theta = polar_angles[..., 1:2]
|
| 115 |
+
|
| 116 |
+
# image borders wrt the 6 projective cameras
|
| 117 |
+
phis_rel = ((phi + np.pi/4) % (np.pi/2)) / (np.pi/2)
|
| 118 |
+
lower_borders = np.arccos(1/((phis_rel*2-1)**2 + 2)**0.5)
|
| 119 |
+
upper_borders = np.pi - lower_borders
|
| 120 |
+
|
| 121 |
+
# omni image regions from the 6 projective cameras
|
| 122 |
+
xy_region = np.logical_and(theta <= upper_borders, theta >= lower_borders)
|
| 123 |
+
pos_x = np.logical_and(xy_region, np.logical_or(phi <= 45 * np.pi/180, phi >= (45 + 270) * np.pi/180))
|
| 124 |
+
pos_y = np.logical_and(xy_region,
|
| 125 |
+
np.logical_and(phi >= 45 * np.pi/180, phi <= (45 + 90) * np.pi/180))
|
| 126 |
+
neg_x = np.logical_and(xy_region,
|
| 127 |
+
np.logical_and(phi >= (45 + 90) * np.pi/180, phi <= (45 + 180) * np.pi/180))
|
| 128 |
+
neg_y = np.logical_and(xy_region,
|
| 129 |
+
np.logical_and(phi >= (45 + 180) * np.pi/180, phi <= (45 + 270) * np.pi/180))
|
| 130 |
+
pos_z = theta <= lower_borders
|
| 131 |
+
neg_z = theta >= upper_borders
|
| 132 |
+
|
| 133 |
+
# trig terms for conversion
|
| 134 |
+
sin_theta = np.sin(theta)
|
| 135 |
+
cos_theta = np.cos(theta)
|
| 136 |
+
sin_phi = np.sin(phi)
|
| 137 |
+
cos_phi = np.cos(phi)
|
| 138 |
+
|
| 139 |
+
# reciprocal terms
|
| 140 |
+
recip_sin_theta_cos_phi = 1/(sin_theta * cos_phi + MIN_DIVISOR)
|
| 141 |
+
recip_sin_theta_sin_phi = 1/(sin_theta * sin_phi + MIN_DIVISOR)
|
| 142 |
+
recip_cos_theta = 1/(cos_theta + MIN_DIVISOR)
|
| 143 |
+
|
| 144 |
+
# planar to radial depth scalars
|
| 145 |
+
return np.where(pos_x, recip_sin_theta_cos_phi,
|
| 146 |
+
np.where(neg_x, -recip_sin_theta_cos_phi,
|
| 147 |
+
np.where(pos_y, recip_sin_theta_sin_phi,
|
| 148 |
+
np.where(neg_y, -recip_sin_theta_sin_phi,
|
| 149 |
+
np.where(pos_z, recip_cos_theta,
|
| 150 |
+
np.where(neg_z, -recip_cos_theta,
|
| 151 |
+
np.zeros_like(recip_cos_theta)))))))[..., 0]
|
| 152 |
+
|
| 153 |
+
# Public #
|
| 154 |
+
# -------#
|
| 155 |
+
|
| 156 |
+
def handle_explicitly(self) -> None:
|
| 157 |
+
"""Handle spherical vision sensor explicitly.
|
| 158 |
+
|
| 159 |
+
This enables capturing image (e.g., capture_rgb())
|
| 160 |
+
without PyRep.step().
|
| 161 |
+
"""
|
| 162 |
+
utils.script_call('handleSpherical@PyRep', PYREP_SCRIPT_TYPE,
|
| 163 |
+
[self._sensor_depth._handle, self._sensor_rgb._handle] + self._six_sensor_handles, [], [], [])
|
| 164 |
+
|
| 165 |
+
def capture_rgb(self) -> np.ndarray:
|
| 166 |
+
"""Retrieves the rgb-image of a spherical vision sensor.
|
| 167 |
+
|
| 168 |
+
:return: A numpy array of size (width, height, 3)
|
| 169 |
+
"""
|
| 170 |
+
return self._sensor_rgb.capture_rgb()
|
| 171 |
+
|
| 172 |
+
def capture_depth(self, in_meters=False) -> np.ndarray:
|
| 173 |
+
"""Retrieves the depth-image of a spherical vision sensor.
|
| 174 |
+
|
| 175 |
+
:param in_meters: Whether the depth should be returned in meters.
|
| 176 |
+
:return: A numpy array of size (width, height)
|
| 177 |
+
"""
|
| 178 |
+
planar_depth = self._sensor_depth.capture_rgb()[:, :, 0]*self._clipping_plane_diff + self._near_clipping_plane
|
| 179 |
+
radial_depth = planar_depth * self._planar_to_radial_depth_scalars
|
| 180 |
+
if in_meters:
|
| 181 |
+
return radial_depth
|
| 182 |
+
return (radial_depth - self._near_clipping_plane)/self._clipping_plane_diff
|
| 183 |
+
|
| 184 |
+
def get_resolution(self) -> List[int]:
|
| 185 |
+
""" Return the spherical vision sensor's resolution.
|
| 186 |
+
|
| 187 |
+
:return: Resolution [x, y]
|
| 188 |
+
"""
|
| 189 |
+
return self._sensor_depth.get_resolution()
|
| 190 |
+
|
| 191 |
+
def set_resolution(self, resolution: List[int]) -> None:
|
| 192 |
+
""" Set the spherical vision sensor's resolution.
|
| 193 |
+
|
| 194 |
+
:param resolution: New resolution [x, y]
|
| 195 |
+
"""
|
| 196 |
+
if resolution[0] != 2*resolution[1]:
|
| 197 |
+
raise Exception('Spherical vision sensors must have an X resolution 2x the Y resolution')
|
| 198 |
+
if resolution[1] % 2 != 0:
|
| 199 |
+
raise Exception('Spherical vision sensors must have an even Y resolution')
|
| 200 |
+
box_sensor_resolution = [int(resolution[1]/2), int(resolution[1]/2)]
|
| 201 |
+
for sensor in self._six_sensors:
|
| 202 |
+
sensor.set_resolution(box_sensor_resolution)
|
| 203 |
+
self._sensor_depth.set_resolution(resolution)
|
| 204 |
+
self._sensor_rgb.set_resolution(resolution)
|
| 205 |
+
|
| 206 |
+
def get_render_mode(self) -> RenderMode:
|
| 207 |
+
""" Retrieves the spherical vision sensor's rendering mode
|
| 208 |
+
|
| 209 |
+
:return: RenderMode for the current rendering mode.
|
| 210 |
+
"""
|
| 211 |
+
return self._sensor_depth.get_render_mode()
|
| 212 |
+
|
| 213 |
+
def set_render_mode(self, render_mode: RenderMode) -> None:
|
| 214 |
+
""" Set the spherical vision sensor's rendering mode
|
| 215 |
+
|
| 216 |
+
:param render_mode: The new sensor rendering mode, one of:
|
| 217 |
+
RenderMode.OPENGL
|
| 218 |
+
RenderMode.OPENGL_AUXILIARY
|
| 219 |
+
RenderMode.OPENGL_COLOR_CODED
|
| 220 |
+
RenderMode.POV_RAY
|
| 221 |
+
RenderMode.EXTERNAL
|
| 222 |
+
RenderMode.EXTERNAL_WINDOWED
|
| 223 |
+
RenderMode.OPENGL3
|
| 224 |
+
RenderMode.OPENGL3_WINDOWED
|
| 225 |
+
"""
|
| 226 |
+
self._sensor_depth.set_render_mode(render_mode)
|
| 227 |
+
self._sensor_rgb.set_render_mode(render_mode)
|
| 228 |
+
for sensor in self._six_sensors:
|
| 229 |
+
sensor.set_render_mode(render_mode)
|
| 230 |
+
|
| 231 |
+
def get_windowed_size(self) -> Sequence[int]:
|
| 232 |
+
"""Get the size of windowed rendering for the omni-directional image.
|
| 233 |
+
|
| 234 |
+
:return: The (x, y) resolution of the window. 0 for full-screen.
|
| 235 |
+
"""
|
| 236 |
+
return self._sensor_depth.get_windowed_size()
|
| 237 |
+
|
| 238 |
+
def set_windowed_size(self, resolution: Sequence[int] = (0, 0)) -> None:
|
| 239 |
+
"""Set the size of windowed rendering for the omni-directional image.
|
| 240 |
+
|
| 241 |
+
:param resolution: The (x, y) resolution of the window.
|
| 242 |
+
0 for full-screen.
|
| 243 |
+
"""
|
| 244 |
+
self._sensor_depth.set_windowed_size(resolution)
|
| 245 |
+
self._sensor_rgb.set_windowed_size(resolution)
|
| 246 |
+
|
| 247 |
+
def get_near_clipping_plane(self) -> float:
|
| 248 |
+
""" Get the spherical depth sensor's near clipping plane.
|
| 249 |
+
|
| 250 |
+
:return: Near clipping plane (metres)
|
| 251 |
+
"""
|
| 252 |
+
return self._front.get_near_clipping_plane()
|
| 253 |
+
|
| 254 |
+
def set_near_clipping_plane(self, near_clipping: float) -> None:
|
| 255 |
+
""" Set the spherical depth sensor's near clipping plane.
|
| 256 |
+
|
| 257 |
+
:param near_clipping: New near clipping plane (in metres)
|
| 258 |
+
"""
|
| 259 |
+
self._near_clipping_plane = near_clipping
|
| 260 |
+
self._clipping_plane_diff = self._far_clipping_plane - near_clipping
|
| 261 |
+
for sensor in self._six_sensors:
|
| 262 |
+
sensor.set_near_clipping_plane(near_clipping)
|
| 263 |
+
|
| 264 |
+
def get_far_clipping_plane(self) -> float:
|
| 265 |
+
""" Get the Sensor's far clipping plane.
|
| 266 |
+
|
| 267 |
+
:return: Near clipping plane (metres)
|
| 268 |
+
"""
|
| 269 |
+
return self._front.get_far_clipping_plane()
|
| 270 |
+
|
| 271 |
+
def set_far_clipping_plane(self, far_clipping: float) -> None:
|
| 272 |
+
""" Set the spherical depth sensor's far clipping plane.
|
| 273 |
+
|
| 274 |
+
:param far_clipping: New far clipping plane (in metres)
|
| 275 |
+
"""
|
| 276 |
+
self._far_clipping_plane = far_clipping
|
| 277 |
+
self._clipping_plane_diff = far_clipping - self._near_clipping_plane
|
| 278 |
+
for sensor in self._six_sensors:
|
| 279 |
+
sensor.set_far_clipping_plane(far_clipping)
|
| 280 |
+
|
| 281 |
+
def set_entity_to_render(self, entity_to_render: int) -> None:
|
| 282 |
+
""" Set the entity to render to the spherical vision sensor,
|
| 283 |
+
this can be an object or more usefully a collection. -1 to render all objects in scene.
|
| 284 |
+
|
| 285 |
+
:param entity_to_render: Handle of the entity to render
|
| 286 |
+
"""
|
| 287 |
+
self._sensor_depth.set_entity_to_render(entity_to_render)
|
| 288 |
+
self._sensor_rgb.set_entity_to_render(entity_to_render)
|
| 289 |
+
|
| 290 |
+
def get_entity_to_render(self) -> None:
|
| 291 |
+
""" Get the entity to render to the spherical vision sensor,
|
| 292 |
+
this can be an object or more usefully a collection. -1 if all objects in scene are rendered.
|
| 293 |
+
|
| 294 |
+
:return: Handle of the entity to render
|
| 295 |
+
"""
|
| 296 |
+
return self._sensor_depth.get_entity_to_render()
|
third_party/AnyBimanual/third_party/PyRep/pyrep/textures/__init__.py
ADDED
|
File without changes
|
third_party/AnyBimanual/third_party/PyRep/pyrep/textures/texture.py
ADDED
|
@@ -0,0 +1,18 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
class Texture(object):
|
| 2 |
+
"""Base class for V-REP textures."""
|
| 3 |
+
|
| 4 |
+
def __init__(self, texture_id: int):
|
| 5 |
+
self._texture_id = texture_id
|
| 6 |
+
|
| 7 |
+
def __eq__(self, other: object):
|
| 8 |
+
if not isinstance(other, Texture):
|
| 9 |
+
raise NotImplementedError
|
| 10 |
+
return self.get_texture_id() == other.get_texture_id()
|
| 11 |
+
|
| 12 |
+
def get_texture_id(self) -> int:
|
| 13 |
+
"""Gets the texture id.
|
| 14 |
+
|
| 15 |
+
:return: The internal texture id.
|
| 16 |
+
"""
|
| 17 |
+
return self._texture_id
|
| 18 |
+
|
third_party/AnyBimanual/third_party/PyRep/requirements.txt
ADDED
|
@@ -0,0 +1,2 @@
|
|
|
|
|
|
|
|
|
|
| 1 |
+
numpy
|
| 2 |
+
cffi==1.14.2
|
third_party/AnyBimanual/third_party/PyRep/robot_ttms/README.md
ADDED
|
@@ -0,0 +1,5 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
Here are a number of modified model files (.ttm) ready to be used.
|
| 2 |
+
|
| 3 |
+
If using the ```pyrep.robots``` module, then use these models.
|
| 4 |
+
|
| 5 |
+
Models have been adapted from V-REP.
|
third_party/AnyBimanual/third_party/PyRep/robot_ttms/end_effectors/BaxterSuctionCup.ttm
ADDED
|
Binary file (62.1 kB). View file
|
|
|
third_party/AnyBimanual/third_party/PyRep/setup.py
ADDED
|
@@ -0,0 +1,64 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import codecs
|
| 2 |
+
import os
|
| 3 |
+
import os.path
|
| 4 |
+
from os.path import join
|
| 5 |
+
|
| 6 |
+
from setuptools import setup
|
| 7 |
+
|
| 8 |
+
import cffi_build.cffi_build as cffi_build
|
| 9 |
+
|
| 10 |
+
|
| 11 |
+
def read(rel_path):
|
| 12 |
+
here = os.path.abspath(os.path.dirname(__file__))
|
| 13 |
+
with codecs.open(os.path.join(here, rel_path), 'r') as fp:
|
| 14 |
+
return fp.read()
|
| 15 |
+
|
| 16 |
+
|
| 17 |
+
def get_version(rel_path):
|
| 18 |
+
for line in read(rel_path).splitlines():
|
| 19 |
+
if line.startswith('__version__'):
|
| 20 |
+
delim = '"' if '"' in line else "'"
|
| 21 |
+
return line.split(delim)[1]
|
| 22 |
+
else:
|
| 23 |
+
raise RuntimeError("Unable to find version string.")
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
# Temp fix for CoppeliaSim 4.1
|
| 27 |
+
if 'COPPELIASIM_ROOT' not in os.environ:
|
| 28 |
+
raise RuntimeError('COPPELIASIM_ROOT not defined.')
|
| 29 |
+
|
| 30 |
+
usrset_file = os.path.join(os.environ['COPPELIASIM_ROOT'], 'system', 'usrset.txt')
|
| 31 |
+
usrset = ''
|
| 32 |
+
if os.path.isfile(usrset_file):
|
| 33 |
+
with open(usrset_file, 'r') as f:
|
| 34 |
+
usrset = f.read()
|
| 35 |
+
|
| 36 |
+
if 'allowOldEduRelease' not in usrset:
|
| 37 |
+
with open(usrset_file, 'a+') as f:
|
| 38 |
+
f.write('\nallowOldEduRelease=7501\n')
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
setup(name='PyRep',
|
| 42 |
+
# Version A.B.C.D.
|
| 43 |
+
# A.B.C info corresponds to the CoppeliaSim version needed.
|
| 44 |
+
# D info corresponds to the PyRep version.
|
| 45 |
+
version=get_version("pyrep/__init__.py"),
|
| 46 |
+
description='Python CoppeliaSim wrapper',
|
| 47 |
+
author='Stephen James',
|
| 48 |
+
author_email='slj12@ic.ac.uk',
|
| 49 |
+
url='https://www.doc.ic.ac.uk/~slj12',
|
| 50 |
+
packages=['pyrep',
|
| 51 |
+
'pyrep.backend',
|
| 52 |
+
'pyrep.objects',
|
| 53 |
+
'pyrep.sensors',
|
| 54 |
+
'pyrep.robots',
|
| 55 |
+
'pyrep.robots.arms',
|
| 56 |
+
'pyrep.robots.end_effectors',
|
| 57 |
+
'pyrep.robots.mobiles',
|
| 58 |
+
'pyrep.robots.configuration_paths',
|
| 59 |
+
'pyrep.textures',
|
| 60 |
+
'pyrep.misc',
|
| 61 |
+
],
|
| 62 |
+
ext_modules=[cffi_build.ffibuilder.distutils_extension(
|
| 63 |
+
join('build', 'pyrep', 'backend'))],
|
| 64 |
+
)
|
third_party/AnyBimanual/third_party/PyRep/system/settings.dat
ADDED
|
Binary file (39 Bytes). View file
|
|
|
third_party/AnyBimanual/third_party/PyRep/system/usrset.txt
ADDED
|
@@ -0,0 +1,199 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
// Debugging
|
| 2 |
+
// =================================================
|
| 3 |
+
alwaysShowConsole = false
|
| 4 |
+
debugInternalFunctionAccess = false // will also heavily slow down CoppeliaSim
|
| 5 |
+
debugCApiAccess = false // will also drastically slow down CoppeliaSim
|
| 6 |
+
debugLuaApiAccess = false // will also slow down CoppeliaSim
|
| 7 |
+
sendDebugInformationToFile = false // if true, debug info is sent to debugLog.txt
|
| 8 |
+
|
| 9 |
+
|
| 10 |
+
// Rendering
|
| 11 |
+
// =================================================
|
| 12 |
+
idleFps = 8
|
| 13 |
+
threadedRenderingDuringSimulation = 0 // recommended to keep 0 (-1=disabled, 0=pre-enabled, 1=enabled).
|
| 14 |
+
desiredOpenGlMajor = -1 // recommended to keep -1.
|
| 15 |
+
desiredOpenGlMinor = -1 // recommended to keep -1.
|
| 16 |
+
offscreenContextType = -1 // recommended to keep -1 (-1=default, 0=Qt offscreen, 1=QGLWidget/QOpenGLWidget visible, 2=QGLWidget/QOpenGLWidget invisible).
|
| 17 |
+
fboType = -1 // recommended to keep -1 (-1=default, 0=native, 1=QOpenGLFramebufferObject).
|
| 18 |
+
forceFboViaExt = false // recommended to keep false.
|
| 19 |
+
vboOperation = -1 // recommended to keep -1 (-1=default, 0=always off, 1=on when available).
|
| 20 |
+
vboPersistenceInMs = 5000 // recommended to keep 5000.
|
| 21 |
+
oglCompatibilityTweak1 = false // recommended to keep false since it causes small memory leaks.
|
| 22 |
+
visionSensorsUseGuiThread_windowed = -1 // recommended to keep -1 (-1=default, 0=GUI when not otherwise possible, 1=always GUI).
|
| 23 |
+
visionSensorsUseGuiThread_headless = -1 // recommended to keep -1 (-1=default, 0=GUI when not otherwise possible, 1=always GUI).
|
| 24 |
+
useGlFinish = false // recommended to keep false. Graphic card dependent.
|
| 25 |
+
useGlFinish_visionSensors = false // recommended to keep false. Graphic card dependent.
|
| 26 |
+
vsync = 0 // recommended to keep at 0. Graphic card dependent.
|
| 27 |
+
debugOpenGl = false
|
| 28 |
+
stereoDist = 0 // 0=no stereo, otherwise the intra occular distance (0.0635 for the human eyes).
|
| 29 |
+
highResDisplay = -1 // -1=automatic, 0=disabled, 1=enabled.
|
| 30 |
+
noEdgesWhenMouseDownInCameraView = false // if true, rendering is faster during mouse/view interaction
|
| 31 |
+
noTexturesWhenMouseDownInCameraView = false // if true, rendering is faster during mouse/view interaction
|
| 32 |
+
noCustomUisWhenMouseDownInCameraView = true // if true, rendering is faster during mouse/view interaction
|
| 33 |
+
hierarchyRefreshCnt = 3
|
| 34 |
+
|
| 35 |
+
|
| 36 |
+
// Visual
|
| 37 |
+
// =================================================
|
| 38 |
+
renderingSurfaceVShift = 0
|
| 39 |
+
renderingSurfaceVResize = 0
|
| 40 |
+
displayWorldRef = true
|
| 41 |
+
antialiasing = false
|
| 42 |
+
displayBoundingBoxeWhenObjectSelected = true
|
| 43 |
+
guiFontSize_Win = -1 // -1=default
|
| 44 |
+
guiFontSize_Mac = -1 // -1=default
|
| 45 |
+
guiFontSize_Linux = 13 // -1=default
|
| 46 |
+
statusbarInitiallyVisible = true
|
| 47 |
+
modelBrowserInitiallyVisible = true
|
| 48 |
+
sceneHierarchyInitiallyVisible = true
|
| 49 |
+
sceneHierarchyHiddenDuringSimulation = false
|
| 50 |
+
scriptEditorFont = // empty=default.
|
| 51 |
+
scriptEditorFontSize = -1 // -1=default.
|
| 52 |
+
auxConsoleFont = // empty=default
|
| 53 |
+
auxConsoleFontSize = -1 // -1=default.
|
| 54 |
+
mainScriptColor_background = 190,175,175
|
| 55 |
+
mainScriptColor_selection = 128,128,255
|
| 56 |
+
mainScriptColor_comment = 0,140,0
|
| 57 |
+
mainScriptColor_number = 220,0,220
|
| 58 |
+
mainScriptColor_string = 255,255,0
|
| 59 |
+
mainScriptColor_character = 255,255,0
|
| 60 |
+
mainScriptColor_operator = 0,0,0
|
| 61 |
+
mainScriptColor_preprocessor = 0,128,128
|
| 62 |
+
mainScriptColor_identifier = 64,64,64
|
| 63 |
+
mainScriptColor_word = 0,0,255
|
| 64 |
+
mainScriptColor_word2 = 152,0,0
|
| 65 |
+
mainScriptColor_word3 = 220,80,20
|
| 66 |
+
mainScriptColor_word4 = 152,64,0
|
| 67 |
+
nonThreadedChildScriptColor_background = 190,190,190
|
| 68 |
+
nonThreadedChildScriptColor_selection = 128,128,255
|
| 69 |
+
nonThreadedChildScriptColor_comment = 0,140,0
|
| 70 |
+
nonThreadedChildScriptColor_number = 220,0,220
|
| 71 |
+
nonThreadedChildScriptColor_string = 255,255,0
|
| 72 |
+
nonThreadedChildScriptColor_character = 255,255,0
|
| 73 |
+
nonThreadedChildScriptColor_operator = 0,0,0
|
| 74 |
+
nonThreadedChildScriptColor_preprocessor = 0,128,128
|
| 75 |
+
nonThreadedChildScriptColor_identifier = 64,64,64
|
| 76 |
+
nonThreadedChildScriptColor_word = 0,0,255
|
| 77 |
+
nonThreadedChildScriptColor_word2 = 152,0,0
|
| 78 |
+
nonThreadedChildScriptColor_word3 = 220,80,20
|
| 79 |
+
nonThreadedChildScriptColor_word4 = 152,64,0
|
| 80 |
+
threadedChildScriptColor_background = 180,190,190
|
| 81 |
+
threadedChildScriptColor_selection = 128,128,255
|
| 82 |
+
threadedChildScriptColor_comment = 0,140,0
|
| 83 |
+
threadedChildScriptColor_number = 220,0,220
|
| 84 |
+
threadedChildScriptColor_string = 255,255,0
|
| 85 |
+
threadedChildScriptColor_character = 255,255,0
|
| 86 |
+
threadedChildScriptColor_operator = 0,0,0
|
| 87 |
+
threadedChildScriptColor_preprocessor = 0,128,128
|
| 88 |
+
threadedChildScriptColor_identifier = 64,64,64
|
| 89 |
+
threadedChildScriptColor_word = 0,0,255
|
| 90 |
+
threadedChildScriptColor_word2 = 152,0,0
|
| 91 |
+
threadedChildScriptColor_word3 = 220,80,20
|
| 92 |
+
threadedChildScriptColor_word4 = 152,64,0
|
| 93 |
+
customizationScriptColor_background = 150,150,150
|
| 94 |
+
customizationScriptColor_selection = 128,128,255
|
| 95 |
+
customizationScriptColor_comment = 0,120,0
|
| 96 |
+
customizationScriptColor_number = 220,0,220
|
| 97 |
+
customizationScriptColor_string = 205,205,0
|
| 98 |
+
customizationScriptColor_character = 205,205,0
|
| 99 |
+
customizationScriptColor_operator = 0,0,0
|
| 100 |
+
customizationScriptColor_preprocessor = 0,128,128
|
| 101 |
+
customizationScriptColor_identifier = 70,70,70
|
| 102 |
+
customizationScriptColor_word = 0,0,255
|
| 103 |
+
customizationScriptColor_word2 = 152,0,0
|
| 104 |
+
customizationScriptColor_word3 = 220,80,20
|
| 105 |
+
customizationScriptColor_word4 = 152,64,0
|
| 106 |
+
|
| 107 |
+
|
| 108 |
+
// Directories
|
| 109 |
+
// =================================================
|
| 110 |
+
defaultDirectoryForScenes = // absolute path, e.g. d:/myScenes (or leave empty for default path)
|
| 111 |
+
defaultDirectoryForModels = // absolute path, e.g. d:/myModels (or leave empty for default path)
|
| 112 |
+
defaultDirectoryForCadFiles = // absolute path, e.g. d:/myCadFiles (or leave empty for default path)
|
| 113 |
+
defaultDirectoryForMiscFiles = // absolute path, e.g. d:/myMiscFiles (or leave empty for default path)
|
| 114 |
+
defaultDirectoryForExternalScriptEditor = // absolute path, e.g. d:/myScriptTempFiles (or leave empty for default path)
|
| 115 |
+
defaultDirectoryForRemoteApiFiles = // absolute path, e.g. d:/myRemoteApiTransfers (or leave empty for default path)
|
| 116 |
+
|
| 117 |
+
|
| 118 |
+
// Serialization
|
| 119 |
+
// =================================================
|
| 120 |
+
autoSaveDelay = 2 // in minutes. 0 to disable.
|
| 121 |
+
doNotWritePersistentData = false
|
| 122 |
+
compressFiles = true
|
| 123 |
+
saveApproxNormals = true
|
| 124 |
+
packIndices = true
|
| 125 |
+
fileDialogs = -1 // recommended to keep -1 (-1=default, 0=native dialogs, 1=Qt dialogs).
|
| 126 |
+
bulletSerializationBuffer = 10000000
|
| 127 |
+
|
| 128 |
+
|
| 129 |
+
// Undo/Redo
|
| 130 |
+
// =================================================
|
| 131 |
+
undoRedoEnabled = true
|
| 132 |
+
undoRedoOnlyPartialWithCameras = true
|
| 133 |
+
undoRedoLevelCount = 16
|
| 134 |
+
undoRedoMaxBufferSize = 80000000
|
| 135 |
+
timeInMsForUndoRedoTooLongWarning = 200
|
| 136 |
+
|
| 137 |
+
|
| 138 |
+
// Messaging
|
| 139 |
+
// =================================================
|
| 140 |
+
doNotShowCrashRecoveryMessage = false
|
| 141 |
+
doNotShowUpdateCheckMessage = false
|
| 142 |
+
doNotShowSceneSelectionThumbnails = false
|
| 143 |
+
doNotShowProgressBars = false
|
| 144 |
+
doNotShowAcknowledgmentMessages = false
|
| 145 |
+
doNotShowVideoCompressionLibraryLoadError = false
|
| 146 |
+
redirectStatusbarMsgToConsoleInHeadlessMode = false
|
| 147 |
+
suppressStartupDialogs = false
|
| 148 |
+
suppressXmlOverwriteMsg = false
|
| 149 |
+
|
| 150 |
+
|
| 151 |
+
// Compatibility
|
| 152 |
+
// =================================================
|
| 153 |
+
middleMouseButtonSwitchesModes = false // has only an effect if navigationBackwardCompatibility is true.
|
| 154 |
+
navigationBackwardCompatibility = false // recommended to keep false.
|
| 155 |
+
colorAdjust_backCompatibility = 1 // recommended to keep 1.0
|
| 156 |
+
specificGpuTweak = false
|
| 157 |
+
enableOldCalcModuleGuis = false
|
| 158 |
+
useAlternateSerialPortRoutines = false
|
| 159 |
+
enableOpenGlBasedCustomUiEditor = false
|
| 160 |
+
changeScriptCodeForNewApiNotation = 1 // 1=enabled, 0=disabled.
|
| 161 |
+
supportOldApiNotation = true
|
| 162 |
+
enableOldMillObjects = false
|
| 163 |
+
enableOldMirrorObjects = false
|
| 164 |
+
|
| 165 |
+
|
| 166 |
+
// Various
|
| 167 |
+
// =================================================
|
| 168 |
+
mouseWheelZoomFactor = 1
|
| 169 |
+
dynamicActivityRange = 1000
|
| 170 |
+
objectTranslationStepSize = 0.0250000004
|
| 171 |
+
objectRotationStepSize = 5
|
| 172 |
+
processorCoreAffinity = 0 // recommended to keep 0 (-1:os default, 0:all threads on same core, m: affinity mask (bit1=core1, bit2=core2, etc.))
|
| 173 |
+
freeServerPortStart = 20000
|
| 174 |
+
freeServerPortRange = 2000
|
| 175 |
+
abortScriptExecutionButton = 3 // in seconds. Zero to disable.
|
| 176 |
+
triCountInOBB = 8
|
| 177 |
+
removeIdenticalVertices = true
|
| 178 |
+
identicalVerticesTolerance = 9.99999975e-05
|
| 179 |
+
removeIdenticalTriangles = true
|
| 180 |
+
triangleWindingCheck = true
|
| 181 |
+
runCustomizationScripts = true
|
| 182 |
+
test1 = false // recommended to keep false.
|
| 183 |
+
orderHierarchyAlphabetically = false
|
| 184 |
+
macChildDialogType = -1 // -1=default.
|
| 185 |
+
useExternalLuaLibrary = false // if true, will call all Lua functions via the simLua library ('simLua.dll', 'libsimLua.so' or 'libsimLua.dylib')
|
| 186 |
+
additionalLuaPath = // e.g. d:/myLuaRoutines
|
| 187 |
+
desktopRecordingIndex = 0
|
| 188 |
+
desktopRecordingWidth = -1 // -1=default.
|
| 189 |
+
externalScriptEditor =
|
| 190 |
+
xmlExportSplitSize = 0 // 0=generate a single file.
|
| 191 |
+
xmlExportKnownFormats = true // true=if several files are generated, mesh and image files are saved under known formats.
|
| 192 |
+
|
| 193 |
+
|
| 194 |
+
// Floating license / dongle license
|
| 195 |
+
// =================================================
|
| 196 |
+
floatingLicenseEnabled = false
|
| 197 |
+
floatingLicenseServer = 127.0.0.1
|
| 198 |
+
floatingLicensePort = 20249
|
| 199 |
+
keepDongleOpen = false
|
third_party/AnyBimanual/third_party/PyRep/tests/__init__.py
ADDED
|
File without changes
|
third_party/AnyBimanual/third_party/PyRep/tests/assets/cracker_box/textured_simple.obj
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
third_party/AnyBimanual/third_party/PyRep/tests/assets/cracker_box/textured_simple.obj.mtl
ADDED
|
@@ -0,0 +1,14 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#
|
| 2 |
+
# Wavefront material file
|
| 3 |
+
# Converted by Meshlab Group
|
| 4 |
+
#
|
| 5 |
+
|
| 6 |
+
newmtl material_0
|
| 7 |
+
Ka 0.200000 0.200000 0.200000
|
| 8 |
+
Kd 0.639216 0.639216 0.639216
|
| 9 |
+
Ks 1.000000 1.000000 1.000000
|
| 10 |
+
Tr 1.000000
|
| 11 |
+
illum 2
|
| 12 |
+
Ns 0.000000
|
| 13 |
+
map_Kd texture_map.png
|
| 14 |
+
|
third_party/AnyBimanual/third_party/PyRep/tests/assets/loadable_model.ttm
ADDED
|
Binary file (28.6 kB). View file
|
|
|
third_party/AnyBimanual/third_party/PyRep/tests/assets/test_mesh_bowl.obj
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
third_party/AnyBimanual/third_party/PyRep/tests/assets/wood_texture.jpg
ADDED
|
third_party/AnyBimanual/third_party/PyRep/tests/core.py
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import unittest
|
| 2 |
+
import pyrep
|
| 3 |
+
from pyrep import PyRep
|
| 4 |
+
from os import path
|
| 5 |
+
|
| 6 |
+
ASSET_DIR = path.join(path.dirname(path.abspath(__file__)), 'assets')
|
| 7 |
+
pyrep.testing = True
|
| 8 |
+
|
| 9 |
+
|
| 10 |
+
class TestCore(unittest.TestCase):
|
| 11 |
+
|
| 12 |
+
def setUp(self):
|
| 13 |
+
self.pyrep = PyRep()
|
| 14 |
+
self.pyrep.launch(path.join(ASSET_DIR, 'test_scene.ttt'), headless=True)
|
| 15 |
+
self.pyrep.step()
|
| 16 |
+
self.pyrep.start()
|
| 17 |
+
|
| 18 |
+
def tearDown(self):
|
| 19 |
+
self.pyrep.stop()
|
| 20 |
+
self.pyrep.step_ui()
|
| 21 |
+
self.pyrep.shutdown()
|
third_party/AnyBimanual/third_party/PyRep/tests/test_accelerometer.py
ADDED
|
@@ -0,0 +1,18 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import unittest
|
| 2 |
+
from tests.core import TestCore
|
| 3 |
+
from pyrep.sensors.accelerometer import Accelerometer
|
| 4 |
+
|
| 5 |
+
|
| 6 |
+
class TestAccelerometer(TestCore):
|
| 7 |
+
|
| 8 |
+
def setUp(self):
|
| 9 |
+
super().setUp()
|
| 10 |
+
self.sensor = Accelerometer('accelerometer')
|
| 11 |
+
|
| 12 |
+
def test_read(self):
|
| 13 |
+
accelerations = self.sensor.read()
|
| 14 |
+
self.assertEqual(len(accelerations), 3)
|
| 15 |
+
|
| 16 |
+
|
| 17 |
+
if __name__ == '__main__':
|
| 18 |
+
unittest.main()
|
third_party/AnyBimanual/third_party/PyRep/tests/test_arms_and_configuration_paths.py
ADDED
|
@@ -0,0 +1,188 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import unittest
|
| 2 |
+
from tests.core import TestCore
|
| 3 |
+
from pyrep import PyRep
|
| 4 |
+
from pyrep.objects.cartesian_path import CartesianPath
|
| 5 |
+
from pyrep.objects.dummy import Dummy
|
| 6 |
+
import numpy as np
|
| 7 |
+
from os import path
|
| 8 |
+
from pyrep.errors import ConfigurationPathError
|
| 9 |
+
|
| 10 |
+
from pyrep.robots.arms.panda import Panda
|
| 11 |
+
from pyrep.robots.arms.mico import Mico
|
| 12 |
+
from pyrep.robots.arms.jaco import Jaco
|
| 13 |
+
from pyrep.robots.arms.sawyer import Sawyer
|
| 14 |
+
from pyrep.robots.arms.baxter import BaxterLeft, BaxterRight
|
| 15 |
+
from pyrep.robots.arms.lbr_iiwa_7_r800 import LBRIwaa7R800
|
| 16 |
+
from pyrep.robots.arms.lbr_iiwa_14_r820 import LBRIwaa14R820
|
| 17 |
+
from pyrep.robots.arms.ur3 import UR3
|
| 18 |
+
from pyrep.robots.arms.ur5 import UR5
|
| 19 |
+
from pyrep.robots.arms.ur10 import UR10
|
| 20 |
+
from pyrep.robots.arms.dobot import Dobot
|
| 21 |
+
from pyrep.robots.arms.xarm7 import XArm7
|
| 22 |
+
|
| 23 |
+
ASSET_DIR = path.join(path.dirname(path.abspath(__file__)), 'assets')
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
ARMS = [
|
| 27 |
+
('Panda', Panda),
|
| 28 |
+
('Sawyer', Sawyer),
|
| 29 |
+
('Mico', Mico),
|
| 30 |
+
('Jaco', Jaco),
|
| 31 |
+
('Baxter_leftArm', BaxterLeft),
|
| 32 |
+
('Baxter_rightArm', BaxterRight),
|
| 33 |
+
('LBR_iiwa_7_R800', LBRIwaa7R800),
|
| 34 |
+
('LBR_iiwa_14_R820', LBRIwaa14R820),
|
| 35 |
+
('UR3', UR3),
|
| 36 |
+
('UR5', UR5),
|
| 37 |
+
('UR10', UR10),
|
| 38 |
+
('Dobot', Dobot),
|
| 39 |
+
('XArm7', XArm7),
|
| 40 |
+
]
|
| 41 |
+
|
| 42 |
+
|
| 43 |
+
class TestArmsAndConfigurationPaths(TestCore):
|
| 44 |
+
|
| 45 |
+
def setUp(self):
|
| 46 |
+
self.pyrep = PyRep()
|
| 47 |
+
self.pyrep.launch(path.join(
|
| 48 |
+
ASSET_DIR, 'test_scene_robots.ttt'), headless=True)
|
| 49 |
+
self.pyrep.step()
|
| 50 |
+
self.pyrep.start()
|
| 51 |
+
|
| 52 |
+
# It is enough to only test the constructor of each arm (in there we make
|
| 53 |
+
# assumptions about the structure of the arm model). All other tests
|
| 54 |
+
# can be run on one arm.
|
| 55 |
+
def test_get_arm(self):
|
| 56 |
+
for arm_name, arm_type in ARMS:
|
| 57 |
+
with self.subTest(arm=arm_name):
|
| 58 |
+
arm = arm_type()
|
| 59 |
+
self.assertIsInstance(arm, arm_type)
|
| 60 |
+
|
| 61 |
+
def test_set_ik_element_properties(self):
|
| 62 |
+
arm = Panda()
|
| 63 |
+
arm.set_ik_element_properties(constraint_gamma=False)
|
| 64 |
+
arm.set_ik_element_properties()
|
| 65 |
+
|
| 66 |
+
def test_set_ik_group_properties(self):
|
| 67 |
+
arm = Panda()
|
| 68 |
+
arm.set_ik_group_properties('damped_least_squares')
|
| 69 |
+
arm.set_ik_group_properties('pseudo_inverse')
|
| 70 |
+
|
| 71 |
+
def test_solve_ik_via_jacobian(self):
|
| 72 |
+
arm = Panda()
|
| 73 |
+
waypoint = Dummy('Panda_waypoint')
|
| 74 |
+
new_config = arm.solve_ik_via_jacobian(
|
| 75 |
+
waypoint.get_position(), waypoint.get_orientation())
|
| 76 |
+
arm.set_joint_positions(new_config)
|
| 77 |
+
self.assertTrue(np.allclose(
|
| 78 |
+
arm.get_tip().get_pose(), waypoint.get_pose(), atol=0.001))
|
| 79 |
+
|
| 80 |
+
def test_solve_ik_via_sampling(self):
|
| 81 |
+
arm = Panda()
|
| 82 |
+
waypoint = Dummy('Panda_waypoint')
|
| 83 |
+
configs = arm.solve_ik_via_sampling(
|
| 84 |
+
waypoint.get_position(), waypoint.get_orientation(), max_configs=5)
|
| 85 |
+
self.assertIsNotNone(configs)
|
| 86 |
+
self.assertEqual(len(configs), 5)
|
| 87 |
+
current_config = arm.get_joint_positions()
|
| 88 |
+
# Checks correct config (last)
|
| 89 |
+
arm.set_joint_positions(configs[-1])
|
| 90 |
+
self.assertTrue(np.allclose(
|
| 91 |
+
arm.get_tip().get_pose(), waypoint.get_pose(), atol=0.001))
|
| 92 |
+
# Checks correct config (first)
|
| 93 |
+
arm.set_joint_positions(configs[0])
|
| 94 |
+
self.assertTrue(np.allclose(
|
| 95 |
+
arm.get_tip().get_pose(), waypoint.get_pose(), atol=0.001))
|
| 96 |
+
# Checks order
|
| 97 |
+
prev_config_dist = 0
|
| 98 |
+
for config in configs:
|
| 99 |
+
config_dist = sum(
|
| 100 |
+
[(c - f)**2 for c, f in zip(current_config, config)])
|
| 101 |
+
# This test requires that the metric scale for each joint remains at
|
| 102 |
+
# 1.0 in _getConfigDistance lua function
|
| 103 |
+
self.assertLessEqual(prev_config_dist, config_dist)
|
| 104 |
+
prev_config_dist = config_dist
|
| 105 |
+
|
| 106 |
+
def test_get_path_from_cartesian_path(self):
|
| 107 |
+
arm = Panda()
|
| 108 |
+
cartesian_path = CartesianPath('Panda_cartesian_path')
|
| 109 |
+
path = arm.get_path_from_cartesian_path(cartesian_path)
|
| 110 |
+
self.assertIsNotNone(path)
|
| 111 |
+
|
| 112 |
+
def test_get_linear_path(self):
|
| 113 |
+
arm = Panda()
|
| 114 |
+
waypoint = Dummy('Panda_waypoint')
|
| 115 |
+
path = arm.get_linear_path(
|
| 116 |
+
waypoint.get_position(), waypoint.get_orientation())
|
| 117 |
+
self.assertIsNotNone(path)
|
| 118 |
+
|
| 119 |
+
def test_get_linear_path_relative(self):
|
| 120 |
+
arm = Panda()
|
| 121 |
+
path = arm.get_linear_path([0, 0, 0.01], [0, 0, 0],
|
| 122 |
+
relative_to=arm.get_tip())
|
| 123 |
+
self.assertIsNotNone(path)
|
| 124 |
+
|
| 125 |
+
def test_get_nonlinear_path(self):
|
| 126 |
+
arm = Panda()
|
| 127 |
+
waypoint = Dummy('Panda_waypoint')
|
| 128 |
+
path = arm.get_nonlinear_path(
|
| 129 |
+
waypoint.get_position(), waypoint.get_orientation())
|
| 130 |
+
self.assertIsNotNone(path)
|
| 131 |
+
|
| 132 |
+
def test_get_nonlinear_path_out_of_reach(self):
|
| 133 |
+
arm = Panda()
|
| 134 |
+
with self.assertRaises(ConfigurationPathError):
|
| 135 |
+
arm.get_nonlinear_path([99, 99, 99], [0.] * 3)
|
| 136 |
+
|
| 137 |
+
def test_get_linear_path_out_of_reach(self):
|
| 138 |
+
arm = Panda()
|
| 139 |
+
with self.assertRaises(ConfigurationPathError):
|
| 140 |
+
arm.get_linear_path([99, 99, 99], [0.] * 3)
|
| 141 |
+
|
| 142 |
+
def test_get_linear_path_and_step(self):
|
| 143 |
+
arm = Panda()
|
| 144 |
+
waypoint = Dummy('Panda_waypoint')
|
| 145 |
+
path = arm.get_linear_path(
|
| 146 |
+
waypoint.get_position(), waypoint.get_orientation())
|
| 147 |
+
self.assertIsNotNone(path)
|
| 148 |
+
done = False
|
| 149 |
+
while not done:
|
| 150 |
+
done = path.step()
|
| 151 |
+
self.pyrep.step()
|
| 152 |
+
self.assertTrue(np.allclose(
|
| 153 |
+
arm.get_tip().get_position(), waypoint.get_position(), atol=0.01))
|
| 154 |
+
|
| 155 |
+
def test_get_linear_path_and_get_end(self):
|
| 156 |
+
arm = Panda()
|
| 157 |
+
waypoint = Dummy('Panda_waypoint')
|
| 158 |
+
path = arm.get_linear_path(
|
| 159 |
+
waypoint.get_position(), waypoint.get_orientation())
|
| 160 |
+
path.set_to_end()
|
| 161 |
+
self.assertTrue(np.allclose(
|
| 162 |
+
arm.get_tip().get_position(), waypoint.get_position(), atol=0.001))
|
| 163 |
+
|
| 164 |
+
def test_get_linear_path_visualize(self):
|
| 165 |
+
arm = Panda()
|
| 166 |
+
waypoint = Dummy('Panda_waypoint')
|
| 167 |
+
path = arm.get_linear_path(
|
| 168 |
+
waypoint.get_position(), waypoint.get_orientation())
|
| 169 |
+
# Check that it does not error
|
| 170 |
+
path.visualize()
|
| 171 |
+
|
| 172 |
+
def test_get_duplicate_arm(self):
|
| 173 |
+
arm = UR3(1)
|
| 174 |
+
self.assertIsInstance(arm, UR3)
|
| 175 |
+
|
| 176 |
+
def test_copy_arm(self):
|
| 177 |
+
arm = UR10()
|
| 178 |
+
new_arm = arm.copy()
|
| 179 |
+
self.assertNotEqual(arm, new_arm)
|
| 180 |
+
|
| 181 |
+
def test_get_jacobian(self):
|
| 182 |
+
arm = Panda()
|
| 183 |
+
jacobian = arm.get_jacobian()
|
| 184 |
+
self.assertEqual(jacobian.shape, (7, 6))
|
| 185 |
+
|
| 186 |
+
|
| 187 |
+
if __name__ == '__main__':
|
| 188 |
+
unittest.main()
|
third_party/AnyBimanual/third_party/PyRep/tests/test_cameras.py
ADDED
|
@@ -0,0 +1,38 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import unittest
|
| 2 |
+
from tests.core import TestCore
|
| 3 |
+
from pyrep.const import ObjectType
|
| 4 |
+
from pyrep.objects.camera import Camera
|
| 5 |
+
from pyrep.objects.dummy import Dummy
|
| 6 |
+
from pyrep.objects.object import Object
|
| 7 |
+
import numpy as np
|
| 8 |
+
|
| 9 |
+
|
| 10 |
+
class TestCameras(TestCore):
|
| 11 |
+
|
| 12 |
+
def setUp(self):
|
| 13 |
+
super().setUp()
|
| 14 |
+
self.camera = Camera('DefaultCamera')
|
| 15 |
+
self.dummy = Dummy('dummy')
|
| 16 |
+
|
| 17 |
+
def test_create(self):
|
| 18 |
+
with self.assertRaises(NotImplementedError):
|
| 19 |
+
_ = Camera.create()
|
| 20 |
+
|
| 21 |
+
def test_get_set_position(self):
|
| 22 |
+
position = self.camera.get_position()
|
| 23 |
+
self.assertIsInstance(position, np.ndarray)
|
| 24 |
+
self.assertEqual(position.shape, (3,))
|
| 25 |
+
|
| 26 |
+
self.camera.set_position([0.1, 0.1, 0.1], self.dummy)
|
| 27 |
+
self.assertTrue(np.allclose(
|
| 28 |
+
self.camera.get_position(self.dummy), [0.1, 0.1, 0.1]))
|
| 29 |
+
self.camera.set_position([0.2, 0.2, 0.2])
|
| 30 |
+
self.assertTrue(np.allclose(self.camera.get_position(), [0.2, 0.2, 0.2]))
|
| 31 |
+
|
| 32 |
+
def test_get_object_type(self):
|
| 33 |
+
self.assertEqual(Object.get_object_type('DefaultCamera'),
|
| 34 |
+
ObjectType.CAMERA)
|
| 35 |
+
|
| 36 |
+
|
| 37 |
+
if __name__ == '__main__':
|
| 38 |
+
unittest.main()
|
third_party/AnyBimanual/third_party/PyRep/tests/test_cartesian_paths.py
ADDED
|
@@ -0,0 +1,28 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import unittest
|
| 2 |
+
from tests.core import TestCore
|
| 3 |
+
from pyrep.objects.cartesian_path import CartesianPath
|
| 4 |
+
|
| 5 |
+
|
| 6 |
+
class TestCartesianPaths(TestCore):
|
| 7 |
+
|
| 8 |
+
def setUp(self):
|
| 9 |
+
super().setUp()
|
| 10 |
+
self.cart_path = CartesianPath('cartesian_path')
|
| 11 |
+
|
| 12 |
+
def test_create_cartesian_path(self):
|
| 13 |
+
p = CartesianPath.create()
|
| 14 |
+
self.assertIsInstance(p, CartesianPath)
|
| 15 |
+
|
| 16 |
+
def test_get_pose_on_path(self):
|
| 17 |
+
pos, ori = self.cart_path.get_pose_on_path(0.5)
|
| 18 |
+
self.assertEqual(len(pos), 3)
|
| 19 |
+
self.assertEqual(len(ori), 3)
|
| 20 |
+
|
| 21 |
+
def test_insert_control_points(self):
|
| 22 |
+
points = [[0.1] * 6]
|
| 23 |
+
# Just check that it does not throw an exception.
|
| 24 |
+
self.cart_path.insert_control_points(points)
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
if __name__ == '__main__':
|
| 28 |
+
unittest.main()
|
third_party/AnyBimanual/third_party/PyRep/tests/test_dummies.py
ADDED
|
@@ -0,0 +1,17 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import unittest
|
| 2 |
+
from tests.core import TestCore
|
| 3 |
+
from pyrep.objects.dummy import Dummy
|
| 4 |
+
|
| 5 |
+
|
| 6 |
+
class TestDummies(TestCore):
|
| 7 |
+
|
| 8 |
+
def test_get_dummy(self):
|
| 9 |
+
Dummy('dummy')
|
| 10 |
+
|
| 11 |
+
def test_create_dummy(self):
|
| 12 |
+
d = Dummy.create(0.01)
|
| 13 |
+
self.assertIsInstance(d, Dummy)
|
| 14 |
+
|
| 15 |
+
|
| 16 |
+
if __name__ == '__main__':
|
| 17 |
+
unittest.main()
|
third_party/AnyBimanual/third_party/PyRep/tests/test_force_sensors.py
ADDED
|
@@ -0,0 +1,23 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import unittest
|
| 2 |
+
from tests.core import TestCore
|
| 3 |
+
from pyrep.objects.force_sensor import ForceSensor
|
| 4 |
+
|
| 5 |
+
|
| 6 |
+
class TestForceSensors(TestCore):
|
| 7 |
+
|
| 8 |
+
def setUp(self):
|
| 9 |
+
super().setUp()
|
| 10 |
+
self.sensor = ForceSensor('force_sensor')
|
| 11 |
+
|
| 12 |
+
def test_read(self):
|
| 13 |
+
force_vector, torque_vector = self.sensor.read()
|
| 14 |
+
self.assertEqual(len(force_vector), 3)
|
| 15 |
+
self.assertEqual(len(torque_vector), 3)
|
| 16 |
+
|
| 17 |
+
def test_create(self):
|
| 18 |
+
sensor = ForceSensor.create()
|
| 19 |
+
sensor.remove()
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
if __name__ == '__main__':
|
| 23 |
+
unittest.main()
|