| """The 6 base-policy eval tasks (the tasks the OSCAR 6-task LoRA was trained on). |
| |
| One entry per task: gym id, the STATELESS base-ACT checkpoint + dataset stats, |
| the OSCAR text prompt (MUST match oscar-public/finetune/captions.py, the caption |
| the skeleton LoRA was trained with), and the per-task rollout horizon. |
| |
| `n_exec` / `max_steps` mirror the per-task `<task>_oscar_wm.yaml` configs: |
| each WM `imagine` produces wm_chunk=25 frames; we keep the first `n_exec` and |
| re-plan, so the episode is `ceil(max_steps / n_exec)` WM calls long. |
| """ |
| from __future__ import annotations |
|
|
| BASE = "/root/pwm_local/data0/base_policies" |
|
|
| |
| _CAP = { |
| "Open-Drawer": "A humanoid robot with two dexterous multi-fingered hands opens a drawer on a wooden table in a bright indoor room.", |
| "Open-Laptop": "A humanoid robot with two dexterous multi-fingered hands opens a laptop on a wooden table in a bright indoor room.", |
| "Pour-Balls": "A humanoid robot with two dexterous multi-fingered hands pours small balls from a cup into a bowl on a wooden table in a bright indoor room.", |
| "Push-Box": "A humanoid robot with two dexterous multi-fingered hands pushes a box across a wooden table in a bright indoor room.", |
| "Stack-Can": "A humanoid robot with two dexterous multi-fingered hands stacks a soda can on top of another can on a wooden table in a bright indoor room.", |
| "Unload-Cans": "A humanoid robot with two dexterous multi-fingered hands unloads soda cans from a holder onto a wooden table in a bright indoor room.", |
| } |
|
|
| TASKS = { |
| "open_drawer": dict( |
| gym="Humanoid-Open-Drawer-v0", |
| ckpt=f"{BASE}/open_drawer_30demo_stateless/policy_epoch_2000_seed_0.ckpt", |
| stats=f"{BASE}/open_drawer_30demo_stateless/dataset_stats.pkl", |
| prompt=_CAP["Open-Drawer"], max_steps=300, n_exec=24, |
| table_idx=1, room_idx=1, scene_seed=6, |
| ), |
| "open_laptop": dict( |
| gym="Humanoid-Open-Laptop-v0", |
| ckpt=f"{BASE}/open_laptop_30demo_stateless/policy.ckpt", |
| stats=f"{BASE}/open_laptop_30demo_stateless/dataset_stats.pkl", |
| prompt=_CAP["Open-Laptop"], max_steps=300, n_exec=12, |
| table_idx=1, room_idx=1, scene_seed=8, |
| ), |
| "pour_balls": dict( |
| gym="Humanoid-Pour-Balls-v0", |
| ckpt=f"{BASE}/pour_balls_50demo_stateless/policy.ckpt", |
| stats=f"{BASE}/pour_balls_50demo_stateless/dataset_stats.pkl", |
| prompt=_CAP["Pour-Balls"], max_steps=400, n_exec=12, |
| table_idx=1, room_idx=1, scene_seed=3, |
| ), |
| "push_box": dict( |
| gym="Humanoid-Push-Box-v0", |
| ckpt=f"{BASE}/push_box_30demo_stateless/policy.ckpt", |
| stats=f"{BASE}/push_box_30demo_stateless/dataset_stats.pkl", |
| prompt=_CAP["Push-Box"], max_steps=300, n_exec=24, |
| table_idx=1, room_idx=1, scene_seed=5, |
| ), |
| "stack_can": dict( |
| gym="Humanoid-Stack-Can-v0", |
| ckpt=f"{BASE}/stack_can_50demo_stateless/policy.ckpt", |
| stats=f"{BASE}/stack_can_50demo_stateless/dataset_stats.pkl", |
| prompt=_CAP["Stack-Can"], max_steps=300, n_exec=12, |
| table_idx=1, room_idx=1, scene_seed=9, |
| ), |
| "unload_cans": dict( |
| gym="Humanoid-Unload-Cans-v0", |
| ckpt=f"{BASE}/unload_cans_50demo_stateless/policy.ckpt", |
| stats=f"{BASE}/unload_cans_50demo_stateless/dataset_stats.pkl", |
| prompt=_CAP["Unload-Cans"], max_steps=300, n_exec=24, |
| table_idx=1, room_idx=1, scene_seed=10, |
| ), |
| } |
|
|
| |
| INIT_POSES_PATH = "/root/pwm_local/PWM/EgoVLA_Release/init_poses_fixed_set_100traj.pkl" |
|
|
