File size: 5,413 Bytes
7399b6f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | """Run a registered YAM task: build scene -> solve -> evaluate -> video.
python scripts/yam_task.py --list
python scripts/yam_task.py --task grape_box --seed 0
python scripts/yam_task.py --task pot_dual_lift --seed 4 --no-randomize
One Python file per task under source/bimanual/yam/tasks/, in the shape ManiSkill uses: a
registered class with _load_scene / _initialize_episode / solve / evaluate.
"""
import argparse, json, os, sys
from pathlib import Path
REPO = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO/"source"))
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--task")
ap.add_argument("--list", action="store_true")
ap.add_argument("--seed", type=int, default=0)
ap.add_argument("--no-randomize", action="store_true")
ap.add_argument("--episode", type=int, default=-1)
ap.add_argument("--video", default="")
ap.add_argument("--set", action="append", default=[], metavar="KEY=VALUE",
help="override a task parameter, on top of tasks/configs/<task>.yaml")
ap.add_argument("--first-frame", default="",
help="build the scene, save ONE frame here, and exit without solving")
args, unknown = ap.parse_known_args()
from bimanual.yam import registry
import bimanual.yam.tasks # noqa: F401 (registers every task)
if args.list or not args.task:
rows = registry.list_tasks()
w = max([len(r["name"]) for r in rows]+[4])
print(f"{'name'.ljust(w)} {'class'.ljust(20)} tags")
for r in rows:
print(f"{r['name'].ljust(w)} {r['class'].ljust(20)} {r['tags']}")
print(f"\n{len(rows)} task(s) registered")
return 0
cls = registry.REGISTRY.get(args.task)
if cls is None:
raise SystemExit(f"unknown task {args.task!r}. Registered: {sorted(registry.REGISTRY)}")
# YAML + --set have to be resolved BEFORE env_vars(): which USDs get registered and how hard
# the gripper clamps are decided before the simulator starts, so a config that changed
# rw_objects would otherwise be read too late to have any effect.
from bimanual.yam.envs import config as taskcfg
overrides = dict(taskcfg.load(args.task))
overrides.update(taskcfg.parse_cli(args.set))
# asset registration and clamp force must be set BEFORE the simulator starts
for k, v in cls.env_vars(overrides).items():
os.environ[k] = v
print(f"[run] {k}={v}", flush=True)
os.environ.setdefault("OMNI_KIT_ACCEPT_EULA", "YES")
from isaaclab.app import AppLauncher
ap2 = argparse.ArgumentParser(); AppLauncher.add_app_launcher_args(ap2)
largs = ap2.parse_args(unknown); largs.headless = True; largs.enable_cameras = True
app = AppLauncher(largs).app
import gymnasium as gym
import bimanual.tasks.manager_based.yam # noqa: F401
from isaaclab_tasks.utils import parse_env_cfg
cfg = parse_env_cfg("Template-YAM-Play-v0", device="cuda:0", num_envs=1)
# one long scripted sequence: the 12 s episode limit would auto-reset mid-run and snap the
# arm back to its home joints
cfg.episode_length_s = 1.0e6
try:
cfg.terminations.time_out = None
except Exception as e:
print("[run] time_out disable failed:", e)
try:
cfg.viewer.eye = tuple(cls.viewer_eye); cfg.viewer.lookat = tuple(cls.viewer_lookat)
cfg.viewer.resolution = (720, 540)
except Exception as e:
print("[run] viewer cfg:", e)
env = gym.make("Template-YAM-Play-v0", cfg=cfg, render_mode="rgb_array")
env.reset()
origin = env.unwrapped.scene.env_origins[0].cpu().numpy()
task = cls(seed=args.seed, randomize=not args.no_randomize, episode=args.episode,
video=args.video, overrides=taskcfg.parse_cli(args.set))
task.build(env, origin)
if args.first_frame:
# Scene-only render: building is the expensive-but-short part, solving is the long part.
# This makes a 30-task overview minutes rather than hours.
import imageio.v2 as _iio
# The renderer needs warm-up: the frames right after a reset come back empty or with the
# wrong camera, so a single render() straight after build() writes a black image (all 29
# tiles of the first overview were black). Step AND render repeatedly, keep the last.
frame = None
for _ in range(60):
task.step()
img = env.render()
if img is not None:
frame = img
os.makedirs(os.path.dirname(args.first_frame) or ".", exist_ok=True)
_iio.imwrite(args.first_frame, frame)
print(f"[run] first frame -> {args.first_frame} {frame.shape}", flush=True)
env.close(); app.close()
print("YAM_TASK_DONE", flush=True)
return 0
# Re-baseline the joints RIGHT BEFORE solving. Recording them at build time counted the
# fixture's own settling as progress: switch_toggle "passed" on 0.44 rad of gravity and
# 0.015 rad of actual robot work.
task._start_joints = task._joint_state()
info = task.solve()
print(f"[run] solver: {json.dumps(info, default=str)}", flush=True)
result = task.evaluate()
task.save_video(args.video or None)
env.close(); app.close()
print("YAM_TASK_DONE", flush=True)
return 0 if result.get("success") else 1
if __name__ == "__main__":
sys.exit(main() or 0)
|