diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..d15da12388a5c7a1a0a458745106e5027870a995 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,3 +33,43 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.zip filter=lfs diff=lfs merge=lfs -text *.zst filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/backgrounds/indoors/machine_shop_01_2k.hdr filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/factory_bolt_m12_loose.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/factory_bolt_m12_tight.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/factory_bolt_m16_loose.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/factory_bolt_m16_tight.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/factory_bolt_m20_loose.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/factory_bolt_m20_tight.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/factory_bolt_m8_loose.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/factory_bolt_m8_tight.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/factory_nut_m12_loose.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/factory_nut_m12_tight.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/factory_nut_m16_loose.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/factory_nut_m16_tight.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/factory_nut_m20_loose.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/factory_nut_m20_tight.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/factory_nut_m8_loose.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/factory_nut_m8_tight.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/gen_bolt_m12_loose.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/gen_bolt_m12_tight.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/gen_bolt_m16_loose.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/gen_bolt_m16_tight.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/gen_bolt_m20_loose.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/gen_bolt_m20_tight.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/gen_bolt_m4_loose.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/gen_bolt_m4_tight.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/gen_bolt_m8_loose.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/gen_bolt_m8_tight.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/gen_gear_large.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/gen_gear_medium.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/gen_gear_small.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/gen_nut_m12_loose.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/gen_nut_m12_tight.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/gen_nut_m16_loose.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/gen_nut_m16_tight.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/gen_nut_m20_loose.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/gen_nut_m20_tight.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/gen_nut_m4_loose.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/gen_nut_m4_tight.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/gen_nut_m8_loose.usd filter=lfs diff=lfs merge=lfs -text +assembly_bench/assets/parts/gen_nut_m8_tight.usd filter=lfs diff=lfs merge=lfs -text diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..dcd349280473c45d550f2d25a612436e2ea1410e --- /dev/null +++ b/.gitignore @@ -0,0 +1,5 @@ +__pycache__/ +*.egg-info/ +.venv/ +*.etli +.cache/ diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..1fcd3de47f247e09faff88004061db5c335a4f0b --- /dev/null +++ b/README.md @@ -0,0 +1,111 @@ +# env-assembly-bench + +[EnvHub](https://huggingface.co/docs/lerobot/en/envhub) package for the NIST +taskboard **assembly** benchmark on +[Isaac Lab Arena](https://github.com/isaac-sim/IsaacLab-Arena). + +Load with one line (executes Hub code — pin a commit in production): + +```python +from lerobot.envs import make_env +from lerobot.envs.configs import IsaaclabArenaEnv + +cfg = IsaaclabArenaEnv( + hub_path="lukasskellijs/env_assembly_bench", + environment="assembly_bench", + embodiment="droid_abs_joint_pos_softmimic", + headless=True, + enable_cameras=True, + state_keys="joint_pos,gripper_pos,eef_pos,eef_quat", + camera_keys="front_cam_rgb,wrist_camera_rgb", + state_dim=15, + action_dim=8, + enable_pinocchio=False, + kwargs={"variant": "peg_round_M1_loose", "reward": "none"}, +) +suites = make_env(cfg, n_envs=1, trust_remote_code=True) +env = next(iter(next(iter(suites.values())).values())) +``` + +`cfg.task` is the natural-language prompt (LeRobot convention). The Arena +variant id is `kwargs.variant` (for example `peg_round_M1_loose`). + +## Prerequisites + +Same stack as [Isaac Lab Arena + EnvHub](https://huggingface.co/docs/lerobot/en/envhub_isaaclab_arena): + +- NVIDIA GPU, Isaac Sim 6.x, Isaac Lab + Isaac Lab Arena installed locally +- `pip install -r requirements.txt` (Hub glue only) +- Optional: `pip install -e .` so `assembly_bench` imports without a Hub snapshot + +EULA (once per shell): + +```bash +export ACCEPT_EULA=Y PRIVACY_CONSENT=Y OMNI_KIT_ACCEPT_EULA=YES +``` + +## Variants + +15 loose-clearance tasks plus `debug` (apple→bowl). Peg sizes use suite codes +`S/M1/M2/L` (4/8/12/16 mm): + +| family | examples | +|---|---| +| peg insert | `peg_round_M1_loose`, `peg_square_L_loose` | +| gear mesh | `gear_small`, `gear_medium`, `gear_large` | +| nut thread | `nut_M8`, `nut_M12`, `nut_M16`, `nut_M20` | + +Default embodiment: `droid_abs_joint_pos_softmimic` (Franka + Robotiq, 8-D +absolute joint targets, DROID cameras). + +## LeRobot eval + +```bash +lerobot-eval \ + --env.type=isaaclab_arena \ + --env.hub_path=lukasskellijs/env_assembly_bench \ + --env.environment=assembly_bench \ + --env.embodiment=droid_abs_joint_pos_softmimic \ + --env.headless=true \ + --env.enable_cameras=true \ + --env.state_keys=joint_pos,gripper_pos,eef_pos,eef_quat \ + --env.camera_keys=front_cam_rgb,wrist_camera_rgb \ + --env.state_dim=15 \ + --env.action_dim=8 \ + --env.enable_pinocchio=false \ + --env.kwargs='{"variant": "peg_round_M1_loose"}' \ + --trust_remote_code=True \ + --eval.batch_size=1 \ + --policy.path= +``` + +## HUD + +[`hud_env.py`](hud_env.py) is the standard HUD robot entry: it snapshot-downloads +this Hub repo, calls `env.py:make_env`, and exposes the sim via `env.gym(...)`. + +```bash +OMNI_KIT_ACCEPT_EULA=YES ACCEPT_EULA=Y PRIVACY_CONSENT=Y \ + python -m hud.environment.server hud_env.py --port 8765 +``` + +Point a `RobotAgent` / `Runtime("tcp://127.0.0.1:8765")` at that port as in the +[HUD robots guide](https://hud.ai). + +## Layout + +``` +env_assembly_bench/ +├── env.py # EnvHub make_env +├── hud_env.py # HUD Environment (Hub download path) +├── isaaclab_env_wrapper.py +├── example_envs.yaml +├── contract.json # HUD obs/action contract +└── assembly_bench/ # Arena package + USD/HDR assets +``` + +## Security + +`trust_remote_code=True` runs Hub Python. Review `env.py`, pin +`lukasskellijs/env_assembly_bench@`, and prefer this org only if you +trust it. diff --git a/__pycache__/env.cpython-312.pyc b/__pycache__/env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8e31041b85ff836feb4f37c66885c904018fb876 Binary files /dev/null and b/__pycache__/env.cpython-312.pyc differ diff --git a/assembly_bench/__init__.py b/assembly_bench/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/assembly_bench/__pycache__/__init__.cpython-312.pyc b/assembly_bench/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..60dc7020f3ca54d866106f9777baf48f99474cf1 Binary files /dev/null and b/assembly_bench/__pycache__/__init__.cpython-312.pyc differ diff --git a/assembly_bench/assets/backgrounds/indoors/machine_shop_01_2k.hdr b/assembly_bench/assets/backgrounds/indoors/machine_shop_01_2k.hdr new file mode 100644 index 0000000000000000000000000000000000000000..14dd907b731d1fb086e7bca1c690d384268a3bc7 --- /dev/null +++ b/assembly_bench/assets/backgrounds/indoors/machine_shop_01_2k.hdr @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2873e0e8a7fdff45d1815f4a8c3cb60c9f425b74c88d7981bf14c32d0464f29a +size 7004194 diff --git a/assembly_bench/assets/parts/factory_bolt_m12_loose.usd b/assembly_bench/assets/parts/factory_bolt_m12_loose.usd new file mode 100644 index 0000000000000000000000000000000000000000..fd98999090d0f90af743b852d03cd29bfad6ffa3 --- /dev/null +++ b/assembly_bench/assets/parts/factory_bolt_m12_loose.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0be5fee6c1be6deaa3b65ca1811e87cb1e248f50e9dc8b463000df17910a60d8 +size 275376 diff --git a/assembly_bench/assets/parts/factory_bolt_m12_tight.usd b/assembly_bench/assets/parts/factory_bolt_m12_tight.usd new file mode 100644 index 0000000000000000000000000000000000000000..57eff5e2794897a4e4ff142c44cf2e876aea86f8 --- /dev/null +++ b/assembly_bench/assets/parts/factory_bolt_m12_tight.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f6ecb8064c36bddc6178727f48b1166e4f9ae25682ff179486ce11706efd5ea2 +size 275663 diff --git a/assembly_bench/assets/parts/factory_bolt_m16_loose.usd b/assembly_bench/assets/parts/factory_bolt_m16_loose.usd new file mode 100644 index 0000000000000000000000000000000000000000..9e76088d8f419470df1469f1019a7d88ba1d32ee --- /dev/null +++ b/assembly_bench/assets/parts/factory_bolt_m16_loose.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2ce3231074de61c243b32756dc06082ff811364f00535cd1b33ffcc47bdcb57d +size 297047 diff --git a/assembly_bench/assets/parts/factory_bolt_m16_tight.usd b/assembly_bench/assets/parts/factory_bolt_m16_tight.usd new file mode 100644 index 0000000000000000000000000000000000000000..9dcd16382ab5305946b148e7354846abdfbc792b --- /dev/null +++ b/assembly_bench/assets/parts/factory_bolt_m16_tight.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3f2203f17be7d203ba14a86449ce1519ac6b194585a881da6144f4227d9bed42 +size 297143 diff --git a/assembly_bench/assets/parts/factory_bolt_m20_loose.usd b/assembly_bench/assets/parts/factory_bolt_m20_loose.usd new file mode 100644 index 0000000000000000000000000000000000000000..8f53f4bf026c5d31dc06cfe50795a0dd51adcb61 --- /dev/null +++ b/assembly_bench/assets/parts/factory_bolt_m20_loose.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0d2220a220ae5c8aece05306f2dcb195542597f99c45eb73e09fa3dc7279e269 +size 419096 diff --git a/assembly_bench/assets/parts/factory_bolt_m20_tight.usd b/assembly_bench/assets/parts/factory_bolt_m20_tight.usd new file mode 100644 index 0000000000000000000000000000000000000000..71897c1d010d577fc768e3c5f6e88d64efc376d4 --- /dev/null +++ b/assembly_bench/assets/parts/factory_bolt_m20_tight.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:54f80109080f586accce7314bf7593f018b167a611191b6cb7a93fd7695670eb +size 419392 diff --git a/assembly_bench/assets/parts/factory_bolt_m8_loose.usd b/assembly_bench/assets/parts/factory_bolt_m8_loose.usd new file mode 100644 index 0000000000000000000000000000000000000000..f17b2707fd2ee25c828736e91ba4645850770de9 --- /dev/null +++ b/assembly_bench/assets/parts/factory_bolt_m8_loose.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a12a82aeb73cedfa86223e668db2acea64969abb05af9a81e9b7d61afd4e8ab4 +size 341225 diff --git a/assembly_bench/assets/parts/factory_bolt_m8_tight.usd b/assembly_bench/assets/parts/factory_bolt_m8_tight.usd new file mode 100644 index 0000000000000000000000000000000000000000..d1d3afbbd33eda1364d6bb8445e6cc035bd3a196 --- /dev/null +++ b/assembly_bench/assets/parts/factory_bolt_m8_tight.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c5bb650c634719dd0b596e3a5c635417b1e45bf18684c31396eb24247b65700e +size 341209 diff --git a/assembly_bench/assets/parts/factory_nut_m12_loose.usd b/assembly_bench/assets/parts/factory_nut_m12_loose.usd new file mode 100644 index 0000000000000000000000000000000000000000..eca9b1c9eda1dba5e16c9879656a4e6cc9517cee --- /dev/null +++ b/assembly_bench/assets/parts/factory_nut_m12_loose.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:3f1f3d7a57f9a981af4de03156b65236a2a34129771a3e651a180ff6cc598020 +size 158812 diff --git a/assembly_bench/assets/parts/factory_nut_m12_tight.usd b/assembly_bench/assets/parts/factory_nut_m12_tight.usd new file mode 100644 index 0000000000000000000000000000000000000000..bc9d872770bfa3aaeebfb9458478554eb528f5f7 --- /dev/null +++ b/assembly_bench/assets/parts/factory_nut_m12_tight.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b6a6b558a95ec65d2eed4d8d8051e23eeebf01d0fc0ef6cfd297b1698abc0dac +size 159164 diff --git a/assembly_bench/assets/parts/factory_nut_m16_loose.usd b/assembly_bench/assets/parts/factory_nut_m16_loose.usd new file mode 100644 index 0000000000000000000000000000000000000000..b210324a6b92371cd08687ac96b900cea9b833b7 --- /dev/null +++ b/assembly_bench/assets/parts/factory_nut_m16_loose.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1aeb2f552cbb820a5e460259ac9ac65512571e4b71c85197281cfb8c355259c2 +size 176764 diff --git a/assembly_bench/assets/parts/factory_nut_m16_tight.usd b/assembly_bench/assets/parts/factory_nut_m16_tight.usd new file mode 100644 index 0000000000000000000000000000000000000000..7f3a1f32daa89084e1f5afccfb03244a64659f5f --- /dev/null +++ b/assembly_bench/assets/parts/factory_nut_m16_tight.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5bce51b5de177629101a5ccccb36f0cc6c336a1bf6c91d116d3bf9fb843ad3f6 +size 177395 diff --git a/assembly_bench/assets/parts/factory_nut_m20_loose.usd b/assembly_bench/assets/parts/factory_nut_m20_loose.usd new file mode 100644 index 0000000000000000000000000000000000000000..9704fa2d10e0f0afdb894d8e6ca69427525085b9 --- /dev/null +++ b/assembly_bench/assets/parts/factory_nut_m20_loose.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a4a63f2acb99c161fce53ae7823bcecc5cbb8adde622fe9736d613ee092a16a4 +size 173240 diff --git a/assembly_bench/assets/parts/factory_nut_m20_tight.usd b/assembly_bench/assets/parts/factory_nut_m20_tight.usd new file mode 100644 index 0000000000000000000000000000000000000000..4c1f1bc336f1e54bb55428007f5376c26895216a --- /dev/null +++ b/assembly_bench/assets/parts/factory_nut_m20_tight.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:52188235a830e7370c0e5794752cb7110af1b04bda01555be8cfa1ea97cda39f +size 173740 diff --git a/assembly_bench/assets/parts/factory_nut_m8_loose.usd b/assembly_bench/assets/parts/factory_nut_m8_loose.usd new file mode 100644 index 0000000000000000000000000000000000000000..0bc4e14474de2150c5403cc0d4358f48951af351 --- /dev/null +++ b/assembly_bench/assets/parts/factory_nut_m8_loose.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9c14e4b13967d011bb8b134e41d6167d2352f6bfab43b17b37d41a224304f2c4 +size 148128 diff --git a/assembly_bench/assets/parts/factory_nut_m8_tight.usd b/assembly_bench/assets/parts/factory_nut_m8_tight.usd new file mode 100644 index 0000000000000000000000000000000000000000..1f77d0d88cb45c09a85e99ac88876e4e0ee78536 --- /dev/null +++ b/assembly_bench/assets/parts/factory_nut_m8_tight.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:456cf2434c57cb3312f5f5fd86572dd6e4166e9e82db9dba44a391e596ebd033 +size 148460 diff --git a/assembly_bench/assets/parts/gen_bolt_m12_loose.usd b/assembly_bench/assets/parts/gen_bolt_m12_loose.usd new file mode 100644 index 0000000000000000000000000000000000000000..1fc69cd3476123479ad6321929865586fbe4e39a --- /dev/null +++ b/assembly_bench/assets/parts/gen_bolt_m12_loose.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:69b11668d6727f14ceeeae0bb990208ec5e47f292c71981dee86bd82b680f4b4 +size 521247 diff --git a/assembly_bench/assets/parts/gen_bolt_m12_tight.usd b/assembly_bench/assets/parts/gen_bolt_m12_tight.usd new file mode 100644 index 0000000000000000000000000000000000000000..fa746897c87191efd7044eff2e056d8ccb6e3468 --- /dev/null +++ b/assembly_bench/assets/parts/gen_bolt_m12_tight.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:975a4af372760f2220b9b4a5af5084e403bf6de9ca62a21343ecf79cb1b6238a +size 521495 diff --git a/assembly_bench/assets/parts/gen_bolt_m16_loose.usd b/assembly_bench/assets/parts/gen_bolt_m16_loose.usd new file mode 100644 index 0000000000000000000000000000000000000000..82aca8287f9ce97c043d8097491a57db91374e0f --- /dev/null +++ b/assembly_bench/assets/parts/gen_bolt_m16_loose.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:726aeee3c8b8bdd10407dbdfb293c44d6cd6902e1654655521499f05ecb57e33 +size 569341 diff --git a/assembly_bench/assets/parts/gen_bolt_m16_tight.usd b/assembly_bench/assets/parts/gen_bolt_m16_tight.usd new file mode 100644 index 0000000000000000000000000000000000000000..73dd72e397cee3e105077281de5318e654f333dd --- /dev/null +++ b/assembly_bench/assets/parts/gen_bolt_m16_tight.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:35aa83f2c90df7659617205f408a3cbaadec737eec5d72750c031d8d57684c41 +size 569142 diff --git a/assembly_bench/assets/parts/gen_bolt_m20_loose.usd b/assembly_bench/assets/parts/gen_bolt_m20_loose.usd new file mode 100644 index 0000000000000000000000000000000000000000..1f327ae114acc2f37be1305c77f471280b4f3d92 --- /dev/null +++ b/assembly_bench/assets/parts/gen_bolt_m20_loose.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:457e06296e2a892df944662b0f5b3c24bd0830d584bc0e6cd4c77434bff27f18 +size 835534 diff --git a/assembly_bench/assets/parts/gen_bolt_m20_tight.usd b/assembly_bench/assets/parts/gen_bolt_m20_tight.usd new file mode 100644 index 0000000000000000000000000000000000000000..f8a78c21684bf3579b8fc35c338757a216901697 --- /dev/null +++ b/assembly_bench/assets/parts/gen_bolt_m20_tight.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:99f6ab053559d7220c60746e63b3d8646d1c291248aa3d058273280fa783cef4 +size 835278 diff --git a/assembly_bench/assets/parts/gen_bolt_m4_loose.usd b/assembly_bench/assets/parts/gen_bolt_m4_loose.usd new file mode 100644 index 0000000000000000000000000000000000000000..b8dc06e3939462681e8675139099fa70bf92e39a --- /dev/null +++ b/assembly_bench/assets/parts/gen_bolt_m4_loose.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6c565397801360d35fa1c6b3dfcd9837b930227ffad144b64af074960b38cd99 +size 1046815 diff --git a/assembly_bench/assets/parts/gen_bolt_m4_tight.usd b/assembly_bench/assets/parts/gen_bolt_m4_tight.usd new file mode 100644 index 0000000000000000000000000000000000000000..927e3eaa45a79db9b1c57f41e9c2e09601a9c424 --- /dev/null +++ b/assembly_bench/assets/parts/gen_bolt_m4_tight.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5816f4aede18772db0b1a96afc841b94332d9c3b03f2479ed693beb785272640 +size 1047872 diff --git a/assembly_bench/assets/parts/gen_bolt_m8_loose.usd b/assembly_bench/assets/parts/gen_bolt_m8_loose.usd new file mode 100644 index 0000000000000000000000000000000000000000..7682a064d862a37c01f7a3ecbb30f94482e07b58 --- /dev/null +++ b/assembly_bench/assets/parts/gen_bolt_m8_loose.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d0c42d85ecba284e365a583507b558e10989923b95b3cb0cde82bf8078e88800 +size 659384 diff --git a/assembly_bench/assets/parts/gen_bolt_m8_tight.usd b/assembly_bench/assets/parts/gen_bolt_m8_tight.usd new file mode 100644 index 0000000000000000000000000000000000000000..58fb69af24dea7498dcff8423662532af08455ad --- /dev/null +++ b/assembly_bench/assets/parts/gen_bolt_m8_tight.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7180e4226431aab9a03788c7455dea3aa0d381b7c602bae522e579b85e5eaa48 +size 660271 diff --git a/assembly_bench/assets/parts/gen_gear_large.usd b/assembly_bench/assets/parts/gen_gear_large.usd new file mode 100644 index 0000000000000000000000000000000000000000..058b0dab70c004d2d546a538ae59e7ce422231fd --- /dev/null +++ b/assembly_bench/assets/parts/gen_gear_large.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b69d5b694135c66b647b099f9afe27557c4425f9862e30640822f52fdf514a3f +size 602697 diff --git a/assembly_bench/assets/parts/gen_gear_medium.usd b/assembly_bench/assets/parts/gen_gear_medium.usd new file mode 100644 index 0000000000000000000000000000000000000000..555b5ab857e89d8e8a70ab817b9e5d4a2990063d --- /dev/null +++ b/assembly_bench/assets/parts/gen_gear_medium.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9ed91de58de2489e4d2b94aae37f3b2162b014b85695387deee16a03846d0016 +size 190331 diff --git a/assembly_bench/assets/parts/gen_gear_small.usd b/assembly_bench/assets/parts/gen_gear_small.usd new file mode 100644 index 0000000000000000000000000000000000000000..cb61f698d584495bfbfb70ec8ac7d007bdddf73d --- /dev/null +++ b/assembly_bench/assets/parts/gen_gear_small.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9c702f7b99c8909d99f6fc57aa35d6c9b4b501587b2a9b0171c6d07eb9b4ae31 +size 277850 diff --git a/assembly_bench/assets/parts/gen_nut_m12_loose.usd b/assembly_bench/assets/parts/gen_nut_m12_loose.usd new file mode 100644 index 0000000000000000000000000000000000000000..6b6749ba96283f73feea69707fedcaa981296a69 --- /dev/null +++ b/assembly_bench/assets/parts/gen_nut_m12_loose.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:620005dbd69ad8b43bc5fa1c0a999361652106cdaf0a5c22ea1fa7fdb6f8a909 +size 260871 diff --git a/assembly_bench/assets/parts/gen_nut_m12_tight.usd b/assembly_bench/assets/parts/gen_nut_m12_tight.usd new file mode 100644 index 0000000000000000000000000000000000000000..6899de513edc0660e74acc2bbcd70cf26ad278ea --- /dev/null +++ b/assembly_bench/assets/parts/gen_nut_m12_tight.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:eed7006ac13ad1922aa8980d5e349e24ca804832a40231539fbceebb7c969e3e +size 260784 diff --git a/assembly_bench/assets/parts/gen_nut_m16_loose.usd b/assembly_bench/assets/parts/gen_nut_m16_loose.usd new file mode 100644 index 0000000000000000000000000000000000000000..9a1707564a992d56adbc08fd1cdb599f1328a300 --- /dev/null +++ b/assembly_bench/assets/parts/gen_nut_m16_loose.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:4aae0444c24de0169c4cd434b021e8c52642e01fc8376985f164c278c3d80d3a +size 297845 diff --git a/assembly_bench/assets/parts/gen_nut_m16_tight.usd b/assembly_bench/assets/parts/gen_nut_m16_tight.usd new file mode 100644 index 0000000000000000000000000000000000000000..58208bf9b297372ef14475786b3c5fe33879b9c1 --- /dev/null +++ b/assembly_bench/assets/parts/gen_nut_m16_tight.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:24bdba4f898f69488df6f9d1a5149e6500f33ce70afdddf1e5dd5dc6224b48dc +size 297669 diff --git a/assembly_bench/assets/parts/gen_nut_m20_loose.usd b/assembly_bench/assets/parts/gen_nut_m20_loose.usd new file mode 100644 index 0000000000000000000000000000000000000000..473bb3c913c90893c2e6c2a004cff869453f182a --- /dev/null +++ b/assembly_bench/assets/parts/gen_nut_m20_loose.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f3cc4c2ee51e17a0ea5db830f956917a2f934de373ca44d2af75b3289bb623fe +size 289495 diff --git a/assembly_bench/assets/parts/gen_nut_m20_tight.usd b/assembly_bench/assets/parts/gen_nut_m20_tight.usd new file mode 100644 index 0000000000000000000000000000000000000000..d2b18b83a9245ece151326823d15e975f6e6af04 --- /dev/null +++ b/assembly_bench/assets/parts/gen_nut_m20_tight.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ec2cf629c58b46b7d833cf68b33485da0370ec29968b79a75740c8830ca94d93 +size 289159 diff --git a/assembly_bench/assets/parts/gen_nut_m4_loose.usd b/assembly_bench/assets/parts/gen_nut_m4_loose.usd new file mode 100644 index 0000000000000000000000000000000000000000..6b45c38827a7ac7e24b0395df0027c8a0b3f2a9e --- /dev/null +++ b/assembly_bench/assets/parts/gen_nut_m4_loose.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8a278b6cb944dc64bae7ddab79078b67f9e9264e65ee86d8f7f9cdb8769f78c6 +size 207967 diff --git a/assembly_bench/assets/parts/gen_nut_m4_tight.usd b/assembly_bench/assets/parts/gen_nut_m4_tight.usd new file mode 100644 index 0000000000000000000000000000000000000000..745095d6c239d7c0f159bb15da5983f6369756db --- /dev/null +++ b/assembly_bench/assets/parts/gen_nut_m4_tight.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2752565815ce67891a482b1caf20bb6d3e4c9e31b91627cf55313f6ad9fca5f7 +size 208031 diff --git a/assembly_bench/assets/parts/gen_nut_m8_loose.usd b/assembly_bench/assets/parts/gen_nut_m8_loose.usd new file mode 100644 index 0000000000000000000000000000000000000000..534c94029db10583998c4aa3556760abc3669eb1 --- /dev/null +++ b/assembly_bench/assets/parts/gen_nut_m8_loose.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0c4ecb413cfc71dd10342ba8be0ad70c83459e47f7064567a4de6bb2a4238e6a +size 236359 diff --git a/assembly_bench/assets/parts/gen_nut_m8_tight.usd b/assembly_bench/assets/parts/gen_nut_m8_tight.usd new file mode 100644 index 0000000000000000000000000000000000000000..e16a1e4a95654167d487217802a3135e6020da94 --- /dev/null +++ b/assembly_bench/assets/parts/gen_nut_m8_tight.usd @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b2d7d943804a2e640a24217c8e212853a84f48f51733742f44aaafcd3b7c8fa1 +size 236600 diff --git a/assembly_bench/assets/parts/gen_rect_hole_12mm.usd b/assembly_bench/assets/parts/gen_rect_hole_12mm.usd new file mode 100644 index 0000000000000000000000000000000000000000..514658964b8898e543dd30fb73e941373ccb1556 Binary files /dev/null and b/assembly_bench/assets/parts/gen_rect_hole_12mm.usd differ diff --git a/assembly_bench/assets/parts/gen_rect_hole_16mm.usd b/assembly_bench/assets/parts/gen_rect_hole_16mm.usd new file mode 100644 index 0000000000000000000000000000000000000000..6e68969ca8595de8a7d4a9ec18d70976bdc6dc2b Binary files /dev/null and b/assembly_bench/assets/parts/gen_rect_hole_16mm.usd differ diff --git a/assembly_bench/assets/parts/gen_rect_hole_4mm.usd b/assembly_bench/assets/parts/gen_rect_hole_4mm.usd new file mode 100644 index 0000000000000000000000000000000000000000..ca55e5e25baf721a42649729c58d5265aff5bcc2 Binary files /dev/null and b/assembly_bench/assets/parts/gen_rect_hole_4mm.usd differ diff --git a/assembly_bench/assets/parts/gen_rect_hole_8mm.usd b/assembly_bench/assets/parts/gen_rect_hole_8mm.usd new file mode 100644 index 0000000000000000000000000000000000000000..4007037129a3cf02219499e1ea9c7429c48af821 Binary files /dev/null and b/assembly_bench/assets/parts/gen_rect_hole_8mm.usd differ diff --git a/assembly_bench/assets/parts/gen_rect_peg_12mm_loose.usd b/assembly_bench/assets/parts/gen_rect_peg_12mm_loose.usd new file mode 100644 index 0000000000000000000000000000000000000000..70b5a69b8c14bb96cca8ceced75856dc93abbb95 Binary files /dev/null and b/assembly_bench/assets/parts/gen_rect_peg_12mm_loose.usd differ diff --git a/assembly_bench/assets/parts/gen_rect_peg_12mm_tight.usd b/assembly_bench/assets/parts/gen_rect_peg_12mm_tight.usd new file mode 100644 index 0000000000000000000000000000000000000000..b0271f98bcf5eb9f52df195796c0e97c1d6837a0 Binary files /dev/null and b/assembly_bench/assets/parts/gen_rect_peg_12mm_tight.usd differ diff --git a/assembly_bench/assets/parts/gen_rect_peg_16mm_loose.usd b/assembly_bench/assets/parts/gen_rect_peg_16mm_loose.usd new file mode 100644 index 0000000000000000000000000000000000000000..a2a4eb002ee0640b9a8e8bdb22741351767a13aa Binary files /dev/null and b/assembly_bench/assets/parts/gen_rect_peg_16mm_loose.usd differ diff --git a/assembly_bench/assets/parts/gen_rect_peg_16mm_tight.usd b/assembly_bench/assets/parts/gen_rect_peg_16mm_tight.usd new file mode 100644 index 0000000000000000000000000000000000000000..2b129bc1e251e705619ca2e0b4c26fa758e47f32 Binary files /dev/null and b/assembly_bench/assets/parts/gen_rect_peg_16mm_tight.usd differ diff --git a/assembly_bench/assets/parts/gen_rect_peg_4mm_loose.usd b/assembly_bench/assets/parts/gen_rect_peg_4mm_loose.usd new file mode 100644 index 0000000000000000000000000000000000000000..99f162b34bf05a88cb6da11d5729580becb572b1 Binary files /dev/null and b/assembly_bench/assets/parts/gen_rect_peg_4mm_loose.usd differ diff --git a/assembly_bench/assets/parts/gen_rect_peg_4mm_tight.usd b/assembly_bench/assets/parts/gen_rect_peg_4mm_tight.usd new file mode 100644 index 0000000000000000000000000000000000000000..6e3a90fe9b17bb6491b47952f35d062ffe9768bd Binary files /dev/null and b/assembly_bench/assets/parts/gen_rect_peg_4mm_tight.usd differ diff --git a/assembly_bench/assets/parts/gen_rect_peg_8mm_loose.usd b/assembly_bench/assets/parts/gen_rect_peg_8mm_loose.usd new file mode 100644 index 0000000000000000000000000000000000000000..8e474dd2ba96c41a05f719028b000c04848842e1 Binary files /dev/null and b/assembly_bench/assets/parts/gen_rect_peg_8mm_loose.usd differ diff --git a/assembly_bench/assets/parts/gen_rect_peg_8mm_tight.usd b/assembly_bench/assets/parts/gen_rect_peg_8mm_tight.usd new file mode 100644 index 0000000000000000000000000000000000000000..67788eb616e3d393132f3646dd6ba04f641de035 Binary files /dev/null and b/assembly_bench/assets/parts/gen_rect_peg_8mm_tight.usd differ diff --git a/assembly_bench/assets/parts/gen_round_hole_12mm.usd b/assembly_bench/assets/parts/gen_round_hole_12mm.usd new file mode 100644 index 0000000000000000000000000000000000000000..c6f85b2368b1fad7b736b56ec02ea05e563d8f6a Binary files /dev/null and b/assembly_bench/assets/parts/gen_round_hole_12mm.usd differ diff --git a/assembly_bench/assets/parts/gen_round_hole_16mm.usd b/assembly_bench/assets/parts/gen_round_hole_16mm.usd new file mode 100644 index 0000000000000000000000000000000000000000..1d6e458bb28e286c7a7948708f44ee4157a20a36 Binary files /dev/null and b/assembly_bench/assets/parts/gen_round_hole_16mm.usd differ diff --git a/assembly_bench/assets/parts/gen_round_hole_4mm.usd b/assembly_bench/assets/parts/gen_round_hole_4mm.usd new file mode 100644 index 0000000000000000000000000000000000000000..c6e4af9bd7a9c03514840ea63f94c491943f7253 Binary files /dev/null and b/assembly_bench/assets/parts/gen_round_hole_4mm.usd differ diff --git a/assembly_bench/assets/parts/gen_round_hole_8mm.usd b/assembly_bench/assets/parts/gen_round_hole_8mm.usd new file mode 100644 index 0000000000000000000000000000000000000000..12333f032c80309423f511685acbffde9c8ac657 Binary files /dev/null and b/assembly_bench/assets/parts/gen_round_hole_8mm.usd differ diff --git a/assembly_bench/assets/parts/gen_round_peg_12mm_loose.usd b/assembly_bench/assets/parts/gen_round_peg_12mm_loose.usd new file mode 100644 index 0000000000000000000000000000000000000000..bf4e7143239e3de4d2c39e6d3f57539131078d4d Binary files /dev/null and b/assembly_bench/assets/parts/gen_round_peg_12mm_loose.usd differ diff --git a/assembly_bench/assets/parts/gen_round_peg_12mm_tight.usd b/assembly_bench/assets/parts/gen_round_peg_12mm_tight.usd new file mode 100644 index 0000000000000000000000000000000000000000..287bfa101edc72c76b12a389811d7edec0673ea4 Binary files /dev/null and b/assembly_bench/assets/parts/gen_round_peg_12mm_tight.usd differ diff --git a/assembly_bench/assets/parts/gen_round_peg_16mm_loose.usd b/assembly_bench/assets/parts/gen_round_peg_16mm_loose.usd new file mode 100644 index 0000000000000000000000000000000000000000..2a989e0ba97154f7daf906605817ad438a2bdafc Binary files /dev/null and b/assembly_bench/assets/parts/gen_round_peg_16mm_loose.usd differ diff --git a/assembly_bench/assets/parts/gen_round_peg_16mm_tight.usd b/assembly_bench/assets/parts/gen_round_peg_16mm_tight.usd new file mode 100644 index 0000000000000000000000000000000000000000..2fdc89b097ae4c36c866365f5a47c1245baa63db Binary files /dev/null and b/assembly_bench/assets/parts/gen_round_peg_16mm_tight.usd differ diff --git a/assembly_bench/assets/parts/gen_round_peg_4mm_loose.usd b/assembly_bench/assets/parts/gen_round_peg_4mm_loose.usd new file mode 100644 index 0000000000000000000000000000000000000000..7ca8958655f9beeb4bd45dfbb1b5bf60e942c0d0 Binary files /dev/null and b/assembly_bench/assets/parts/gen_round_peg_4mm_loose.usd differ diff --git a/assembly_bench/assets/parts/gen_round_peg_4mm_tight.usd b/assembly_bench/assets/parts/gen_round_peg_4mm_tight.usd new file mode 100644 index 0000000000000000000000000000000000000000..277dd304b0a453292f10e3e059e7a23b2585d0b2 Binary files /dev/null and b/assembly_bench/assets/parts/gen_round_peg_4mm_tight.usd differ diff --git a/assembly_bench/assets/parts/gen_round_peg_8mm_loose.usd b/assembly_bench/assets/parts/gen_round_peg_8mm_loose.usd new file mode 100644 index 0000000000000000000000000000000000000000..e2ad17ce9402e84fe8d7ac4f981e1b940c5f76ad Binary files /dev/null and b/assembly_bench/assets/parts/gen_round_peg_8mm_loose.usd differ diff --git a/assembly_bench/assets/parts/gen_round_peg_8mm_tight.usd b/assembly_bench/assets/parts/gen_round_peg_8mm_tight.usd new file mode 100644 index 0000000000000000000000000000000000000000..00d549826cd47530695d067b0874e263733fd714 Binary files /dev/null and b/assembly_bench/assets/parts/gen_round_peg_8mm_tight.usd differ diff --git a/assembly_bench/assets/parts/nist_gmc_base.usd b/assembly_bench/assets/parts/nist_gmc_base.usd new file mode 100644 index 0000000000000000000000000000000000000000..5b866b571c551b6a6d995352317abb161772854c Binary files /dev/null and b/assembly_bench/assets/parts/nist_gmc_base.usd differ diff --git a/assembly_bench/environments/__init__.py b/assembly_bench/environments/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/assembly_bench/environments/__pycache__/__init__.cpython-312.pyc b/assembly_bench/environments/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..07212c54049c999244ee5a7cfeea781832252c6b Binary files /dev/null and b/assembly_bench/environments/__pycache__/__init__.cpython-312.pyc differ diff --git a/assembly_bench/environments/assembly/__init__.py b/assembly_bench/environments/assembly/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/assembly_bench/environments/assembly/__pycache__/__init__.cpython-312.pyc b/assembly_bench/environments/assembly/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ffa16c4b9c259a056f4b66ea481549219d0e96b5 Binary files /dev/null and b/assembly_bench/environments/assembly/__pycache__/__init__.cpython-312.pyc differ diff --git a/assembly_bench/environments/assembly/__pycache__/assembly.cpython-312.pyc b/assembly_bench/environments/assembly/__pycache__/assembly.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7a91be31728dc203709490fd0beb2c584d0cde9e Binary files /dev/null and b/assembly_bench/environments/assembly/__pycache__/assembly.cpython-312.pyc differ diff --git a/assembly_bench/environments/assembly/__pycache__/cameras.cpython-312.pyc b/assembly_bench/environments/assembly/__pycache__/cameras.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7ca6f71f86b161addd71eb907b4d83f56aa35f96 Binary files /dev/null and b/assembly_bench/environments/assembly/__pycache__/cameras.cpython-312.pyc differ diff --git a/assembly_bench/environments/assembly/__pycache__/embodiments.cpython-312.pyc b/assembly_bench/environments/assembly/__pycache__/embodiments.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7f1a7e420bb717aa66c607a3a9c2cc147b84e720 Binary files /dev/null and b/assembly_bench/environments/assembly/__pycache__/embodiments.cpython-312.pyc differ diff --git a/assembly_bench/environments/assembly/__pycache__/observations.cpython-312.pyc b/assembly_bench/environments/assembly/__pycache__/observations.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0a33bf81273598077ae10301ee5d49514c283581 Binary files /dev/null and b/assembly_bench/environments/assembly/__pycache__/observations.cpython-312.pyc differ diff --git a/assembly_bench/environments/assembly/__pycache__/scene.cpython-312.pyc b/assembly_bench/environments/assembly/__pycache__/scene.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..246f45c8d7dcd2be7d507dd76f649378db04e49e Binary files /dev/null and b/assembly_bench/environments/assembly/__pycache__/scene.cpython-312.pyc differ diff --git a/assembly_bench/environments/assembly/__pycache__/tasks.cpython-312.pyc b/assembly_bench/environments/assembly/__pycache__/tasks.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ca324da3de37cd6126f9f77311a325eb699569f5 Binary files /dev/null and b/assembly_bench/environments/assembly/__pycache__/tasks.cpython-312.pyc differ diff --git a/assembly_bench/environments/assembly/__pycache__/variants.cpython-312.pyc b/assembly_bench/environments/assembly/__pycache__/variants.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bbca2e6a865f1fdbb144d8ea7388f587a7a0dfcf Binary files /dev/null and b/assembly_bench/environments/assembly/__pycache__/variants.cpython-312.pyc differ diff --git a/assembly_bench/environments/assembly/assembly.py b/assembly_bench/environments/assembly/assembly.py new file mode 100644 index 0000000000000000000000000000000000000000..27707fba0f20f0a7a323a3fe49a1077be86773f7 --- /dev/null +++ b/assembly_bench/environments/assembly/assembly.py @@ -0,0 +1,267 @@ +"""NIST assembly benchmark environment for Isaac Lab Arena (externally defined). + +A Franka faces the NIST-taskboard workspace on the ``table`` background, with +one of 27 task variants (peg insert / gear mesh / nut thread) +selected via ``--task``. Run with, e.g.:: + + python isaaclab_arena/evaluation/policy_runner.py \\ + --policy_type zero_action --num_episodes 1 \\ + --external_environment_class_path \\ + assembly_bench.environments.assembly.assembly:AssemblyBenchEnvironment \\ + assembly_bench --task peg_round_M1_loose +""" + +import argparse + +from isaaclab_arena_environments.example_environment_base import ExampleEnvironmentBase + +from assembly_bench.environments.assembly.variants import TABLE_TOP_Z, VARIANTS + + +def assembly_bench_env_cfg_callback(env_cfg): + """Arena's assembly sim settings, at the benchmark's 15 Hz control rate. + + Arena's callback sets 60 Hz physics with decimation 2 (30 Hz control); + pi0.5-DROID runs at 15 Hz, so decimate by 4 and render once per policy + step (camera observations and recorded videos are 15 fps). + """ + from isaaclab.sim import RenderCfg + from isaaclab_arena_environments import mdp + + env_cfg = mdp.assembly_env_cfg_callback(env_cfg) + env_cfg.decimation = 4 + env_cfg.sim.render_interval = env_cfg.decimation + # Grasp reliability: "max" friction combine takes the higher of the two + # contacting materials' coefficients, so the grippy part (held_friction up to + # 1.0) governs the pad contact regardless of the pad material -- the Isaac Lab + # consensus fix for objects slipping out of a grasp (default "average" dilutes + # a high part friction against a low pad friction). + env_cfg.sim.physics_material.friction_combine_mode = "max" + # Contact-velocity iterations: Arena ships max_velocity_iteration_count=1, + # which under-resolves contact velocities -- a hard ram into the table then + # diverges (joint vels blow to 1e30+). Raise it so TGS actually damps the + # contact velocity and ramming stays bounded. + env_cfg.sim.physics.max_velocity_iteration_count = 4 + # DLAA (native-res temporal AA): its history is also the denoiser -- FXAA + # A/B measured temporal grain ~3.0 vs DLAA's ~0.15 (unusable boil), and + # per-frame knobs (spp, DL denoiser) proved inert at runtime. The per-wave + # teleport ghosts come from RTX geometry streaming holding stale transforms + # (kit warning: readTransformsFromFabricInRenderDelegate + geometry + # streaming "dynamic objects not streaming correctly"), so ghosting is + # fixed at the source: run_expert disables geometry streaming via carb + # right after app launch, before the scene loads. + env_cfg.sim.render = RenderCfg(antialiasing_mode="DLAA") + return env_cfg + + +class AssemblyBenchEnvironment(ExampleEnvironmentBase): + + name: str = "assembly_bench" + + def get_env(self, args_cli: argparse.Namespace): + import os + + import isaaclab.sim as sim_utils + import torch + + from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment + from isaaclab_arena.scene.scene import Scene + from isaaclab_arena.utils.pose import Pose + from isaaclab_arena_environments import mdp + + # Importing scene registers the asm_* assets with the AssetRegistry; + # importing embodiments registers droid_abs_joint_pos_softmimic. + import assembly_bench.environments.assembly.embodiments # noqa: F401 + import assembly_bench.environments.assembly.scene # noqa: F401 + from assembly_bench.environments.assembly.cameras import make_assembly_camera_cfg + from assembly_bench.environments.assembly.tasks import NISTAssemblyTask + + variant = VARIANTS[args_cli.task] + + # Background (the SeattleLabTable, as in the source benchmark), light and + # embodiment from the registry. An HDR-registry HDRI becomes the dome + # light's environment texture ("none" keeps a plain dome light). + # Table pose as in the Isaac Lab Factory env (90-degree yaw; this + # IsaacLab stack is xyzw end to end, so the tuple reads as authored). + background = self.asset_registry.get_asset_by_name("table")() + background.set_initial_pose(Pose(position_xyz=(0.55, 0.0, 0.0), rotation_xyzw=(0.0, 0.0, 0.70711, 0.70711))) + light = self.asset_registry.get_asset_by_name("light")( + spawner_cfg=sim_utils.DomeLightCfg(intensity=args_cli.light_intensity), + ) + if args_cli.hdr != "none": + light.add_hdr(self.hdr_registry.get_hdr_by_name(args_cli.hdr)()) + embodiment = self.asset_registry.get_asset_by_name(args_cli.embodiment)( + enable_cameras=args_cli.enable_cameras, + ) + # The assembly benchmark's DROID contact-stability tuning (softened + # Robotiq mimic overlay + solver/PD) lives in the registered + # `droid_abs_joint_pos_softmimic` embodiment (see embodiments.py), the + # default below -- no imperative USD authoring in the build path. The + # calibrated DROID asset (wrist camera mounted frame-for-frame against + # the source demos) is preserved via a referencing overlay. + if args_cli.embodiment == "droid_differential_ik": + # Teleop embodiment fixes (Arena's DIK cfg is untested for DROID): + # - body_name: upstream says "panda_link0" (the ARM BASE -- fixed-base + # jacobian indexing then wraps to the last body; the arm can't servo). + # Control the Robotiq "base_link" instead: same frame the eef_pos/ + # eef_quat observations and the scripted expert's IK servo use. + # - gripper: Se3Keyboard emits +1 open / -1 close; the ZeroToOne term + # (>0.5 = close) inverts that. Use the stock binary term (<0 = close), + # which matches the keyboard. NOTE: raw DIK gripper actions are thus + # +/-1, not the canonical 0/1 -- the canonical label stream is the + # recorded `abs_joint_action` (see recorders.py), which is derived + # from processed joint targets and convention-independent. + from isaaclab.envs.mdp.actions.actions_cfg import BinaryJointPositionActionCfg + + embodiment.action_config.arm_action.body_name = "base_link" + embodiment.action_config.gripper_action = BinaryJointPositionActionCfg( + asset_name="robot", + joint_names=["finger_joint"], + open_command_expr={"finger_joint": 0.0}, + close_command_expr={"finger_joint": torch.pi / 4}, + ) + if "franka" in args_cli.embodiment: + # The Factory-tuned high-PD arm the Arena assembly examples use. + embodiment.scene_config.robot = mdp.FRANKA_PANDA_ASSEMBLY_HIGH_PD_CFG.replace( + prim_path="{ENV_REGEX_NS}/Robot" + ) + # Cameras: the embodiment's calibrated wrist camera + one frontal + # exterior view, DROID-native 1280x720 (policy adapters resize). + embodiment.camera_config = make_assembly_camera_cfg(embodiment) + + # Parts at explicit poses (the reset event re-jitters them per episode). + # The peg's presentation stand is a second bore at the held position. + held = self.asset_registry.get_asset_by_name(variant.held)( + instance_name="held_part", initial_pose=Pose(position_xyz=variant.held_pos)) + fixed = self.asset_registry.get_asset_by_name(variant.fixed)( + instance_name="fixed_part", initial_pose=Pose(position_xyz=variant.fixed_pos)) + stand = None + if variant.stand is not None: + # "peg_stand", not "stand": the DROID embodiment's robot stand + # already occupies the "stand" scene entity name. + stand = self.asset_registry.get_asset_by_name(variant.stand)( + instance_name="peg_stand", + initial_pose=Pose(position_xyz=(variant.held_pos[0], variant.held_pos[1], TABLE_TOP_Z))) + extras = [ + self.asset_registry.get_asset_by_name(asset_name)( + instance_name=f"extra_{i}", + initial_pose=Pose(position_xyz=tuple(p + o for p, o in zip(variant.fixed_pos, off)))) + for i, (asset_name, off) in enumerate(variant.extras) + ] + + parts = [a for a in (held, fixed, stand, *extras) if a is not None] + scene = Scene(assets=[background, light, *parts]) + + # Privileged part poses on the wire (policy/held_part_pose, policy/ + # fixed_part_pose) for the PA-RL critic; policy adapters ignore them. + from isaaclab.managers import ObservationTermCfg as ObsTerm + from isaaclab.managers import SceneEntityCfg + + from assembly_bench.environments.assembly.observations import asset_root_pose + + embodiment.observation_config.policy.held_part_pose = ObsTerm( + func=asset_root_pose, params={"asset_cfg": SceneEntityCfg("held_part")}) + embodiment.observation_config.policy.fixed_part_pose = ObsTerm( + func=asset_root_pose, params={"asset_cfg": SceneEntityCfg("fixed_part")}) + + task = NISTAssemblyTask(variant=variant, held=held, fixed=fixed, stand=stand, extras=extras, + reward_mode=getattr(args_cli, "reward", None)) + + # Optional teleoperation (demo collection). Only the differential-IK + # embodiment can consume SE(3) devices; a keyboard cannot emit absolute + # joint targets, so fail loudly instead of moving nothing. + teleop_device = None + teleop_device_name = getattr(args_cli, "teleop_device", None) + if teleop_device_name is not None: + if args_cli.embodiment != "droid_differential_ik": + raise ValueError( + f"--teleop_device {teleop_device_name} requires --embodiment droid_differential_ik " + f"(got {args_cli.embodiment}): SE(3) teleop devices cannot drive absolute joint-position " + "actions." + ) + # Arena's registry builds the device with a low default sensitivity + # (0.05) and ignores teleop.py's --sensitivity flag. Raise it here so + # each keypress moves the arm more per step; override per session with + # TELEOP_POS_SENS / TELEOP_ROT_SENS. Keep rotation a touch lower than + # translation -- roll/pitch/yaw at high gain make fine alignment jumpy. + pos_sens = float(os.environ.get("TELEOP_POS_SENS", "0.1")) + rot_sens = float(os.environ.get("TELEOP_ROT_SENS", "0.08")) + # The streamed viewport looks at the workspace from the operator's + # side, mirrored vs. the robot base frame, so W/S, A/D and Q/E all + # read backwards. All three translation axes scale by pos_sensitivity, + # so negating it flips them together (rotation unaffected). Set + # TELEOP_INVERT_XLATE=0 to restore the raw base-frame directions. + if os.environ.get("TELEOP_INVERT_XLATE", "1") == "1": + pos_sens = -pos_sens + teleop_device = self.device_registry.get_device_by_name(teleop_device_name)( + pos_sensitivity=pos_sens, rot_sensitivity=rot_sens + ) + + return IsaacLabArenaEnvironment( + name=self.name, + embodiment=embodiment, + scene=scene, + task=task, + teleop_device=teleop_device, + env_cfg_callback=assembly_bench_env_cfg_callback, + ) + + @staticmethod + def add_cli_args(parser: argparse.ArgumentParser) -> None: + parser.add_argument("--task", type=str, default="peg_round_M1_loose", choices=sorted(VARIANTS)) + # DROID platform (Franka + Robotiq 2F-85) with absolute joint-position + # actions -- the source benchmark's droid_jointpos / pi0.5-DROID contract. + # The _softmimic variant adds the benchmark's contact-stability tuning + # (see embodiments.py); plain droid_abs_joint_pos is the untuned stock. + parser.add_argument("--embodiment", type=str, default="droid_abs_joint_pos_softmimic") + parser.add_argument("--hdr", type=str, default="asm_machine_shop", + help='HDR name from the registry (e.g. "asm_machine_shop", ' + '"carpentry_shop_robolab"), or "none"') + parser.add_argument("--light_intensity", type=float, default=1500.0) + # "staged" enables dense idempotent RL shaping (rewards.py); default + # "none" keeps the env reward-free for eval. + parser.add_argument("--reward", type=str, default="none", choices=["none", "staged"]) + # Teleop demo collection (Arena teleop.py / record_demos.py read this). + # Requires --embodiment droid_differential_ik; default None keeps eval/ + # training paths teleop-free (no retargeter exists for abs joint pos). + parser.add_argument("--teleop_device", type=str, default=None, choices=["keyboard", "spacemouse"], + help="SE(3) teleop device for demo collection " + "(requires --embodiment droid_differential_ik)") + + +def make_assembly_env( + task: str = "peg_round_M1_loose", + num_envs: int = 1, + embodiment: str = "droid_abs_joint_pos_softmimic", + hdr: str = "asm_machine_shop", + light_intensity: float = 1500.0, + reward: str = "none", +): + """Build the assembly Arena gym env for one variant (the Isaac app must be up). + + Shared factory for the HUD server (``env.py``) / ``train.rl`` collect path. + ``reward="staged"`` turns on dense idempotent shaping; "none" (default) keeps + the env reward-free for eval. ``num_envs`` is the vectorization width. + """ + import carb + from isaaclab_arena.cli.isaaclab_arena_cli import ( + arena_env_builder_cfg_from_argparse, + get_isaaclab_arena_cli_parser, + ) + from isaaclab_arena.environments.arena_env_builder import ArenaEnvBuilder + + # Rebuild guard: env construction calls rep.set_global_seed(cfg.seed), which + # touches the replicator graph; pre-sync the carb setting to the Arena default + # so a mid-run rebuild short-circuits instead of erroring. + carb.settings.get_settings().set_int("/omni/replicator/globalSeed", 42) + + parser = get_isaaclab_arena_cli_parser() + AssemblyBenchEnvironment.add_cli_args(parser) + args, _ = parser.parse_known_args([]) + args.task, args.embodiment, args.hdr = task, embodiment, hdr + args.light_intensity, args.enable_cameras = light_intensity, True + args.num_envs = num_envs + args.reward = reward + arena_env = AssemblyBenchEnvironment().get_env(args) + builder_cfg = arena_env_builder_cfg_from_argparse(args) + return ArenaEnvBuilder(arena_env, builder_cfg).make_registered(render_mode="rgb_array") diff --git a/assembly_bench/environments/assembly/cameras.py b/assembly_bench/environments/assembly/cameras.py new file mode 100644 index 0000000000000000000000000000000000000000..7a040d91af9bb78096be328e4ecd7bf39b0c1ec3 --- /dev/null +++ b/assembly_bench/environments/assembly/cameras.py @@ -0,0 +1,95 @@ +"""Benchmark cameras: the embodiment's own wrist camera + one frontal exterior. + +The DROID embodiment already ships the right cameras (calibrated Robotiq wrist +mount, 16:9 DROID intrinsics), so this module keeps their pose/intrinsics and +only swaps the two over-shoulder exterior views for a single frontal one, and +drops both cameras' render resolution to 640x360 (see RENDER_W/H). Frames stream +at 16:9 -- model-input sizing (openpi's ``resize_with_pad`` to 224x224) is the +policy adapter's job, exactly as in the chess bench's pi0.5 eval. +""" + +import numpy as np + +import isaaclab.sim as sim_utils +from isaaclab.sensors import CameraCfg + +from isaaclab_arena.utils.cameras import ArenaCameraCfg +from isaaclab_arena.utils.configclass import make_configclass + +# DROID intrinsics (2.8 mm focal, 5.376 x 3.024 mm aperture, ~88 x 57 deg FOV), +# matching Arena's DroidCameraCfg. RENDER res is 640x360 (half the DROID-native +# 1280x720, same 16:9 FOV/intrinsics -- only pixel count changes): we store +# 320x180 (DROID-RLDS) so rendering 1280x720 was 16x wasteful and the cameras +# are the per-step cost. Recorder area-downscales 640x360 -> 320x180 (clean 2x). +RENDER_W, RENDER_H = 640, 360 +_SPAWN = dict(focal_length=2.8, focus_distance=28.0, + horizontal_aperture=5.376, vertical_aperture=3.024, + clipping_range=(0.01, 6.0)) + +# Frontal exterior view (LIBERO agentview convention): close in front of the +# workspace on the robot midline, low enough that the parts read as 3D shapes, +# with the arm entering from the far side. (Framing picked from a rendered +# candidate sweep; see scripts/preview_assembly.py to re-render.) +FRONT_CAM_EYE = (0.55, 0.0, 0.30) +FRONT_CAM_TARGET = (0.37, 0.0, 0.03) + + +def lookat_opengl_quat_xyzw(eye, target, up=(0.0, 0.0, 1.0)): + """Quaternion (x,y,z,w) aiming an OpenGL/USD-convention camera (looks down + local -Z, +Y up) from ``eye`` to ``target``. Baked into the OffsetCfg: it is + per-env-relative, so one offset works for all envs.""" + eye, target, up = (np.asarray(v, dtype=float) for v in (eye, target, up)) + f = target - eye + f /= np.linalg.norm(f) # forward (view dir); camera -Z + r = np.cross(f, up) + r /= np.linalg.norm(r) # right; camera +X + u = np.cross(r, f) # true up; camera +Y + R = np.column_stack([r, u, -f]) # camera->world basis + t = np.trace(R) + if t > 0: + s = np.sqrt(t + 1.0) * 2 + w, x, y, z = 0.25 * s, (R[2, 1] - R[1, 2]) / s, (R[0, 2] - R[2, 0]) / s, (R[1, 0] - R[0, 1]) / s + else: + i = int(np.argmax(np.diag(R))) + j, k = (i + 1) % 3, (i + 2) % 3 + s = np.sqrt(1.0 + R[i, i] - R[j, j] - R[k, k]) * 2 + q = [0.0, 0.0, 0.0] + w = (R[k, j] - R[j, k]) / s + q[i] = 0.25 * s + q[j] = (R[j, i] + R[i, j]) / s + q[k] = (R[k, i] + R[i, k]) / s + x, y, z = q + return (float(x), float(y), float(z), float(w)) + + +def front_camera(eye=FRONT_CAM_EYE, target=FRONT_CAM_TARGET, name="front_cam") -> CameraCfg: + """The fixed frontal exterior camera (DROID-native 16:9).""" + return CameraCfg( + prim_path="{ENV_REGEX_NS}/" + name, + offset=CameraCfg.OffsetCfg( + pos=eye, rot=lookat_opengl_quat_xyzw(eye, target), convention="opengl"), + update_period=0.0, + height=RENDER_H, + width=RENDER_W, + data_types=["rgb"], + spawn=sim_utils.PinholeCameraCfg(**_SPAWN), + ) + + +def make_assembly_camera_cfg(embodiment): + """The embodiment's camera config with its exterior views replaced by + ``front_cam``. Wrist cameras keep the calibrated Robotiq mount pose/intrinsics + (verified frame-for-frame against the source benchmark's demos) but their + render resolution is dropped to match the front cam (RENDER_W/H).""" + fields = [] + for name in getattr(embodiment.camera_config, "__dataclass_fields__", {}): + cam = getattr(embodiment.camera_config, name) + if isinstance(cam, CameraCfg) and "wrist" in name: + # Match the front cam's render res (same FOV, fewer pixels): the + # calibrated mount/intrinsics are unchanged, only resolution drops. + cam.height, cam.width = RENDER_H, RENDER_W + fields.append((name, CameraCfg, cam)) + fields.append(("front_cam", CameraCfg, front_camera())) + # Arena embodiments require camera_config to be an ArenaCameraCfg subclass + # (get_cfg() returns tiled/untiled). make_configclass defaults to no bases. + return make_configclass("AssemblyCameraCfg", fields, bases=(ArenaCameraCfg,))() diff --git a/assembly_bench/environments/assembly/embodiments.py b/assembly_bench/environments/assembly/embodiments.py new file mode 100644 index 0000000000000000000000000000000000000000..5f36988b0e65af45365f6bd25a9e88de31934a33 --- /dev/null +++ b/assembly_bench/environments/assembly/embodiments.py @@ -0,0 +1,158 @@ +"""Assembly-tuned DROID embodiment for Isaac Lab Arena. + +Registers ``droid_abs_joint_pos_softmimic``: the stock Arena DROID absolute- +joint-position embodiment plus the assembly benchmark's contact-stability +tuning, authored ONCE at construction instead of mutated imperatively inside +``AssemblyBenchEnvironment.get_env`` on every env build. + +Why the overlay exists +---------------------- +The DROID hand is a Robotiq 2F-85 whose two fingers are mechanically linked; +PhysX models that linkage as a "mimic joint" -- a penalty spring with a +``naturalFrequency`` (stiffness) and ``dampingRatio``. The stock asset authors +it near-rigid and undamped (naturalFrequency 1e6, dampingRatio ~0). At the +benchmark's 60 Hz step that spring is numerically unstable (omega*dt >> 1), so a +hard ram self-amplifies and the gripper explodes (velocities -> NaN). + +The calibrated DROID USD must not be edited (its wrist camera is mounted +frame-for-frame against the source demos), so we author a thin USD *overlay* +that references the calibrated asset and overrides only the five mimic springs +to a stiff, well-damped coupling (naturalFrequency 1000, dampingRatio 0.7). +PhysX parses joint params at spawn, so this has to live on the joint prim, not +an actuator. + +Why a registered embodiment (not imperative mutation) +----------------------------------------------------- +The overlay is pure static data (it depends only on the source USD path), so it +is authored once per process and reused. ``Usd.Stage.CreateNew`` raises if the +layer identifier is already registered -- which is exactly what broke the served +flow, where the env is built twice (serve-time build + first ``reset`` rebuild). +Binding the tuning to a registered embodiment subclass makes the build path +declarative and rebuild-safe, per Arena's external-package pattern. +""" + +from __future__ import annotations + +import os +import tempfile +from typing import Any + +from isaaclab.managers import ObservationTermCfg as ObsTerm, SceneEntityCfg + +from isaaclab_arena.assets.register import register_asset +from isaaclab_arena.embodiments.droid.droid import DroidAbsoluteJointPositionEmbodiment + +from assembly_bench.environments.assembly.observations import arm_joint_vel, part_pose + +# Robotiq 2F-85 mimic-joint prims (relative to the flattened DROID asset) and +# the PhysX penalty-spring axis each one couples. +_MIMIC_JOINTS = ( + ("right_outer_knuckle_joint", "rotZ"), + ("right_inner_finger_joint", "rotX"), + ("right_inner_finger_knuckle_joint", "rotX"), + ("left_inner_finger_knuckle_joint", "rotX"), + ("left_inner_finger_joint", "rotX"), +) +_MIMIC_NATURAL_FREQUENCY = 1000.0 +_MIMIC_DAMPING_RATIO = 0.7 + +# Authored overlays, keyed by source USD path: the overlay is static, so build +# it once per process and reuse it across every env (re)build. A second +# Usd.Stage.CreateNew on the same layer identifier would raise. +_OVERLAY_CACHE: dict[str, str] = {} + + +def _softened_mimic_overlay(source_usd_path: str) -> str: + """Path to a USD overlay referencing *source_usd_path* with the Robotiq + mimic springs softened. Authored once per process, then reused.""" + cached = _OVERLAY_CACHE.get(source_usd_path) + if cached is not None and os.path.exists(cached): + return cached + + from pxr import Sdf, Usd + + overlay = os.path.join(tempfile.gettempdir(), "asm_droid_softened_mimic.usd") + # USD caches layers by identifier; reuse an already-open layer (authored by + # an earlier build in this process) instead of re-creating it. + if Sdf.Layer.Find(overlay) is None: + if os.path.exists(overlay): + os.remove(overlay) + stage = Usd.Stage.CreateNew(overlay) + root = stage.OverridePrim("/panda") + root.GetReferences().AddReference(source_usd_path) + stage.SetDefaultPrim(root.GetPrim()) + for joint, axis in _MIMIC_JOINTS: + prim = stage.OverridePrim(f"/panda/Gripper/Robotiq_2F_85/Joints/{joint}") + prim.CreateAttribute( + f"physxMimicJoint:{axis}:naturalFrequency", Sdf.ValueTypeNames.Float + ).Set(_MIMIC_NATURAL_FREQUENCY) + prim.CreateAttribute( + f"physxMimicJoint:{axis}:dampingRatio", Sdf.ValueTypeNames.Float + ).Set(_MIMIC_DAMPING_RATIO) + stage.GetRootLayer().Save() + + _OVERLAY_CACHE[source_usd_path] = overlay + return overlay + + +def apply_assembly_droid_tuning(embodiment: Any) -> None: + """Apply the assembly benchmark's contact-stability tuning to a DROID + embodiment in place: softened-mimic spawn overlay, solver-iteration + headroom, a compliant arm, and joint velocity in the policy observation. + + Kept as a module function so any DROID variant (e.g. a teleop embodiment) + can opt into the same tuning without duplicating it. + """ + robot = embodiment.scene_config.robot + + # Point the spawn at the softened-mimic overlay (keeps the calibrated asset, + # overrides only the mimic springs). The USD's native firm finger drive + # stays as-is: it stops the pad firmly at the part surface, and the mimic + # coupling -- not the finger drive -- was the divergence root cause. + robot.spawn.usd_path = _softened_mimic_overlay(robot.spawn.usd_path) + + # Ramming the table should stall, not explode: keep depenetration bounded, + # add articulation velocity iterations (0 -> 8) so the solver damps contact + # velocity, and raise position iterations (64 -> 192) to give the stiff + # mimic spring the headroom to hold without diverging. + robot.spawn.rigid_props.max_depenetration_velocity = 5.0 + robot.spawn.articulation_props.solver_velocity_iteration_count = 8 + robot.spawn.articulation_props.solver_position_iteration_count = 192 + + # Compliant arm: soften PD (400/80 -> 150/40) so the arm yields under + # contact instead of shoving the now-rigid gripper into the part ("arm + # gives, gripper stays"). Restore the real panda joint-speed limits + # (velocity_limit is a no-op on implicit actuators) and add base-panda + # armature to damp high-PD jitter. + for name, vlim in (("panda_shoulder", 2.175), ("panda_forearm", 2.61)): + robot.actuators[name].velocity_limit_sim = vlim + robot.actuators[name].armature = 1e-3 + robot.actuators[name].stiffness = 150.0 + robot.actuators[name].damping = 40.0 + + # Joint velocities alongside joint positions in the policy obs (recorded to + # HDF5 via the flat policy-obs recorder term). + embodiment.observation_config.policy.joint_vel = ObsTerm(func=arm_joint_vel) + # Privileged part poses for the PA-RL critic (not fed to the VLA). + embodiment.observation_config.policy.held_part_pose = ObsTerm( + func=part_pose, params={"asset_cfg": SceneEntityCfg("held_part")}) + embodiment.observation_config.policy.fixed_part_pose = ObsTerm( + func=part_pose, params={"asset_cfg": SceneEntityCfg("fixed_part")}) + + +@register_asset +class DroidAbsoluteJointPositionSoftMimicEmbodiment(DroidAbsoluteJointPositionEmbodiment): + """DROID absolute-joint-position embodiment with the assembly benchmark's + contact-stability tuning baked in (softened Robotiq mimic + solver/PD). + + This is the benchmark's default embodiment; the pi0.5-DROID contract (8-D + action: 7 joint targets + binary gripper) is unchanged from the stock + absolute-joint-position embodiment it subclasses. + """ + + name = "droid_abs_joint_pos_softmimic" + tags = ["embodiment", "assembly"] + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + apply_assembly_droid_tuning(self) diff --git a/assembly_bench/environments/assembly/observations.py b/assembly_bench/environments/assembly/observations.py new file mode 100644 index 0000000000000000000000000000000000000000..1dfdc8fd8189d3db13ebf628879041b03e68abbf --- /dev/null +++ b/assembly_bench/environments/assembly/observations.py @@ -0,0 +1,37 @@ +"""Extra policy observations for the assembly benchmark. + +Mirrors the DROID embodiment's observation helpers (isaaclab_arena +embodiments/droid/observations.py): env-local reads off the articulation, +Panda arm joints only, in joint-buffer order (asserted panda_joint1..7 +first on this asset by the expert servo). +""" + +import torch + +import warp as wp +from isaaclab.envs import ManagerBasedRLEnv +from isaaclab.managers import SceneEntityCfg + +_PANDA_JOINTS = [f"panda_joint{i}" for i in range(1, 8)] + + +def arm_joint_vel(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor: + """Angular velocities (rad/s) of the 7 Panda arm joints.""" + robot = env.scene[asset_cfg.name] + joint_indices = [i for i, name in enumerate(robot.data.joint_names) if name in _PANDA_JOINTS] + return wp.to_torch(robot.data.joint_vel)[:, joint_indices] + + +def part_pose(env: ManagerBasedRLEnv, asset_cfg: SceneEntityCfg) -> torch.Tensor: + """Env-local root pose of a scene asset: position(3) + quat xyzw(4). + + Privileged observation for the PA-RL critic (same convention as the + reward-labeled demo datasets: env-origin-relative position, xyzw quat). + """ + asset = env.scene[asset_cfg.name] + pos = wp.to_torch(asset.data.root_pos_w) - env.scene.env_origins + return torch.cat([pos, wp.to_torch(asset.data.root_quat_w)], dim=-1) + + +# Alias used by assembly.py privileged ObsTerm wiring. +asset_root_pose = part_pose diff --git a/assembly_bench/environments/assembly/recorders.py b/assembly_bench/environments/assembly/recorders.py new file mode 100644 index 0000000000000000000000000000000000000000..545aef30629c0a5fe6819f4fdfd50cc5d76e3107 --- /dev/null +++ b/assembly_bench/environments/assembly/recorders.py @@ -0,0 +1,46 @@ +"""Recorder terms for teleop demo collection. + +The teleop env acts in differential-IK space (6 relative EE values + gripper), +but the benchmark's canonical policy contract (pi0.5-DROID, contract.json) is +8-D absolute joint targets. ``PostStepAbsJointTargetRecorder`` records that +canonical label stream alongside the raw DIK actions: post-step, the arm's +``joint_pos_target`` buffer holds exactly the absolute target the IK solver +commanded this step (what set_joint_position_target received), and the +finger_joint target normalizes back to the 0/1 gripper convention -- so the +label is independent of the teleop device's raw gripper sign convention. + +Raw ``actions`` stay in the HDF5 untouched: replay, annotation, and Mimic +generation step the DIK env with them. ``abs_joint_action`` is the training +export for the droid_abs_joint_pos contract. +""" + +import torch + +import warp as wp +from isaaclab.managers import RecorderTerm, RecorderTermCfg +from isaaclab.utils import configclass + +_PANDA_JOINTS = [f"panda_joint{i}" for i in range(1, 8)] +_GRIPPER_CLOSE_RAD = torch.pi / 4 # finger_joint close target (droid embodiments) + + +class PostStepAbsJointTargetRecorder(RecorderTerm): + """Records the canonical 8-D [7 abs joint targets, gripper 0..1] per step.""" + + def record_post_step(self): + robot = self._env.scene["robot"] + if not hasattr(self, "_arm_idx"): + names = list(robot.data.joint_names) + self._arm_idx = [names.index(j) for j in _PANDA_JOINTS] + self._finger_idx = names.index("finger_joint") + targets = wp.to_torch(robot.data.joint_pos_target) + arm = targets[:, self._arm_idx] + grip = (targets[:, self._finger_idx : self._finger_idx + 1] / _GRIPPER_CLOSE_RAD).clamp(0.0, 1.0) + return "abs_joint_action", torch.cat([arm, grip], dim=-1) + + +@configclass +class PostStepAbsJointTargetRecorderCfg(RecorderTermCfg): + """Configuration for the absolute joint-target label recorder term.""" + + class_type: type[RecorderTerm] = PostStepAbsJointTargetRecorder diff --git a/assembly_bench/environments/assembly/rewards.py b/assembly_bench/environments/assembly/rewards.py new file mode 100644 index 0000000000000000000000000000000000000000..1777bc6104a94a1adbc3dfa1e179f6df024f7df9 --- /dev/null +++ b/assembly_bench/environments/assembly/rewards.py @@ -0,0 +1,228 @@ +"""Staged assembly reward for RL fine-tuning (opt-in; eval leaves it off). + +Dense, idempotent shaping for PA-RL / FilteredBC: once-fired milestones plus +new-best potentials (pay only progress, never a per-step "being close" rate). + +Phases: + pre-lift -- Φ_grasp: EE approaches the held part + post-lift -- Φ_xy / Φ_depth: held part approaches the seat + nut only -- after engage: once-fired first-turn + Φ_thread (new-best turns) + +Wired only when ``reward_mode="staged"``; eval stays reward-free. Seat geometry +matches ``tasks.part_seated``. State lives on ``ManagerTermBase`` (same pattern +as ``hold_success``). +""" + +import math + +import torch + +import warp as wp +from isaaclab.envs import ManagerBasedRLEnv +from isaaclab.managers import ManagerTermBase, RewardTermCfg, SceneEntityCfg +from isaaclab.utils import configclass +from isaaclab.utils.math import quat_apply + +from assembly_bench.environments.assembly.tasks import part_seated +from assembly_bench.environments.assembly.variants import TABLE_TOP_Z, AssemblyVariant + +# Once-fired milestones. +W_LIFT = 0.5 +W_ENGAGE = 0.4 +W_THREAD_START = 1.5 # nut: first meaningful turn on the bolt +W_SUCCESS = 2.0 + +# Continuous new-best potentials (total contribution capped by weight * 1.0). +W_GRASP_PC = 0.2 +W_ALIGN_PC = 0.3 +W_DEPTH_PC = 0.5 +W_THREAD_PC = 0.8 + +# Grasp approach (flange EE -> held). Falloff radius for Φ_grasp. +GRASP_Z_OFF = 0.16 +GRASP_DIST0 = 0.15 +# Seat xy approach falloff (m); depth uses variant.partial_socket_h. +ALIGN_DIST0 = 0.10 +# Nut thread: first-turn threshold + target cumulative |Δyaw| (~1.5 turns). +THREAD_START_RAD = 0.3 +THREAD_TARGET_RAD = 1.5 * 2.0 * math.pi + +MILESTONES = ("lifted", "engaged", "thread_start", "success") + + +def _align_and_gap(env, held_cfg, fixed_cfg, seat_off, held_base_z_off): + """Seat geometry: (xy distance to target, seat gap). Mirrors ``part_seated``.""" + held = env.scene[held_cfg.name] + fixed = env.scene[fixed_cfg.name] + held_pos = wp.to_torch(held.data.root_pos_w) - env.scene.env_origins + fixed_pos = wp.to_torch(fixed.data.root_pos_w) - env.scene.env_origins + fixed_quat = wp.to_torch(fixed.data.root_quat_w) + off = torch.tensor(seat_off, device=env.device).expand(env.num_envs, 3) + target = fixed_pos + quat_apply(fixed_quat, off) + xy = torch.norm(held_pos[:, :2] - target[:, :2], dim=-1) + gap = (held_pos[:, 2] + held_base_z_off) - target[:, 2] + return xy, gap + + +def _yaw_z(q: torch.Tensor) -> torch.Tensor: + """World-z yaw from xyzw quaternions [N,4].""" + x, y, z, w = q.unbind(-1) + return torch.atan2(2.0 * (w * z + x * y), 1.0 - 2.0 * (y * y + z * z)) + + +class staged_assembly_reward(ManagerTermBase): + """Phase-gated new-best potentials + once-fired milestones.""" + + def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + n, dev = env.num_envs, env.device + self._fired = {k: torch.zeros(n, dtype=torch.bool, device=dev) for k in MILESTONES} + self._best_grasp = torch.zeros(n, device=dev) + self._best_xy = torch.zeros(n, device=dev) + self._best_depth = torch.zeros(n, device=dev) + self._best_thread = torch.zeros(n, device=dev) + self._prev_yaw = torch.zeros(n, device=dev) + self._yaw_valid = torch.zeros(n, dtype=torch.bool, device=dev) + self._turns = torch.zeros(n, device=dev) + + def reset(self, env_ids=None) -> None: + idx = slice(None) if env_ids is None else env_ids + for v in self._fired.values(): + v[idx] = False + self._best_grasp[idx] = 0.0 + self._best_xy[idx] = 0.0 + self._best_depth[idx] = 0.0 + self._best_thread[idx] = 0.0 + self._prev_yaw[idx] = 0.0 + self._yaw_valid[idx] = False + self._turns[idx] = 0.0 + + def __call__( + self, + env: ManagerBasedRLEnv, + held_cfg: SceneEntityCfg, + fixed_cfg: SceneEntityCfg, + robot_cfg: SceneEntityCfg, + seat_off: tuple[float, float, float], + held_base_z_off: float, + align_tol: float, + seat_tol: float, + engage_gap: float, + lift_clear: float, + partial_socket_h: float, + stand_z: float, + ee_body: str, + family: str = "peg_insert", + ) -> torch.Tensor: + held = env.scene[held_cfg.name] + robot = env.scene[robot_cfg.name] + origins = env.scene.env_origins + held_pos = wp.to_torch(held.data.root_pos_w) - origins + held_z = held_pos[:, 2] + xy, gap = _align_and_gap(env, held_cfg, fixed_cfg, seat_off, held_base_z_off) + + # Φ_grasp: flange approaches the held part (pre-lift only). + ee_idx = robot.data.body_names.index(ee_body) + ee_pos = wp.to_torch(robot.data.body_pos_w)[:, ee_idx] - origins + grasp_d = torch.norm(held_pos - (ee_pos - torch.tensor( + [0.0, 0.0, GRASP_Z_OFF], device=env.device)), dim=-1) + phi_grasp = torch.clamp(1.0 - grasp_d / GRASP_DIST0, 0.0, 1.0) + + phi_xy = torch.clamp(1.0 - xy / ALIGN_DIST0, 0.0, 1.0) + phi_depth = torch.clamp(1.0 - gap / partial_socket_h, 0.0, 1.0) + + lifted = (held_z - stand_z) > lift_clear + engaged = (xy < align_tol) & (gap < engage_gap) + success = part_seated(env, held_cfg, fixed_cfg, seat_off, held_base_z_off, + align_tol, seat_tol) + + # Nut: accumulate |Δyaw| while engaged -> Φ_thread + first-turn milestone. + is_nut = family == "nut_thread" + thread_start = torch.zeros(env.num_envs, dtype=torch.bool, device=env.device) + if is_nut: + yaw = _yaw_z(wp.to_torch(held.data.root_quat_w)) + delta = (yaw - self._prev_yaw + torch.pi) % (2.0 * torch.pi) - torch.pi + self._turns = self._turns + torch.where( + engaged & self._yaw_valid, delta.abs(), torch.zeros_like(delta)) + self._prev_yaw = yaw + self._yaw_valid = self._yaw_valid | engaged + thread_start = self._turns > THREAD_START_RAD + phi_thread = torch.clamp(self._turns / THREAD_TARGET_RAD, 0.0, 1.0) + else: + phi_thread = torch.zeros(env.num_envs, device=env.device) + + states = { + "lifted": lifted, + "engaged": engaged, + "thread_start": thread_start if is_nut else torch.zeros_like(lifted), + "success": success, + } + weights = { + "lifted": W_LIFT, + "engaged": W_ENGAGE, + "thread_start": W_THREAD_START, + "success": W_SUCCESS, + } + + r = torch.zeros(env.num_envs, device=env.device) + for k in MILESTONES: + newly = states[k] & ~self._fired[k] + r += weights[k] * newly.float() + self._fired[k] |= states[k] + + # Continuous: grasp approach until lift fires; seat approach after. + pre_lift = ~self._fired["lifted"] + dg = torch.clamp(phi_grasp - self._best_grasp, min=0.0) + r += W_GRASP_PC * dg * pre_lift.float() + self._best_grasp = torch.where(pre_lift, torch.maximum(self._best_grasp, phi_grasp), + self._best_grasp) + + post_lift = self._fired["lifted"] + dxy = torch.clamp(phi_xy - self._best_xy, min=0.0) + dd = torch.clamp(phi_depth - self._best_depth, min=0.0) + r += (W_ALIGN_PC * dxy + W_DEPTH_PC * dd) * post_lift.float() + self._best_xy = torch.where(post_lift, torch.maximum(self._best_xy, phi_xy), + self._best_xy) + self._best_depth = torch.where(post_lift, torch.maximum(self._best_depth, phi_depth), + self._best_depth) + + # Nut thread densify after engage (turns on bolt; depth already covered above). + if is_nut: + on_bolt = self._fired["engaged"] + dt = torch.clamp(phi_thread - self._best_thread, min=0.0) + r += W_THREAD_PC * dt * on_bolt.float() + self._best_thread = torch.where( + on_bolt, torch.maximum(self._best_thread, phi_thread), self._best_thread) + + # RewardManager scales by weight*dt; undo dt so payouts are absolute. + return r / env.step_dt + + +@configclass +class AssemblyRewardsCfg: + staged: RewardTermCfg = None + + +def build_rewards_cfg(variant: AssemblyVariant, held, fixed, robot_name: str = "robot", + ee_body: str = "base_link") -> AssemblyRewardsCfg: + """Staged reward wired to a variant's seat geometry. ``weight=1.0`` + (per-term weights live inside; dt is undone there).""" + return AssemblyRewardsCfg(staged=RewardTermCfg( + func=staged_assembly_reward, + weight=1.0, + params={ + "held_cfg": SceneEntityCfg(held.name), + "fixed_cfg": SceneEntityCfg(fixed.name), + "robot_cfg": SceneEntityCfg(robot_name), + "seat_off": variant.seat_off, + "held_base_z_off": variant.held_base_z_off, + "align_tol": variant.align_tol, + "seat_tol": variant.seat_tol, + "engage_gap": variant.engage_gap, + "lift_clear": variant.lift_clear, + "partial_socket_h": variant.partial_socket_h, + "stand_z": TABLE_TOP_Z, + "ee_body": ee_body, + "family": variant.family, + }, + )) diff --git a/assembly_bench/environments/assembly/scene.py b/assembly_bench/environments/assembly/scene.py new file mode 100644 index 0000000000000000000000000000000000000000..7774ee7d3d394df6f042bb1fbbc42cfb6905b64a --- /dev/null +++ b/assembly_bench/environments/assembly/scene.py @@ -0,0 +1,167 @@ +"""Arena-registered assets for the NIST assembly benchmark. + +Registers every part the variants reference, under an ``asm_`` prefix so the +names never collide with Arena's own library. Local USDs ship in +``assembly_bench/assets/parts`` (generated, physics-validated: watertight SDF +bores, calibrated frames); the M16 nut/bolt and the gear base come from the +Isaac Lab Factory asset dir, exactly as in the source benchmark. + +Asset roles mirror the source env: + held -- free dynamic part (gravity on), grasped and assembled. + fixed -- world-pinned socket / base / bolt (the curated/generated USDs carry + their own fixed joint, so spawning them as articulations pins them). + extras / board -- kinematic context and mesh partners (flanking gears, NIST board). +""" + +from pathlib import Path + +import isaaclab.sim as sim_utils +from isaaclab.assets import ArticulationCfg +from isaaclab_tasks.direct.factory.factory_tasks_cfg import ASSET_DIR + +from isaaclab_arena.assets.hdr_image_library import LibraryHDR +from isaaclab_arena.assets.object_base import ObjectType +from isaaclab_arena.assets.object_library import LibraryObject +from isaaclab_arena.assets.object_utils import RIGID_BODY_PROPS_HIGH_PRECISION +from isaaclab_arena.assets.register import register_asset, register_hdr + +ASSETS_DIR = Path(__file__).resolve().parents[2] / "assets" +PARTS_DIR = ASSETS_DIR / "parts" + + +@register_hdr +class MachineShopHDR(LibraryHDR): + """The source benchmark's machine-shop backdrop (Poly Haven CC0), shipped locally.""" + + name = "asm_machine_shop" + tags = ["indoor", "workshop"] + texture_file = str(ASSETS_DIR / "backgrounds" / "indoors" / "machine_shop_01_2k.hdr") + +# Empty-joint init state for DOF-less articulations (the {".*": 0.0} default +# fails to match when the articulation has no joints). +_EMPTY_INIT = ArticulationCfg.InitialStateCfg(joint_pos={}, joint_vel={}) + +# Contact-impulse cap for gear/nut assets: a mis-phased tooth clash or the +# nut's fine-thread SDF bore can otherwise generate an unbounded impulse -> +# NaN blow-up. The cap must still be high enough that a gripper can actually +# clamp the part: at 1e4 the finger's clamp impulse was clipped and the pad +# sank through the gear. A low cap can only make sink-through worse, so keep it +# high (1e16) -- still finite, so the NaN guard survives a runaway SDF clash, +# but with ample headroom so a firm clamp is never clipped into the part. +GEAR_NUT_IMPULSE_CAP = 1e16 + +GEAR_MASS = {"small": 0.006, "medium": 0.012, "large": 0.025} + + +def _register(name: str, usd: str, *, mass: float | None = None, impulse_cap: float = 1e32, + kinematic: bool = False) -> None: + """Register one part. Free/pinned parts spawn as DOF-less articulations + (the NIST USDs carry articulation roots); ``kinematic`` disables the + articulation root and spawns an immovable rigid body instead.""" + spawn: dict = { + "mass_props": sim_utils.MassPropertiesCfg(mass=mass) if mass is not None else None, + "collision_props": sim_utils.CollisionPropertiesCfg(contact_offset=0.005, rest_offset=0.0), + } + if kinematic: + spawn["rigid_props"] = sim_utils.RigidBodyPropertiesCfg( + kinematic_enabled=True, disable_gravity=True, max_depenetration_velocity=5.0, + solver_position_iteration_count=192, solver_velocity_iteration_count=1, + max_contact_impulse=impulse_cap) + spawn["articulation_props"] = sim_utils.ArticulationRootPropertiesCfg(articulation_enabled=False) + object_type, asset_addon = ObjectType.RIGID, {} + else: + # Free part. Standard high-precision rigid body (same treatment pegs use + # and grasp cleanly with) -- the earlier gear/nut-specific throttle + # (capped depenetration velocity, capped lin/ang velocity, extra damping) + # was fighting the grasp, so it was dropped. Retain only two guards: a + # bounded contact impulse (NaN safety on gear-tooth / nut-thread SDF + # clashes) and velocity solver iterations raised 1 -> 8 to match the robot + # so both sides resolve contact velocity (more iterations only sharpen + # contact, never worsen it). + spawn["rigid_props"] = RIGID_BODY_PROPS_HIGH_PRECISION.replace( + max_contact_impulse=impulse_cap, + solver_velocity_iteration_count=8, + ) + object_type, asset_addon = ObjectType.ARTICULATION, {"init_state": _EMPTY_INIT} + register_asset(type( + name.title().replace("_", ""), + (LibraryObject,), + { + "name": name, + "tags": ["object", "assembly"], + "usd_path": usd, + "object_type": object_type, + "spawn_cfg_addon": spawn, + "asset_cfg_addon": asset_addon, + }, + )) + + +# Peg family: loose-clearance pegs (50 mm long, ~1 mm chamfer mouth) + matching +# bores. The bore doubles as the presentation stand the peg starts standing in. +for _size in (4, 8, 12, 16): + for _stem in ("round", "rect"): + _register(f"asm_peg_{_stem}_{_size}mm_loose", + str(PARTS_DIR / f"gen_{_stem}_peg_{_size}mm_loose.usd"), mass=0.019) + _register(f"asm_hole_{_stem}_{_size}mm", + str(PARTS_DIR / f"gen_{_stem}_hole_{_size}mm.usd"), mass=0.05) + +# Gear family: re-centered gears (frame on the shaft axis) as both the free +# held part and the kinematic flanking mesh partners; base from Factory. +for _k, _m in GEAR_MASS.items(): + _register(f"asm_gear_{_k}", str(PARTS_DIR / f"gen_gear_{_k}.usd"), + mass=_m, impulse_cap=GEAR_NUT_IMPULSE_CAP) + _register(f"asm_gear_{_k}_fixed", str(PARTS_DIR / f"gen_gear_{_k}.usd"), + mass=_m, impulse_cap=GEAR_NUT_IMPULSE_CAP, kinematic=True) +_register("asm_gear_base", f"{ASSET_DIR}/factory_gear_base.usd", + mass=0.05, impulse_cap=GEAR_NUT_IMPULSE_CAP) + +# Nut family on the NIST GMC board: IsaacGymEnvs factory OBJs converted to USD +# (SDF + brass/steel MDL). Loose clearance only. Procedural gen_* pairs do not +# SDF-mate for descent. M4 omitted — pad geometry cannot grasp the 3.2 mm hex. +for _s in (8, 12, 16, 20): + _register(f"asm_nut_m{_s}_loose", str(PARTS_DIR / f"factory_nut_m{_s}_loose.usd"), + impulse_cap=GEAR_NUT_IMPULSE_CAP) + _register(f"asm_bolt_m{_s}_loose", str(PARTS_DIR / f"factory_bolt_m{_s}_loose.usd"), + impulse_cap=GEAR_NUT_IMPULSE_CAP) +_register("asm_nist_board", str(PARTS_DIR / "nist_gmc_base.usd"), + mass=1.0, impulse_cap=GEAR_NUT_IMPULSE_CAP, kinematic=True) + +# DEBUG ONLY — RoboLab apple/bowl for a policy sanity-check pick-and-place. +# Not part of the NIST assembly matrix. Apple is Objaverse (~100x oversized). +DEBUG_DIR = PARTS_DIR / "debug" + + +@register_asset +class AsmDebugApple(LibraryObject): + """DEBUG: RoboLab ``apple_01`` (scaled to ~7 cm).""" + + name = "asm_debug_apple" + tags = ["object", "debug"] + usd_path = str(DEBUG_DIR / "apple_01.usd") + object_type = ObjectType.RIGID + scale = (0.01, 0.01, 0.01) + spawn_cfg_addon = { + "rigid_props": RIGID_BODY_PROPS_HIGH_PRECISION, + "mass_props": sim_utils.MassPropertiesCfg(mass=0.15), + "collision_props": sim_utils.CollisionPropertiesCfg(contact_offset=0.005, rest_offset=0.0), + } + + +@register_asset +class AsmDebugBowl(LibraryObject): + """DEBUG: RoboLab YCB bowl (kinematic receptacle).""" + + name = "asm_debug_bowl" + tags = ["object", "debug"] + usd_path = str(DEBUG_DIR / "bowl.usd") + object_type = ObjectType.RIGID + spawn_cfg_addon = { + "rigid_props": sim_utils.RigidBodyPropertiesCfg( + kinematic_enabled=True, disable_gravity=True, + max_depenetration_velocity=5.0, + solver_position_iteration_count=16, solver_velocity_iteration_count=1, + ), + "mass_props": sim_utils.MassPropertiesCfg(mass=0.2), + "collision_props": sim_utils.CollisionPropertiesCfg(contact_offset=0.005, rest_offset=0.0), + } diff --git a/assembly_bench/environments/assembly/tasks.py b/assembly_bench/environments/assembly/tasks.py new file mode 100644 index 0000000000000000000000000000000000000000..0f3fcb8ce87396670193bb7f3955f447d61f6c01 --- /dev/null +++ b/assembly_bench/environments/assembly/tasks.py @@ -0,0 +1,347 @@ +"""NIST assembly task for Isaac Lab Arena. + +One task class covers all families (peg insert / gear mesh / nut thread); +only the seat geometry differs, carried by the variant. Success +is the source benchmark's outcome check: the held part's base point reaches +the seat target (fixed root + family offset, rotated into the fixed frame) +within the alignment/seat tolerances AND the part is at rest — debounced by +``hold_success`` (N consecutive true steps). The hard contact skill lives in +the policy, not the check: a mis-clocked peg, clashing gear, or cross-threaded +nut cannot descend, so the seat gap stays large. +""" + +import numpy as np +import torch +from dataclasses import MISSING + +import isaaclab.envs.mdp as mdp_isaac_lab +import warp as wp +from isaaclab.envs import ManagerBasedRLEnv +from isaaclab.envs.common import ViewerCfg +from isaaclab.managers import EventTermCfg, ManagerTermBase, SceneEntityCfg, TerminationTermCfg +from isaaclab.utils import configclass +from isaaclab.utils.math import quat_apply + +import isaaclab_arena_environments.mdp as mdp +from isaaclab_arena.assets.object import Object +from isaaclab_arena.assets.register import register_task +from isaaclab_arena.metrics.metric_base import MetricBase +from isaaclab_arena.metrics.object_moved import ObjectMovedRateMetric +from isaaclab_arena.metrics.success_rate import SuccessRateMetric +from isaaclab_arena.tasks.task_base import TaskBase +from isaaclab_arena.utils.cameras import get_viewer_cfg_look_at_object + +from assembly_bench.environments.assembly.variants import TABLE_TOP_Z, AssemblyVariant + +# Success must hold this many consecutive steps (debounces one-frame trues). +SUCCESS_HOLD_STEPS = 3 +# The presentation stand only presents the part upright -- near-frictionless +# so the part lifts off cleanly instead of jamming in the bore. +STAND_FRICTION = 0.01 +# Failure: the held part fell below the tabletop. +DROP_HEIGHT = TABLE_TOP_Z - 0.05 + +# --------------------------------------------------------------------------- +# Success predicate +# --------------------------------------------------------------------------- + + +def part_seated( + env: ManagerBasedRLEnv, + held_cfg: SceneEntityCfg, + fixed_cfg: SceneEntityCfg, + seat_off: tuple[float, float, float], + held_base_z_off: float, + align_tol: float, + seat_tol: float, + max_speed: float = 0.05, +) -> torch.Tensor: + """Held base at the seat target (aligned in xy, descended past the seat + depth) and the part nearly at rest. seat_off is rotated into the fixed + asset's frame so the target tracks any base yaw (rect-peg clocking).""" + held = env.scene[held_cfg.name] + fixed = env.scene[fixed_cfg.name] + held_pos = wp.to_torch(held.data.root_pos_w) - env.scene.env_origins + fixed_pos = wp.to_torch(fixed.data.root_pos_w) - env.scene.env_origins + fixed_quat = wp.to_torch(fixed.data.root_quat_w) + + off = torch.tensor(seat_off, device=env.device).expand(env.num_envs, 3) + target = fixed_pos + quat_apply(fixed_quat, off) + + aligned = torch.norm(held_pos[:, :2] - target[:, :2], dim=-1) < align_tol + seat_gap = (held_pos[:, 2] + held_base_z_off) - target[:, 2] + seated = seat_gap < seat_tol + stable = torch.norm(wp.to_torch(held.data.root_lin_vel_w), dim=-1) < max_speed + return aligned & seated & stable + + +class hold_success(ManagerTermBase): + """Success wrapper: the raw predicate ``func(env, **params)`` must hold for + ``SUCCESS_HOLD_STEPS`` consecutive steps. Class-based so it keeps a per-env + streak counter; the TerminationManager resets it for envs that reset.""" + + def __init__(self, cfg: TerminationTermCfg, env: ManagerBasedRLEnv): + super().__init__(cfg, env) + self._streak = torch.zeros(env.num_envs, dtype=torch.long, device=env.device) + + def reset(self, env_ids=None) -> None: + self._streak[env_ids if env_ids is not None else slice(None)] = 0 + + def __call__(self, env: ManagerBasedRLEnv, func, params: dict) -> torch.Tensor: + self._streak = torch.where(func(env, **params), self._streak + 1, 0) + return self._streak >= SUCCESS_HOLD_STEPS + + +# --------------------------------------------------------------------------- +# Reset randomization (the source benchmark's jitter scheme) +# --------------------------------------------------------------------------- + + +def reset_assembly_workspace( + env, + env_ids: torch.Tensor, + held_group: list[tuple[str, tuple[float, float, float]]], + fixed_group: list[tuple[str, tuple[float, float, float]]], + rand_held_xy: float, + rand_fixed_xy: float, + rand_fixed_yaw: float, +) -> None: + """Teleport the workspace to jittered poses: the held group (part + its + presentation stand) shares one xy jitter so the part stays presented, the + fixed group (socket + rigidly-attached extras) shares an independent one. + The FIRST fixed-group asset optionally gets a yaw jitter (hole clocking).""" + if env_ids is None: + return + n = len(env_ids) + device = env.device + origins = env.scene.env_origins[env_ids] + # NOTE: this IsaacLab stack uses XYZW quaternions throughout (math utils, + # spawners, root states) -- identity is (0,0,0,1), NOT the wxyz (1,0,0,0). + identity = torch.tensor([0.0, 0.0, 0.0, 1.0], device=device).expand(n, 4) + + def place(group, mag, yaw_range=0.0): + jitter = torch.zeros((n, 3), device=device) + jitter[:, :2] = mag * (2 * torch.rand((n, 2), device=device) - 1) + for i, (name, base) in enumerate(group): + pos = torch.tensor(base, device=device) + jitter + origins + quat = identity + if i == 0 and yaw_range > 0.0: + half = 0.5 * yaw_range * (2 * torch.rand(n, device=device) - 1) + quat = torch.stack( + [torch.zeros_like(half), torch.zeros_like(half), torch.sin(half), torch.cos(half)], dim=-1) + asset = env.scene[name] + asset.write_root_pose_to_sim(torch.cat([pos, quat], dim=-1), env_ids=env_ids) + asset.write_root_velocity_to_sim(torch.zeros((n, 6), device=device), env_ids=env_ids) + + place(held_group, rand_held_xy) + place(fixed_group, rand_fixed_xy, rand_fixed_yaw) + + +def settle_and_render(env, env_ids, steps: int = 120, rt_subframes: int = 32) -> None: + """Post-reset warmup: settle the freshly jittered parts into contact, push + the teleported transforms through fabric to the renderer (IsaacLab's + reset() alone returns camera obs of the PREVIOUS state, see + notes/ISSUE_stale_reset_camera_obs.md), then flush the RTX temporal state + (DLAA history) with a Replicator-style subframe pump -- the step-loop's + one-render-per-step did NOT clear teleport ghosts (near-solid ghosts of + pre-reset poses survived 120 warmup frames and faded over ~300 recorded + frames). ``rep.orchestrator.step(rt_subframes=N)`` is NVIDIA's documented + flush for exactly this (SDG pipelines teleporting assets): it pauses the + timeline and re-renders the SAME frame N times across all render products. + Steps the whole sim, so it assumes benchmark-style global resets. + + The raw ``sim.step`` bypasses the manager's ``write_data_to_sim``, so the + home joint-position TARGET staged by ``reset_scene_to_default`` never + reaches PhysX -- the arm teleports home but the stale prior-episode target + drags it back off during settle. Re-flush each step so the PD holds home.""" + # Native 1280x720 dual-camera renders are intentionally expensive. Keep the + # production defaults above, but let expert-development runs shorten this + # reset-only anti-ghosting pass without changing task configuration. + import os + + steps = int(os.environ.get("ASSEMBLY_RESET_WARMUP_STEPS", steps)) + rt_subframes = int(os.environ.get("ASSEMBLY_RESET_RT_SUBFRAMES", rt_subframes)) + print( + f"[reset] settle warmup: physics_frames={steps} rt_subframes={rt_subframes}", + flush=True, + ) + del env_ids + for _ in range(steps): + for art in env.scene.articulations.values(): + art.write_data_to_sim() + env.sim.step(render=True) + for sensor in env.scene.sensors.values(): + sensor.update(dt=0.0, force_recompute=True) + # Hard temporal-history reset: DLAA's static-pixel blend is too sticky to + # flush by re-rendering alone (120 warmup renders + a 32-frame pump only + # FADED teleport ghosts). Toggling the AA mode tears down the accumulation + # buffers -- one FXAA frame has no history at all -- then DLAA rebuilds + # them from the NEW scene; a short pump re-converges quality before the + # first recorded frame. (rep.orchestrator.step is NOT usable here: it + # blocks on Replicator's capture pipeline, which this workflow never runs.) + # Only meaningful when rendering images. Under --disable_cameras (dev A/B + # runs) Replicator isn't loaded and there are no frames to de-ghost, so + # skip the toggle gracefully instead of crashing the reset. + try: + import omni.kit.app + import omni.replicator.core as rep + except ModuleNotFoundError: + return + app = omni.kit.app.get_app() + env.sim.set_setting("/app/player/playSimulations", False) + rep.settings.set_render_rtx_realtime(antialiasing="FXAA") + for _ in range(2): + app.update() + rep.settings.set_render_rtx_realtime(antialiasing="DLAA") + for _ in range(rt_subframes): + app.update() + env.sim.set_setting("/app/player/playSimulations", True) + # Refetch so the first post-reset obs reads the pumped, ghost-free frame. + for sensor in env.scene.sensors.values(): + sensor.update(dt=0.0, force_recompute=True) + + +def _friction_term(asset_name: str, friction: float) -> EventTermCfg: + """Startup term pinning an asset's contact friction (source: set_friction).""" + return EventTermCfg( + func=mdp.randomize_rigid_body_material, + mode="startup", + params={ + "static_friction_range": (friction, friction), + "dynamic_friction_range": (friction, friction), + "restitution_range": (0.0, 0.0), + "num_buckets": 1, + "asset_cfg": SceneEntityCfg(asset_name), + }, + ) + + +# --------------------------------------------------------------------------- +# Task +# --------------------------------------------------------------------------- + + +@configclass +class AssemblyTerminationsCfg: + time_out: TerminationTermCfg = TerminationTermCfg(func=mdp_isaac_lab.time_out) + success: TerminationTermCfg = MISSING + part_dropped: TerminationTermCfg = MISSING + + +@configclass +class AssemblyEventsCfg: + reset_all: EventTermCfg = MISSING + randomize_workspace: EventTermCfg = MISSING + friction_held: EventTermCfg | None = None + friction_fixed: EventTermCfg | None = None + friction_stand: EventTermCfg | None = None + friction_extra_0: EventTermCfg | None = None + friction_extra_1: EventTermCfg | None = None + # LAST field: terms run in declaration order, and the warmup must follow + # every pose-reset above. 120 frames clears DLAA teleport-ghosts (50 did not). + settle_warmup: EventTermCfg = EventTermCfg( + func=settle_and_render, mode="reset", params={"steps": 120}) + + +@register_task +class NISTAssemblyTask(TaskBase): + """Pick -> transport -> align -> seat (insert / mesh / thread / place).""" + + def __init__( + self, + variant: AssemblyVariant, + held: Object, + fixed: Object, + stand: Object | None = None, + extras: list[Object] = [], + reward_mode: str | None = None, + ): + super().__init__(episode_length_s=variant.episode_length_s, task_description=variant.instruction) + self.variant = variant + self.held = held + self.fixed = fixed + self.stand = stand + self.extras = list(extras) + # "staged" turns on dense idempotent RL shaping (rewards.py); None/other + # leaves the env reward-free (eval default -- terminations unchanged). + self.reward_mode = reward_mode + # Our reset event owns all part poses. + for asset in [held, fixed, stand, *self.extras]: + if asset is not None: + asset.disable_reset_pose() + + def get_termination_cfg(self): + v = self.variant + raw = TerminationTermCfg( + func=part_seated, + params={ + "held_cfg": SceneEntityCfg(self.held.name), + "fixed_cfg": SceneEntityCfg(self.fixed.name), + "seat_off": v.seat_off, + "held_base_z_off": v.held_base_z_off, + "align_tol": v.align_tol, + "seat_tol": v.seat_tol, + }, + ) + return AssemblyTerminationsCfg( + success=TerminationTermCfg(func=hold_success, params={"func": raw.func, "params": raw.params}), + part_dropped=TerminationTermCfg( + func=mdp_isaac_lab.root_height_below_minimum, + params={"minimum_height": DROP_HEIGHT, "asset_cfg": SceneEntityCfg(self.held.name)}, + ), + ) + + def get_events_cfg(self): + v = self.variant + # Held part + stand share one jitter; fixed + extras get an independent + # one (extras keep their rigid offsets). The stand bore sits on the table. + held_group = [(self.held.name, v.held_pos)] + if self.stand is not None: + held_group.insert(0, (self.stand.name, (v.held_pos[0], v.held_pos[1], TABLE_TOP_Z))) + fixed_group = [(self.fixed.name, v.fixed_pos)] + [ + (asset.name, tuple(p + o for p, o in zip(v.fixed_pos, off))) + for asset, (_, off) in zip(self.extras, v.extras) + ] + events = AssemblyEventsCfg( + reset_all=EventTermCfg( + func=mdp.reset_scene_to_default, mode="reset", params={"reset_joint_targets": True} + ), + randomize_workspace=EventTermCfg( + func=reset_assembly_workspace, + mode="reset", + params={ + "held_group": held_group, + "fixed_group": fixed_group, + "rand_held_xy": v.rand_xy, + "rand_fixed_xy": v.rand_xy if v.rand_fixed_xy is None else v.rand_fixed_xy, + "rand_fixed_yaw": v.rand_fixed_yaw, + }, + ), + friction_held=_friction_term(self.held.name, v.held_friction), + friction_fixed=_friction_term(self.fixed.name, v.fixed_friction), + ) + if self.stand is not None: + events.friction_stand = _friction_term(self.stand.name, STAND_FRICTION) + for i, asset in enumerate(self.extras): + setattr(events, f"friction_extra_{i}", _friction_term(asset.name, v.fixed_friction)) + return events + + def get_rewards_cfg(self): + """Dense staged RL reward when reward_mode="staged", else none.""" + if self.reward_mode != "staged": + return None + from assembly_bench.environments.assembly.rewards import build_rewards_cfg + return build_rewards_cfg(self.variant, self.held, self.fixed) + + def get_scene_cfg(self): + return None + + def get_mimic_env_cfg(self, arm_mode): + raise NotImplementedError + + def get_metrics(self) -> list[MetricBase]: + return [SuccessRateMetric(), ObjectMovedRateMetric(self.held)] + + def get_viewer_cfg(self) -> ViewerCfg: + # In front of and above the workspace, opposite the robot. + return get_viewer_cfg_look_at_object(lookat_object=self.held, offset=np.array([0.75, -0.55, 0.55])) diff --git a/assembly_bench/environments/assembly/variants.py b/assembly_bench/environments/assembly/variants.py new file mode 100644 index 0000000000000000000000000000000000000000..b3fc8dc69c731915dc3963c44594c2ba66f350bb --- /dev/null +++ b/assembly_bench/environments/assembly/variants.py @@ -0,0 +1,177 @@ +"""The benchmark manifest: NIST-taskboard assembly variants as pure data. + +15 assembly variants (loose-clearance only) plus one DEBUG apple→bowl sanity +check: peg-insert (round/square × S/M1/M2/L), gear-mesh (small/medium/large), +and nut-thread (M8/M12/M16/M20). Each variant fully specifies the scene content +(held / fixed / stand / extra assets and their poses) and the seat geometry that +defines success. Kept free of Isaac imports so the CLI can list ``--task`` +choices before the simulator app launches. +""" + +from dataclasses import dataclass + +# Nominal tabletop height (Factory convention: parts are authored so z=0 is +# the working surface). The SeattleLabTable's collision top is recessed 3 mm +# below the visual surface, so free parts settle at -0.003 -- authored fixture +# base plates (hole/stand bores span z=[-3,0]mm) sink flush by design. +TABLE_TOP_Z = 0.0 + +# Workspace layout (m, env-local; z = TABLE_TOP_Z is the tabletop). Parts sit +# centered under the home gripper so they land in the wrist camera's first frame. +PEG_FIXED_POS = (0.37, 0.07, TABLE_TOP_Z) # insertion hole / gear base +PEG_HELD_POS = (0.37, -0.07, TABLE_TOP_Z) # presentation stand / free part +NUT_BASE_TOP = TABLE_TOP_Z + 0.009 # NIST GMC board top surface + +# Peg suite size codes → stem diameter (mm). Assets stay named by mm. +PEG_SIZE_MM = {"S": 4, "M1": 8, "M2": 12, "L": 16} + +# Gear shaft x-offsets from the gear-base root (re-centered gen_gear USDs). +GEAR_SHAFT = {"small": 0.05075, "medium": 0.02025, "large": -0.03025} + + +@dataclass(frozen=True) +class AssemblyVariant: + family: str # "peg_insert" | "gear_mesh" | "nut_thread" | "debug" + instruction: str + held: str # registered asset name of the grasped part + fixed: str # registered asset name of the socket / base / bolt + held_pos: tuple[float, float, float] + fixed_pos: tuple[float, float, float] + stand: str | None = None # presentation bore the peg stands in (kinematic, slippery) + extras: tuple[tuple[str, tuple[float, float, float]], ...] = () # (asset, offset from fixed) + # Seat geometry: target = fixed_root + seat_off (in the fixed frame); the held + # part's base point (root + held_base_z_off) must reach it within the tolerances. + seat_off: tuple[float, float, float] = (0.0, 0.0, 0.0) + held_base_z_off: float = 0.0 + align_tol: float = 0.0025 # xy distance to target + seat_tol: float = 0.003 # seat gap (one-sided: gap < tol) + # Staged-reward geometry (reward_mode="staged"): engage gap (socket/shaft/ + # thread mouth), depth scale for Φ_depth, lift-clear height off the stand. + engage_gap: float = 0.025 + partial_socket_h: float = 0.025 + lift_clear: float = 0.03 + # Reset randomization: held(+stand) and fixed(+extras) each share one xy jitter. + rand_xy: float = 0.05 + rand_fixed_xy: float | None = None # None -> same as rand_xy; 0.0 pins the fixed asset + rand_fixed_yaw: float = 0.0 # +/- yaw jitter on the fixed asset (rect-peg clocking) + held_friction: float = 0.75 + fixed_friction: float = 0.75 + episode_length_s: float = 32.0 + + +def _peg(size_code: str, geometry: str) -> AssemblyVariant: + """One peg-insert instance (loose clearance). 'square' pegs are rectangular + (USD stem 'rect'), so they are not yaw-symmetric: the hole gets yaw jitter + the policy must match.""" + size = PEG_SIZE_MM[size_code] + stem = {"round": "round", "square": "rect"}[geometry] + return AssemblyVariant( + family="peg_insert", + instruction=f"pick up the {size} mm {geometry} peg and insert it into the hole", + held=f"asm_peg_{stem}_{size}mm_loose", + fixed=f"asm_hole_{stem}_{size}mm", + stand=f"asm_hole_{stem}_{size}mm", # a second bore presents the peg upright + held_pos=PEG_HELD_POS, + fixed_pos=PEG_FIXED_POS, + rand_fixed_yaw=0.6 if geometry == "square" else 0.0, + ) + + +def _gear(size: str) -> AssemblyVariant: + """Mesh the held gear onto its shaft on the 3-shaft base; the other two gears + flank it as kinematic mesh partners, so seating requires tooth-phase alignment.""" + flanks = tuple( + (f"asm_gear_{k}_fixed", (GEAR_SHAFT[k], 0.0, 0.0)) for k in GEAR_SHAFT if k != size + ) + return AssemblyVariant( + family="gear_mesh", + instruction=f"pick up the {size} gear and mesh it onto its shaft", + held=f"asm_gear_{size}", + fixed="asm_gear_base", + held_pos=(*PEG_HELD_POS[:2], TABLE_TOP_Z - 0.005), # gear frame z=5mm -> bottom on the table + fixed_pos=PEG_FIXED_POS, + extras=flanks, + seat_off=(GEAR_SHAFT[size], 0.0, 0.0), # the held gear's own shaft + held_friction=1.0, # grippy pads so the gear cannot slip out of the grasp + # (was 0.4 -- too low, the gear slid through the fingers) + engage_gap=0.02, + partial_socket_h=0.02, + episode_length_s=52.0, + ) + + +VARIANTS: dict[str, AssemblyVariant] = {} + +for _code in ("S", "M1", "M2", "L"): + for _geom in ("round", "square"): + VARIANTS[f"peg_{_geom}_{_code}_loose"] = _peg(_code, _geom) + +for _k in ("small", "medium", "large"): + VARIANTS[f"gear_{_k}"] = _gear(_k) + +# Nut-thread tiers. Bolt/nut dimensions (m) from the authoring script +# (gen_nutbolt_usds.py NUTBOLT table; M16 authored at the curated reference +# dims). The nut is authored at its assembled-start height, so its base sits +# head_h above the root and rests on the board at z = board_top - head_h. +NUTBOLT = { + 8: dict(head_h=0.008, shank=0.018, pitch=0.00125), + 12: dict(head_h=0.012, shank=0.02, pitch=0.00175), + 16: dict(head_h=0.010, shank=0.025, pitch=0.002), + 20: dict(head_h=0.020, shank=0.045, pitch=0.0025), +} + + +def _nut(size: int) -> AssemblyVariant: + """Thread the nut onto its bolt on the NIST GMC board (loose clearance). + Success is only reachable by helical threading (a straight push jams): nut + base descends to head_h + shank - 1.5*pitch, within pitch*0.375 + (FORGE-faithful). The bolt + board are pinned dead-center; only the nut + jitters.""" + d = NUTBOLT[size] + return AssemblyVariant( + family="nut_thread", + instruction=f"pick up the M{size} nut and thread it onto the bolt", + held=f"asm_nut_m{size}_loose", + fixed=f"asm_bolt_m{size}_loose", + held_pos=(0.30, 0.10, NUT_BASE_TOP - d["head_h"]), + fixed_pos=(0.37, 0.0, NUT_BASE_TOP), + extras=(("asm_nist_board", (0.0, 0.0, 0.0)),), + seat_off=(0.0, 0.0, d["head_h"] + d["shank"] - 1.5 * d["pitch"]), + held_base_z_off=d["head_h"], + seat_tol=0.375 * d["pitch"], + engage_gap=0.01, # thread mouth is close; engage = nut on the bolt tip + rand_xy=0.03, + rand_fixed_xy=0.0, + held_friction=1.0, # grippy pads so the nut survives lift/regrip cycles + fixed_friction=0.2, + episode_length_s=120.0, # threading + release-regrip-rotate cycles + ) + + +# M4 dropped: 3.2 mm hex is below the Robotiq pad band when open on the board. +for _size in (8, 12, 16, 20): + VARIANTS[f"nut_M{_size}"] = _nut(_size) + +# --------------------------------------------------------------------------- +# DEBUG ONLY — RoboLab apple→bowl. Not part of the NIST assembly matrix. +# Prompt matches RoboLab's "Put the X in the Y" convention. +# --------------------------------------------------------------------------- +_APPLE_ORIGIN_ABOVE_BOTTOM = 0.0188 # scaled apple_01 AABB +_BOWL_ORIGIN_ABOVE_BOTTOM = 0.0275 # YCB bowl AABB + +VARIANTS["debug"] = AssemblyVariant( + family="debug", + instruction="Put the apple in the bowl", + held="asm_debug_apple", + fixed="asm_debug_bowl", + held_pos=(*PEG_HELD_POS[:2], _APPLE_ORIGIN_ABOVE_BOTTOM), + fixed_pos=(*PEG_FIXED_POS[:2], _BOWL_ORIGIN_ABOVE_BOTTOM), + seat_off=(0.0, 0.0, 0.02), # apple resting inside the bowl + align_tol=0.05, # within the bowl rim (~8 cm radius) + seat_tol=0.04, + rand_xy=0.03, + rand_fixed_xy=0.02, + held_friction=1.0, + fixed_friction=0.6, + episode_length_s=32.0, +) diff --git a/contract.json b/contract.json new file mode 100644 index 0000000000000000000000000000000000000000..a1b2436da7aa4632abe91e91fbafd2acd9d7bfcf --- /dev/null +++ b/contract.json @@ -0,0 +1,101 @@ +{ + "robot_type": "droid_franka_robotiq", + "robot_class": "arm7g1", + "control_rate": 15, + "features": { + "camera_obs/front_cam_rgb": { + "role": "observation", + "type": "rgb" + }, + "camera_obs/wrist_camera_rgb": { + "role": "observation", + "type": "rgb" + }, + "policy/joint_pos": { + "role": "observation", + "state_type": "JOINT_ABS_POS", + "state_representation": "REAL", + "dtype": "float32", + "units": "rad", + "shape": [7], + "names": [ + "panda_joint1", + "panda_joint2", + "panda_joint3", + "panda_joint4", + "panda_joint5", + "panda_joint6", + "panda_joint7" + ] + }, + "policy/gripper_pos": { + "role": "observation", + "state_type": "JOINT_ABS_POS", + "state_representation": "NORM01", + "dtype": "float32", + "units": "none", + "shape": [1], + "names": ["gripper.opening"] + }, + "policy/eef_pos": { + "role": "observation", + "state_type": "EE_ABS_POS", + "state_representation": "XYZ", + "frame": "world", + "dtype": "float32", + "units": "m", + "shape": [3], + "names": ["eef_pos.x", "eef_pos.y", "eef_pos.z"] + }, + "policy/eef_quat": { + "role": "observation", + "state_type": "EE_ABS_ROT", + "state_representation": "QUATWXYZ", + "frame": "world", + "dtype": "float32", + "units": "none", + "shape": [4], + "names": ["eef_quat.w", "eef_quat.x", "eef_quat.y", "eef_quat.z"] + }, + "policy/held_part_pose": { + "role": "observation", + "state_type": "OBJECT_ABS_POSE", + "state_representation": "XYZ_QUATXYZW", + "frame": "env_local", + "dtype": "float32", + "units": "m", + "shape": [7], + "names": ["held.x", "held.y", "held.z", "held.qx", "held.qy", "held.qz", "held.qw"], + "comment": "Privileged (PA-RL critic only); policy adapters must not consume." + }, + "policy/fixed_part_pose": { + "role": "observation", + "state_type": "OBJECT_ABS_POSE", + "state_representation": "XYZ_QUATXYZW", + "frame": "env_local", + "dtype": "float32", + "units": "m", + "shape": [7], + "names": ["fixed.x", "fixed.y", "fixed.z", "fixed.qx", "fixed.qy", "fixed.qz", "fixed.qw"], + "comment": "Privileged (PA-RL critic only); policy adapters must not consume." + }, + "action": { + "role": "action", + "state_type": "JOINT_ABS_POS", + "dtype": "float32", + "shape": [8], + "names": [ + "panda_joint1", + "panda_joint2", + "panda_joint3", + "panda_joint4", + "panda_joint5", + "panda_joint6", + "panda_joint7", + "gripper.open_close" + ], + "comment": "Dims 0-6: absolute joint targets (rad, REAL). Dim 7: gripper (NORM01; >0.5 closes)." + } + }, + "comment": "assembly_bench DROID (droid_abs_joint_pos): Franka Panda 7-DoF + Robotiq 2F-85. Wire keys match Arena obs groups. One action feature per the v0 contract: 7 absolute joint targets + binary gripper. Omits policy/actions and policy/robot_joint_pos (redundant with action echo / full articulation)." +} diff --git a/env.py b/env.py new file mode 100644 index 0000000000000000000000000000000000000000..faa71e87681510039a52b644287bca12d377dfda --- /dev/null +++ b/env.py @@ -0,0 +1,299 @@ +"""Assembly bench EnvHub entry — LeRobot ``make_env`` for Isaac Lab Arena. + +Docs: https://huggingface.co/docs/lerobot/en/envhub_isaaclab_arena + +Variant selection (Arena ``--task``) is *not* ``cfg.task`` — that field is the +natural-language prompt on ``IsaaclabArenaEnv``. Pass the variant via kwargs: + + --env.kwargs='{"variant": "peg_round_M1_loose", "reward": "none"}' +""" + +from __future__ import annotations + +import importlib.util +import logging +import os +import sys +from pathlib import Path +from typing import Any + +import gymnasium as gym +import yaml + +from huggingface_hub import hf_hub_download, snapshot_download + +# Hub id for sibling downloads + full-repo snapshot (assets + package). +HUB_REPO_ID = os.environ.get("ASSEMBLY_BENCH_HUB", "lukasskellijs/env_assembly_bench") +EXAMPLE_ENVS = "example_envs.yaml" + +# Default proprio / cameras match assembly_bench/contract.json (DROID 8-D). +_DEFAULT_STATE_KEYS = "joint_pos,gripper_pos,eef_pos,eef_quat" +_DEFAULT_CAMERA_KEYS = "front_cam_rgb,wrist_camera_rgb" +_DEFAULT_STATE_DIM = 15 # 7+1+3+4 +_DEFAULT_ACTION_DIM = 8 +_DEFAULT_VARIANT = "peg_round_M1_loose" +_DEFAULT_EMBODIMENT = "droid_abs_joint_pos_softmimic" +_MICROWAVE_PROMPT = "Reach out to the microwave and open it." # IsaaclabArenaEnv default + +_REPO_ROOT: str | None = None + + +def _ensure_repo_root() -> str: + """Put this EnvHub checkout/snapshot on sys.path so ``assembly_bench`` imports.""" + global _REPO_ROOT + if _REPO_ROOT is not None: + return _REPO_ROOT + + here = Path(__file__).resolve().parent + if (here / "assembly_bench" / "environments" / "assembly" / "assembly.py").is_file(): + _REPO_ROOT = str(here) + else: + # Loaded as a lone hub file — pull the full repo (USDs + package). + _REPO_ROOT = snapshot_download(repo_id=HUB_REPO_ID) + + if _REPO_ROOT not in sys.path: + sys.path.insert(0, _REPO_ROOT) + return _REPO_ROOT + + +def _download_hub_file(filename: str) -> str: + root = _ensure_repo_root() + local = Path(root) / filename + if local.is_file(): + return str(local) + return hf_hub_download(repo_id=HUB_REPO_ID, filename=filename) + + +def _download_and_import(filename: str): + """Load a sibling module from the hub snapshot (absolute imports).""" + local_path = _download_hub_file(filename) + module_dir = os.path.dirname(local_path) + if module_dir not in sys.path: + sys.path.insert(0, module_dir) + + module_name = filename.replace(".py", "") + with open(local_path) as f: + content = f.read() + content = content.replace("from .errors import", "from errors import") + content = content.replace("from .isaaclab_env_wrapper import", "from isaaclab_env_wrapper import") + + spec = importlib.util.spec_from_file_location(module_name, local_path) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + exec(compile(content, local_path, "exec"), module.__dict__) # noqa: S102 + return module + + +try: + from .errors import IsaacLabArenaCameraKeyError, IsaacLabArenaStateKeyError + from .isaaclab_env_wrapper import IsaacLabEnvWrapper, cleanup_isaaclab +except ImportError: + _errors = _download_and_import("errors.py") + _wrapper = _download_and_import("isaaclab_env_wrapper.py") + IsaacLabEnvWrapper = _wrapper.IsaacLabEnvWrapper + cleanup_isaaclab = _wrapper.cleanup_isaaclab + IsaacLabArenaCameraKeyError = _errors.IsaacLabArenaCameraKeyError + IsaacLabArenaStateKeyError = _errors.IsaacLabArenaStateKeyError + + +def validate_config( + env, + state_keys: tuple[str, ...], + camera_keys: tuple[str, ...], + cfg_state_dim: int, + cfg_action_dim: int, +) -> None: + """Check observation keys and dims against the live Isaac Lab managers.""" + obs_manager = env.observation_manager + active_terms = obs_manager.active_terms + policy_terms = set(active_terms.get("policy", [])) + camera_terms = set(active_terms.get("camera_obs", [])) + + missing_state = [k for k in state_keys if k not in policy_terms] + if missing_state: + raise IsaacLabArenaStateKeyError(missing_state, policy_terms) + + missing_cam = [k for k in camera_keys if k not in camera_terms] + if missing_cam: + raise IsaacLabArenaCameraKeyError(missing_cam, camera_terms) + + env_action_dim = env.action_space.shape[-1] + if cfg_action_dim != env_action_dim: + raise ValueError(f"action_dim mismatch: config={cfg_action_dim}, env={env_action_dim}") + + policy_dims = obs_manager.group_obs_dim.get("policy", []) + policy_names = active_terms.get("policy", []) + term_dims = dict(zip(policy_names, policy_dims, strict=False)) + + expected_state_dim = 0 + for key in state_keys: + if key in term_dims: + shape = term_dims[key] + dim = 1 + for s in shape if isinstance(shape, (tuple, list)) else [shape]: + dim *= s + expected_state_dim += dim + + if cfg_state_dim != expected_state_dim: + raise ValueError( + f"state_dim mismatch: config={cfg_state_dim}, " + f"computed={expected_state_dim}. " + f"Term dims: {term_dims}" + ) + + logging.info(f"Validated: state_keys={state_keys}, camera_keys={camera_keys}") + + +def resolve_environment_alias(environment: str) -> str: + """Map ``--env.environment`` alias → dotted class path.""" + with open(_download_hub_file(EXAMPLE_ENVS)) as f: + envs_mapping = yaml.safe_load(f) + return f"{envs_mapping['repo']['base_module']}.{envs_mapping['repo']['envs'][environment]}" + + +def _cfg_get(cfg: Any, name: str, default: Any = None) -> Any: + if cfg is None: + return default + return getattr(cfg, name, default) + + +def _create_isaaclab_env(config: dict, n_envs: int) -> dict[str, dict[int, gym.vector.VectorEnv]]: + """Boot Kit, build AssemblyBench via Arena, wrap for EnvHub.""" + _ensure_repo_root() + + from isaaclab.app import AppLauncher + + if config.get("enable_pinocchio", False): + import pinocchio # noqa: F401 + + config = {**config, "num_envs": n_envs} + + # Only Kit-relevant fields — extra Arena/EnvHub keys confuse nothing, but a + # minimal Namespace keeps AppLauncher resolution predictable. + logging.info("Launching IsaacLab simulation app...") + app_launcher = AppLauncher( + headless=bool(config.get("headless", True)), + enable_cameras=bool(config.get("enable_cameras", True)), + device=config.get("device") or "cuda:0", + ) + + # Package imports need the hub root on path (assets resolve from scene.py). + from assembly_bench.environments.assembly.assembly import make_assembly_env + from assembly_bench.environments.assembly.variants import VARIANTS + + variant = config.get("variant") or _DEFAULT_VARIANT + if variant not in VARIANTS: + raise ValueError(f"Unknown variant {variant!r}. Choices: {sorted(VARIANTS)}") + # Confirm the EnvHub alias maps to this package (discoverability / typos). + resolve_environment_alias(config.get("environment") or "assembly_bench") + + # Shared factory: Arena CLI defaults + reward/hdr wiring. + raw_env = make_assembly_env( + task=variant, + num_envs=n_envs, + embodiment=config.get("embodiment") or _DEFAULT_EMBODIMENT, + hdr=config.get("hdr") or "asm_machine_shop", + light_intensity=float(config.get("light_intensity") or 1500.0), + reward=config.get("reward") or "none", + ) + + render_mode = "rgb_array" if config.get("enable_cameras", False) else None + if render_mode and hasattr(raw_env, "render_mode"): + raw_env.render_mode = render_mode + + state_keys = tuple(k.strip() for k in (config.get("state_keys") or "").split(",") if k.strip()) + camera_keys = tuple(k.strip() for k in (config.get("camera_keys") or "").split(",") if k.strip()) + try: + validate_config( + raw_env, + state_keys, + camera_keys, + config.get("state_dim", _DEFAULT_STATE_DIM), + config.get("action_dim", _DEFAULT_ACTION_DIM), + ) + except (IsaacLabArenaCameraKeyError, IsaacLabArenaStateKeyError, ValueError): + cleanup_isaaclab(raw_env, app_launcher) + raise + except Exception as e: + logging.error(f"Validation failed with unexpected error: {type(e).__name__}: {e}") + cleanup_isaaclab(raw_env, app_launcher) + raise + + # NL prompt: cfg.task unless still the Arena microwave default → use variant text. + task_prompt = config.get("task") + if not task_prompt or task_prompt == _MICROWAVE_PROMPT: + task_prompt = VARIANTS[variant].instruction + + # 15 Hz control; fall back to the variant's authored duration (seconds). + episode_length = config.get("episode_length") + if episode_length in (None, 300): + episode_length = int(VARIANTS[variant].episode_length_s * 15) + + environment = config.get("environment") or "assembly_bench" + wrapped_env = IsaacLabEnvWrapper( + raw_env, + episode_length=int(episode_length), + task=task_prompt, + render_mode=render_mode, + simulation_app=app_launcher, + ) + logging.info( + f"Created: {environment} variant={variant} n_envs={wrapped_env.num_envs} " + f"render_mode={render_mode}" + ) + return {environment: {0: wrapped_env}} + + +def make_env( + n_envs: int = 1, + use_async_envs: bool = False, # noqa: ARG001 — GPU batching; AsyncVectorEnv unused + cfg: Any | None = None, +) -> dict[str, dict[int, gym.vector.VectorEnv]]: + """EnvHub entry point. ``cfg`` is typically ``IsaaclabArenaEnv`` (or SimpleNamespace).""" + if n_envs < 1: + raise ValueError("`n_envs` must be at least 1") + + environment = _cfg_get(cfg, "environment", "assembly_bench") + if environment is None: + raise ValueError("No 'environment' specified (expected 'assembly_bench').") + + # IsaaclabArenaEnv.__post_init__ promotes kwargs (variant/reward/hdr) onto cfg. + variant = _cfg_get(cfg, "variant", _DEFAULT_VARIANT) + reward = _cfg_get(cfg, "reward", "none") + hdr = _cfg_get(cfg, "hdr", "asm_machine_shop") + light_intensity = _cfg_get(cfg, "light_intensity", 1500.0) + + config = { + "environment": environment, + "embodiment": _cfg_get(cfg, "embodiment", _DEFAULT_EMBODIMENT), + "object": _cfg_get(cfg, "object", None), + "mimic": _cfg_get(cfg, "mimic", False), + "teleop_device": _cfg_get(cfg, "teleop_device", None), + "seed": _cfg_get(cfg, "seed", 42), + "device": _cfg_get(cfg, "device", "cuda:0"), + "disable_fabric": _cfg_get(cfg, "disable_fabric", False), + "enable_cameras": _cfg_get(cfg, "enable_cameras", True), + "headless": _cfg_get(cfg, "headless", True), + "enable_pinocchio": _cfg_get(cfg, "enable_pinocchio", False), + "episode_length": _cfg_get(cfg, "episode_length", None), + "state_dim": _cfg_get(cfg, "state_dim", _DEFAULT_STATE_DIM), + "action_dim": _cfg_get(cfg, "action_dim", _DEFAULT_ACTION_DIM), + "camera_height": _cfg_get(cfg, "camera_height", 720), + "camera_width": _cfg_get(cfg, "camera_width", 1280), + "video": _cfg_get(cfg, "video", False), + "video_length": _cfg_get(cfg, "video_length", 100), + "video_interval": _cfg_get(cfg, "video_interval", 200), + "state_keys": _cfg_get(cfg, "state_keys", _DEFAULT_STATE_KEYS), + "camera_keys": _cfg_get(cfg, "camera_keys", _DEFAULT_CAMERA_KEYS) or "", + "task": _cfg_get(cfg, "task", None), + "variant": variant, + "reward": reward, + "hdr": hdr, + "light_intensity": light_intensity, + } + + logging.info( + f"EnvHub make_env: environment={environment}, variant={variant}, n_envs={n_envs}, " + f"headless={config['headless']}, enable_cameras={config['enable_cameras']}" + ) + return _create_isaaclab_env(config, n_envs) diff --git a/errors.py b/errors.py new file mode 100644 index 0000000000000000000000000000000000000000..d928e618318049fe32cae60d59c5e083084fb013 --- /dev/null +++ b/errors.py @@ -0,0 +1,33 @@ +"""EnvHub errors for assembly_bench (Isaac Lab Arena).""" + + +class IsaacLabArenaError(RuntimeError): + """Base exception for IsaacLab Arena environment errors.""" + + def __init__(self, message: str = "IsaacLab Arena error"): + self.message = message + super().__init__(self.message) + + +class IsaacLabArenaConfigError(IsaacLabArenaError): + """Invalid observation / action configuration against the live env.""" + + def __init__(self, invalid: list, available: list, key_type: str = "keys"): + msg = f"Invalid {key_type}: {invalid}. Available: {sorted(available)}" + super().__init__(msg) + self.invalid = invalid + self.available = available + + +class IsaacLabArenaCameraKeyError(IsaacLabArenaConfigError): + """camera_keys do not match available camera observation terms.""" + + def __init__(self, invalid: list, available: list): + super().__init__(invalid, available, "camera_keys") + + +class IsaacLabArenaStateKeyError(IsaacLabArenaConfigError): + """state_keys do not match available policy observation terms.""" + + def __init__(self, invalid: list, available: list): + super().__init__(invalid, available, "state_keys") diff --git a/example_envs.yaml b/example_envs.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6b1a013b096f843d4286e42fbe54e8f9a96ceeec --- /dev/null +++ b/example_envs.yaml @@ -0,0 +1,6 @@ +# EnvHub aliases → AssemblyBenchEnvironment class path. +# Key is --env.environment; value is module path under base_module. +repo: + base_module: assembly_bench.environments.assembly + envs: + assembly_bench: assembly.AssemblyBenchEnvironment diff --git a/hud_env.py b/hud_env.py new file mode 100644 index 0000000000000000000000000000000000000000..98f37e4fe6c609fbd7fba70d43775d6508a30989 --- /dev/null +++ b/hud_env.py @@ -0,0 +1,118 @@ +"""HUD environment that loads assembly_bench from EnvHub (Hub download path). + +Standard HUD robot flow: ``env.gym(make_env)`` derives the contract/capability +and serves over the robot WebSocket. ``make_env`` pulls this repo from the Hub +(the same way LeRobot EnvHub does), then builds the Isaac Lab Arena sim. + +Serve (isaac conda, GPU):: + + OMNI_KIT_ACCEPT_EULA=YES ACCEPT_EULA=Y PRIVACY_CONSENT=Y \\ + python -m hud.environment.server hud_env.py --port 8765 + +Scene config is the factory signature (``variant`` / ``num_envs`` / …); episodic +args (seed) go through ``sim.reset``. +""" + +from __future__ import annotations + +import importlib.util +import os +import sys +from pathlib import Path +from types import SimpleNamespace + +from hud import Environment + +HUB = os.environ.get("ASSEMBLY_BENCH_HUB", "lukasskellijs/env_assembly_bench") + +# Authored contract next to this file (absolute — cwd-relative would silently drift). +_CONTRACT = Path(__file__).resolve().parent / "contract.json" + + +def make_env( + variant: str = "peg_round_M1_loose", + num_envs: int = int(os.environ.get("ASSEMBLY_NUM_ENVS", "1")), + embodiment: str = os.environ.get("ASSEMBLY_EMBODIMENT", "droid_abs_joint_pos_softmimic"), + reward: str = os.environ.get("ASSEMBLY_REWARD", "none"), +): + """Download EnvHub package (if needed) and return the vectorized Arena env.""" + from huggingface_hub import snapshot_download + + # Fresh-enough Hub snapshot; local checkout of this repo also works when + # ASSEMBLY_BENCH_HUB points at a path... snapshot always hits the Hub id. + root = snapshot_download(repo_id=HUB) + env_py = Path(root) / "env.py" + if not env_py.is_file(): + raise FileNotFoundError(f"Hub repo {HUB} missing env.py (got {root})") + + # Prefer Hub package over any other assembly_bench on sys.path. + if root in sys.path: + sys.path.remove(root) + sys.path.insert(0, root) + + spec = importlib.util.spec_from_file_location("env_assembly_bench_hub_env", env_py) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + + # IsaaclabArenaEnv-shaped config (no lerobot import required in the isaac env). + cfg = SimpleNamespace( + environment="assembly_bench", + embodiment=embodiment, + object=None, + mimic=False, + teleop_device=None, + seed=42, + device=os.environ.get("ASSEMBLY_DEVICE", "cuda:0"), + disable_fabric=False, + enable_cameras=True, + headless=True, + enable_pinocchio=False, + episode_length=300, # hub make_env replaces 300 with variant duration + state_dim=15, + action_dim=8, + camera_height=720, + camera_width=1280, + video=False, + video_length=100, + video_interval=200, + state_keys="joint_pos,gripper_pos,eef_pos,eef_quat", + camera_keys="front_cam_rgb,wrist_camera_rgb", + task=None, # filled from variants.instruction + variant=variant, + reward=reward, + hdr="asm_machine_shop", + light_intensity=1500.0, + ) + suites = module.make_env(n_envs=num_envs, use_async_envs=False, cfg=cfg) + # {suite: {task_id: VectorEnv}} → the live vectorized env + return next(iter(next(iter(suites.values())).values())) + + +env = Environment(name="assembly-bench") +sim = env.gym( + make_env, + contract=_CONTRACT, + variant="peg_round_M1_loose", +) + + +@env.template(id="assembly") +async def assembly( + variant: str = "peg_round_M1_loose", + seed: int = 0, + num_envs: int | None = None, + embodiment: str | None = None, + reward: str | None = None, +): + """One assembly episode. Env-defining args rebuild the scene when they change.""" + reset_kw: dict = {"variant": variant, "seed": seed} + if num_envs is not None: + reset_kw["num_envs"] = num_envs + if embodiment is not None: + reset_kw["embodiment"] = embodiment + if reward is not None: + reset_kw["reward"] = reward + ep = await sim.reset(**reset_kw) + yield {"prompt": ep["prompt"], "robot": {"token": ep["token"]}} + yield await sim.result(token=ep["token"]) diff --git a/isaaclab_env_wrapper.py b/isaaclab_env_wrapper.py new file mode 100644 index 0000000000000000000000000000000000000000..717de6f2fcd853f191e9bbfccfa8933e3f576e82 --- /dev/null +++ b/isaaclab_env_wrapper.py @@ -0,0 +1,215 @@ +"""VectorEnv wrapper adapting Isaac Lab's GPU-batched env for LeRobot EnvHub.""" + +from __future__ import annotations + +import atexit +import logging +import os +import signal +from contextlib import suppress +from typing import Any + +import gymnasium as gym +import numpy as np +import torch + +try: + from .errors import IsaacLabArenaError +except ImportError: + from errors import IsaacLabArenaError + + +def cleanup_isaaclab(env, simulation_app) -> None: + """Close the gym env and Kit app; ignore signals mid-teardown.""" + old_sigint = signal.signal(signal.SIGINT, signal.SIG_IGN) + old_sigterm = signal.signal(signal.SIGTERM, signal.SIG_IGN) + try: + with suppress(Exception): + if env is not None: + env.close() + with suppress(Exception): + if simulation_app is not None: + simulation_app.app.close() + finally: + signal.signal(signal.SIGINT, old_sigint) + signal.signal(signal.SIGTERM, old_sigterm) + + +class IsaacLabEnvWrapper(gym.vector.AsyncVectorEnv): + """Isaac Lab already vectorizes on GPU; this exposes an AsyncVectorEnv API.""" + + metadata = {"render_modes": ["rgb_array"], "render_fps": 15} + _cleanup_in_progress = False + + def __init__( + self, + env, + episode_length: int = 500, + task: str | None = None, + render_mode: str | None = "rgb_array", + simulation_app=None, + ): + self._env = env + self._num_envs = env.num_envs + self._episode_length = episode_length + self._closed = False + self.render_mode = render_mode + self._simulation_app = simulation_app + + self.observation_space = env.observation_space + self.action_space = env.action_space + self.single_observation_space = env.observation_space + self.single_action_space = env.action_space + self.task = task + + if hasattr(env, "metadata") and env.metadata: + self.metadata = {**self.metadata, **env.metadata} + + atexit.register(self._cleanup) + signal.signal(signal.SIGINT, self._signal_handler) + signal.signal(signal.SIGTERM, self._signal_handler) + + def _signal_handler(self, signum, frame): + if IsaacLabEnvWrapper._cleanup_in_progress: + return + IsaacLabEnvWrapper._cleanup_in_progress = True + logging.info(f"Received signal {signum}, cleaning up...") + self._cleanup() + os._exit(0) + + def _check_closed(self): + if self._closed: + raise IsaacLabArenaError() + + @property + def unwrapped(self): + return self + + @property + def num_envs(self) -> int: + return self._num_envs + + @property + def _max_episode_steps(self) -> int: + return self._episode_length + + @property + def device(self) -> str: + return getattr(self._env, "device", "cpu") + + def reset( + self, + *, + seed: int | list[int] | None = None, + options: dict[str, Any] | None = None, + ) -> tuple[dict[str, Any], dict[str, Any]]: + self._check_closed() + if isinstance(seed, (list, tuple, range)): + seed = seed[0] if len(seed) > 0 else None + + obs, info = self._env.reset(seed=seed, options=options) + + if "final_info" not in info: + zeros = np.zeros(self._num_envs, dtype=bool) + info["final_info"] = {"is_success": zeros} + + return obs, info + + def step( + self, actions: np.ndarray | torch.Tensor + ) -> tuple[dict, np.ndarray, np.ndarray, np.ndarray, dict]: + self._check_closed() + if isinstance(actions, np.ndarray): + actions = torch.from_numpy(actions).to(self._env.device) + + obs, reward, terminated, truncated, info = self._env.step(actions) + + reward = reward.cpu().numpy().astype(np.float32) + terminated = terminated.cpu().numpy().astype(bool) + truncated = truncated.cpu().numpy().astype(bool) + + is_success = self._get_success(terminated, truncated) + info["final_info"] = {"is_success": is_success} + + return obs, reward, terminated, truncated, info + + def _get_success(self, terminated: np.ndarray, truncated: np.ndarray) -> np.ndarray: + is_success = np.zeros(self._num_envs, dtype=bool) + + if not hasattr(self._env, "termination_manager"): + return is_success & (terminated | truncated) + + term_manager = self._env.termination_manager + if not hasattr(term_manager, "get_term"): + return is_success & (terminated | truncated) + + success_tensor = term_manager.get_term("success") + if success_tensor is None: + return is_success & (terminated | truncated) + + is_success = success_tensor.cpu().numpy().astype(bool) + return is_success & (terminated | truncated) + + def call(self, method_name: str, *args, **kwargs) -> list[Any]: + if method_name == "_max_episode_steps": + return [self._episode_length] * self._num_envs + if method_name == "task": + return [self.task] * self._num_envs + if method_name == "render": + return self.render_all() + + if hasattr(self._env, method_name): + attr = getattr(self._env, method_name) + result = attr(*args, **kwargs) if callable(attr) else attr + if isinstance(result, list): + return result + return [result] * self._num_envs + + raise AttributeError(f"IsaacLab-Arena has no method/attribute '{method_name}'") + + def render_all(self) -> list[np.ndarray]: + self._check_closed() + frames = self.render() + if frames is None: + placeholder = np.zeros((480, 640, 3), dtype=np.uint8) + return [placeholder] * self._num_envs + return [frames] * self._num_envs + + def render(self) -> np.ndarray | None: + self._check_closed() + if self.render_mode != "rgb_array": + return None + + frames = self._env.render() if hasattr(self._env, "render") else None + if frames is None: + return None + + if isinstance(frames, torch.Tensor): + frames = frames.cpu().numpy() + + return frames[0] if frames.ndim == 4 else frames + + def _cleanup(self) -> None: + if self._closed: + return + self._closed = True + IsaacLabEnvWrapper._cleanup_in_progress = True + logging.info("Cleaning up IsaacLab Arena environment...") + cleanup_isaaclab(self._env, self._simulation_app) + + def close(self) -> None: + self._cleanup() + + @property + def envs(self) -> list[IsaacLabEnvWrapper]: + return [self] * self._num_envs + + def __del__(self): + self._cleanup() + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self._cleanup() + return False diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000000000000000000000000000000000000..cfd171cd4718ba0f90f121dac5e691b3df860ea6 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,25 @@ +[build-system] +requires = ["setuptools>=64"] +build-backend = "setuptools.build_meta" + +[project] +name = "env-assembly-bench" +version = "0.1.0" +description = "NIST assembly benchmark — Isaac Lab Arena EnvHub package" +requires-python = ">=3.10" +dependencies = [ + "huggingface_hub", + "PyYAML", + "gymnasium", + "numpy", +] +# isaaclab / isaaclab_arena are installed separately (see README). + +[tool.setuptools.packages.find] +include = ["assembly_bench*"] + +[tool.setuptools.package-data] +assembly_bench = [ + "assets/parts/*.usd", + "assets/backgrounds/indoors/*.hdr", +] diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..d13110414c7893266bedc7792d42016959e7e0f3 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,6 @@ +# EnvHub glue only. Isaac Sim 6 + Isaac Lab Arena must already be installed +# in the active environment (see README). +huggingface_hub>=0.20 +PyYAML>=6.0 +gymnasium>=0.29 +numpy>=1.26 diff --git a/scripts/smoke_hub_load.py b/scripts/smoke_hub_load.py new file mode 100644 index 0000000000000000000000000000000000000000..3a944cedb7d3e2710a9a1bd8180982e27fb6ac68 --- /dev/null +++ b/scripts/smoke_hub_load.py @@ -0,0 +1,117 @@ +"""Smoke-test: load make_env from a Hub snapshot (or local checkout) and step. + +Usage (isaac conda, GPU):: + + ASSEMBLY_BENCH_HUB=lukasskellijs/env_assembly_bench \\ + ACCEPT_EULA=Y PRIVACY_CONSENT=Y OMNI_KIT_ACCEPT_EULA=YES \\ + python scripts/smoke_hub_load.py --from-hub --steps 5 + + # Local checkout (no Hub round-trip): + python scripts/smoke_hub_load.py --local --steps 5 +""" + +from __future__ import annotations + +import argparse +import importlib.util +import logging +import os +import sys +from pathlib import Path +from types import SimpleNamespace + +ROOT = Path(__file__).resolve().parents[1] + + +def _load_env_module(root: Path): + env_py = root / "env.py" + if root_str := str(root): + if root_str in sys.path: + sys.path.remove(root_str) + sys.path.insert(0, root_str) + spec = importlib.util.spec_from_file_location("env_assembly_bench_env", env_py) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +def main() -> None: + logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s") + p = argparse.ArgumentParser() + p.add_argument("--from-hub", action="store_true", help="snapshot_download from Hub") + p.add_argument("--local", action="store_true", help="use this checkout") + p.add_argument("--steps", type=int, default=5) + p.add_argument("--variant", default="peg_round_M1_loose") + p.add_argument("--num-envs", type=int, default=1) + args = p.parse_args() + if args.from_hub == args.local: + p.error("pass exactly one of --from-hub or --local") + + if args.from_hub: + from huggingface_hub import snapshot_download + + hub = os.environ.get("ASSEMBLY_BENCH_HUB", "lukasskellijs/env_assembly_bench") + # Force a re-fetch when HF_HUB_FORCE_DOWNLOAD is set; otherwise use cache. + root = Path( + snapshot_download( + repo_id=hub, + force_download=os.environ.get("HF_HUB_FORCE_DOWNLOAD", "") == "1", + ) + ) + logging.info("Hub snapshot at %s", root) + else: + root = ROOT + logging.info("Local checkout at %s", root) + + module = _load_env_module(root) + cfg = SimpleNamespace( + environment="assembly_bench", + embodiment="droid_abs_joint_pos_softmimic", + object=None, + mimic=False, + teleop_device=None, + seed=0, + device="cuda:0", + disable_fabric=False, + enable_cameras=True, + headless=True, + enable_pinocchio=False, + episode_length=50, + state_dim=15, + action_dim=8, + camera_height=720, + camera_width=1280, + video=False, + video_length=10, + video_interval=15, + state_keys="joint_pos,gripper_pos,eef_pos,eef_quat", + camera_keys="front_cam_rgb,wrist_camera_rgb", + task=None, + variant=args.variant, + reward="none", + hdr="asm_machine_shop", + light_intensity=1500.0, + ) + suites = module.make_env(n_envs=args.num_envs, use_async_envs=False, cfg=cfg) + env = next(iter(next(iter(suites.values())).values())) + obs, info = env.reset(seed=0) + logging.info("reset ok; task=%r keys=%s", env.task, list(obs.keys()) if isinstance(obs, dict) else type(obs)) + for i in range(args.steps): + action = env.action_space.sample() + obs, reward, terminated, truncated, info = env.step(action) + logging.info( + "step %d reward=%s term=%s trunc=%s", + i, + getattr(reward, "mean", lambda: reward)(), + terminated, + truncated, + ) + if terminated.any() or truncated.any(): + obs, info = env.reset() + env.close() + logging.info("smoke ok") + + +if __name__ == "__main__": + main()