lsnu commited on
Commit
6e9fe6b
·
verified ·
1 Parent(s): c9bdd6d

Add files using upload-large-folder tool

Browse files
This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. code/VLAarchtests4/VLAarchtests2_code/VLAarchtests/code/reveal_vla_bimanual/eval/public_benchmark_package.py +365 -0
  2. code/VLAarchtests4/VLAarchtests2_code/VLAarchtests/code/reveal_vla_bimanual/eval/render_maniskill_pickclutter_benchmark_gifs.py +863 -0
  3. code/VLAarchtests4/VLAarchtests2_code/VLAarchtests/code/reveal_vla_bimanual/eval/run_dexgarmentlab_hybrid_smoke.py +215 -0
  4. code/VLAarchtests4/VLAarchtests2_code/VLAarchtests/code/reveal_vla_bimanual/eval/run_maniskill_bridge_retrieval_smoke.py +2037 -0
  5. code/VLAarchtests4/VLAarchtests2_code/VLAarchtests/code/reveal_vla_bimanual/eval/run_maniskill_pickclutter_smoke.py +2005 -0
  6. code/VLAarchtests4/VLAarchtests2_code/VLAarchtests/code/reveal_vla_bimanual/eval/run_public_benchmark_package.py +369 -0
  7. code/VLAarchtests4/VLAarchtests2_code/VLAarchtests/code/reveal_vla_bimanual/eval/run_rlbench_hybrid_smoke.py +269 -0
  8. third_party/AnyBimanual/third_party/PyRep/.gitignore +102 -0
  9. third_party/AnyBimanual/third_party/PyRep/LICENSE +21 -0
  10. third_party/AnyBimanual/third_party/PyRep/README.md +284 -0
  11. third_party/AnyBimanual/third_party/PyRep/pyrep/__init__.py +5 -0
  12. third_party/AnyBimanual/third_party/PyRep/pyrep/const.py +110 -0
  13. third_party/AnyBimanual/third_party/PyRep/pyrep/errors.py +30 -0
  14. third_party/AnyBimanual/third_party/PyRep/pyrep/misc/__init__.py +1 -0
  15. third_party/AnyBimanual/third_party/PyRep/pyrep/pyrep.py +332 -0
  16. third_party/AnyBimanual/third_party/PyRep/pyrep/robots/__init__.py +0 -0
  17. third_party/AnyBimanual/third_party/PyRep/pyrep/robots/end_effectors/xarm_gripper.py +13 -0
  18. third_party/AnyBimanual/third_party/PyRep/pyrep/robots/mobiles/holonomic_base.py +215 -0
  19. third_party/AnyBimanual/third_party/PyRep/pyrep/robots/mobiles/line_tracer.py +6 -0
  20. third_party/AnyBimanual/third_party/PyRep/pyrep/robots/mobiles/locobot.py +6 -0
  21. third_party/AnyBimanual/third_party/PyRep/pyrep/robots/mobiles/mobile_base.py +202 -0
  22. third_party/AnyBimanual/third_party/PyRep/pyrep/robots/mobiles/nonholonomic_base.py +140 -0
  23. third_party/AnyBimanual/third_party/PyRep/pyrep/robots/mobiles/turtlebot.py +6 -0
  24. third_party/AnyBimanual/third_party/PyRep/pyrep/robots/mobiles/youbot.py +7 -0
  25. third_party/AnyBimanual/third_party/PyRep/pyrep/robots/robot_component.py +290 -0
  26. third_party/AnyBimanual/third_party/PyRep/pyrep/sensors/__init__.py +0 -0
  27. third_party/AnyBimanual/third_party/PyRep/pyrep/sensors/accelerometer.py +30 -0
  28. third_party/AnyBimanual/third_party/PyRep/pyrep/sensors/gyroscope.py +40 -0
  29. third_party/AnyBimanual/third_party/PyRep/pyrep/sensors/spherical_vision_sensor.py +296 -0
  30. third_party/AnyBimanual/third_party/PyRep/pyrep/textures/__init__.py +0 -0
  31. third_party/AnyBimanual/third_party/PyRep/pyrep/textures/texture.py +18 -0
  32. third_party/AnyBimanual/third_party/PyRep/requirements.txt +2 -0
  33. third_party/AnyBimanual/third_party/PyRep/robot_ttms/README.md +5 -0
  34. third_party/AnyBimanual/third_party/PyRep/robot_ttms/end_effectors/BaxterSuctionCup.ttm +0 -0
  35. third_party/AnyBimanual/third_party/PyRep/setup.py +64 -0
  36. third_party/AnyBimanual/third_party/PyRep/system/settings.dat +0 -0
  37. third_party/AnyBimanual/third_party/PyRep/system/usrset.txt +199 -0
  38. third_party/AnyBimanual/third_party/PyRep/tests/__init__.py +0 -0
  39. third_party/AnyBimanual/third_party/PyRep/tests/assets/cracker_box/textured_simple.obj +0 -0
  40. third_party/AnyBimanual/third_party/PyRep/tests/assets/cracker_box/textured_simple.obj.mtl +14 -0
  41. third_party/AnyBimanual/third_party/PyRep/tests/assets/loadable_model.ttm +0 -0
  42. third_party/AnyBimanual/third_party/PyRep/tests/assets/test_mesh_bowl.obj +0 -0
  43. third_party/AnyBimanual/third_party/PyRep/tests/assets/wood_texture.jpg +0 -0
  44. third_party/AnyBimanual/third_party/PyRep/tests/core.py +21 -0
  45. third_party/AnyBimanual/third_party/PyRep/tests/test_accelerometer.py +18 -0
  46. third_party/AnyBimanual/third_party/PyRep/tests/test_arms_and_configuration_paths.py +188 -0
  47. third_party/AnyBimanual/third_party/PyRep/tests/test_cameras.py +38 -0
  48. third_party/AnyBimanual/third_party/PyRep/tests/test_cartesian_paths.py +28 -0
  49. third_party/AnyBimanual/third_party/PyRep/tests/test_dummies.py +17 -0
  50. 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 [![Build Status](https://github.com/stepjam/PyRep/workflows/Build/badge.svg)](https://github.com/stepjam/PyRep/actions) [![Discord](https://img.shields.io/discord/694945313638842378.svg?label=&logo=discord&logoColor=ffffff&color=7389D8&labelColor=6A7EC2)](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()