| |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import sys |
| from pathlib import Path |
|
|
| PROJECT_ROOT = Path(__file__).resolve().parents[1] |
| if str(PROJECT_ROOT) not in sys.path: |
| sys.path.insert(0, str(PROJECT_ROOT)) |
|
|
| from dovla_cil.generation.maniskill_lattice import ( |
| ManiSkillLatticeConfig, |
| generate_maniskill_lattice, |
| ) |
|
|
|
|
| def main(argv: list[str] | None = None) -> int: |
| parser = argparse.ArgumentParser( |
| description="Generate measured same-state CIL branches from ManiSkill demonstrations." |
| ) |
| parser.add_argument("--demo", type=Path, required=True) |
| parser.add_argument("--out", type=Path, required=True) |
| parser.add_argument("--num-groups", type=int, default=128) |
| parser.add_argument( |
| "--group-offset", |
| type=int, |
| default=0, |
| help="Start index in the deterministic global episode-step plan.", |
| ) |
| parser.add_argument("--k", type=int, default=8) |
| parser.add_argument("--horizon", type=int, default=4) |
| parser.add_argument("--seed", type=int, default=0) |
| parser.add_argument("--shard-size", type=int, default=1024) |
| parser.add_argument("--env-id", default="PickCube-v1") |
| parser.add_argument("--obs-mode", default="state") |
| parser.add_argument( |
| "--image-quality", |
| type=int, |
| default=90, |
| help="JPEG quality for observations.h5 when obs-mode includes RGB.", |
| ) |
| parser.add_argument("--control-mode", default="pd_ee_delta_pose") |
| parser.add_argument("--sim-backend", default="physx_cuda") |
| parser.add_argument( |
| "--render-backend", |
| default="gpu", |
| help="ManiSkill render backend. Use 'gpu' for standard task assets and RGB rendering.", |
| ) |
| parser.add_argument( |
| "--parallel-branches", |
| action=argparse.BooleanOptionalAction, |
| default=True, |
| help="Execute K same-state interventions in K vectorized ManiSkill environments.", |
| ) |
| parser.add_argument( |
| "--state-storage", |
| choices=("archive", "files", "none"), |
| default="archive", |
| help="Persist replay states in one archive, separate files, or only by source reference.", |
| ) |
| parser.add_argument( |
| "--state-batch-size", |
| type=int, |
| default=1, |
| help="Number of distinct simulator states to execute together; total envs are G*K.", |
| ) |
| parser.add_argument( |
| "--candidate-mode", |
| choices=("structured", "random"), |
| default="structured", |
| help="Use the proposed intervention lattice or a matched random-negative baseline.", |
| ) |
| args = parser.parse_args(argv) |
| summary = generate_maniskill_lattice( |
| ManiSkillLatticeConfig( |
| demo_path=args.demo, |
| output_dir=args.out, |
| num_groups=args.num_groups, |
| group_offset=args.group_offset, |
| k=args.k, |
| horizon=args.horizon, |
| seed=args.seed, |
| shard_size=args.shard_size, |
| env_id=args.env_id, |
| obs_mode=args.obs_mode, |
| image_quality=args.image_quality, |
| control_mode=args.control_mode, |
| sim_backend=args.sim_backend, |
| render_backend=args.render_backend, |
| parallel_branches=args.parallel_branches, |
| state_storage=args.state_storage, |
| state_batch_size=args.state_batch_size, |
| candidate_mode=args.candidate_mode, |
| ) |
| ) |
| print(json.dumps(summary, indent=2, default=str)) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
