#!/usr/bin/env python 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.eval.maniskill_policy_rollout import ( # noqa: E402 evaluate_maniskill_policy_rollout, ) def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser( description="Execute a DoVLA policy checkpoint from restored ManiSkill CIL states." ) parser.add_argument("--checkpoint", type=Path, required=True) parser.add_argument("--dataset", type=Path, required=True) parser.add_argument("--out", type=Path, required=True) parser.add_argument("--device", default="auto") parser.add_argument("--all-groups", action="store_true") parser.add_argument("--max-groups", type=int, default=None) parser.add_argument("--group-batch-size", type=int, default=16) parser.add_argument( "--sim-backend", default=None, help="Override generation_summary sim_backend, e.g. physx_cuda or physx_cpu.", ) parser.add_argument( "--render-backend", default=None, help="Override generation_summary render_backend. Leave unset for state-only archives.", ) parser.add_argument( "--selection-mode", choices=( "policy", "field", "field_optim", "proposal_lattice", "lattice", "retrieval_lattice", "retrieval_residual", ), default="policy", help="'policy' executes the deterministic policy mean; 'field' scores model-generated " "candidates with the learned interventional field; 'field_optim' additionally " "optimizes model-generated candidates with projected action-space gradient ascent; " "'proposal_lattice' scores the model's typed proposal head; " "'lattice' scores the current state's CIL action lattice without reading rewards; " "'retrieval_lattice' scores the nearest train-state lattice for the current state; " "'retrieval_residual' translates nearest train-state counterfactual residuals around " "the current policy mean before field scoring.", ) parser.add_argument( "--num-candidates", type=int, default=1, help="Candidates per state when selection-mode=field (mean + perturbations).", ) parser.add_argument( "--candidate-sigma", type=float, default=0.2, help="Std-dev of Gaussian perturbations applied to the policy mean for field selection.", ) parser.add_argument( "--selection-seed", type=int, default=0, help="Base RNG seed for candidate perturbations.", ) parser.add_argument( "--selection-margin", type=float, default=0.0, help="Minimum field-potential advantage over candidate 0 required to leave it.", ) parser.add_argument( "--prepend-policy-candidate", action="store_true", help="Prepend the current policy action as candidate 0 for lattice/retrieval_lattice.", ) parser.add_argument( "--proposal-lattice-types", default="", help="Optional comma-separated proposal head candidate types for " "selection-mode=proposal_lattice. Empty uses every type in the checkpoint.", ) parser.add_argument( "--field-optim-steps", type=int, default=0, help="Projected gradient-ascent steps when selection-mode=field_optim.", ) parser.add_argument( "--field-optim-step-size", type=float, default=0.05, help="Action-space step size for selection-mode=field_optim.", ) parser.add_argument( "--field-optim-trust-radius", type=float, default=0.5, help="L-infinity trust radius around the policy proposal for field_optim.", ) parser.add_argument( "--field-optim-l2-penalty", type=float, default=0.0, help="Mean squared action penalty around the policy proposal for field_optim.", ) parser.add_argument( "--retrieval-neighbors", type=int, default=1, help="Nearest train states to use for retrieval_lattice/retrieval_residual proposals.", ) parser.add_argument( "--retrieval-metric", choices=("raw", "zscore", "task_relative"), default="raw", help="State-space metric for retrieval proposals. 'raw' preserves earlier results; " "'zscore' standardizes each task's train-bank features before nearest-neighbor lookup; " "'task_relative' retrieves by target/reference actor pose rather than full robot state.", ) parser.add_argument( "--retrieval-type-min-success", type=float, default=0.0, help="Minimum train-split terminal success rate for a residual candidate family to be " "eligible in retrieval_residual mode. The policy_residual fallback is always kept.", ) parser.add_argument( "--retrieval-type-success-bonus-scale", type=float, default=0.0, help="Scale for adding a train-split task/family terminal-success prior to each " "retrieved residual candidate before field selection.", ) parser.add_argument( "--retrieval-residual-consensus-penalty-scale", type=float, default=0.0, help="Penalty scale for low-consensus train-neighbor residual families. The penalty is " "the local residual dispersion divided by residual energy, using only train actions.", ) parser.add_argument( "--retrieval-residual-min-source-progress", type=float, default=0.0, help="Minimum measured train-source progress for an individual residual candidate to " "be eligible in retrieval_residual mode. The policy_residual fallback is always kept.", ) parser.add_argument( "--retrieval-residual-min-source-advantage", type=float, default=-1.0e9, help="Minimum source reward-score advantage over the source anchor for an individual " "residual candidate. The default disables advantage gating.", ) parser.add_argument( "--retrieval-residual-source-progress-bonus-scale", type=float, default=0.0, help="Scale for adding a train-source progress prior to each retrieved residual " "candidate before field selection. The policy_residual fallback receives zero.", ) parser.add_argument( "--retrieval-residual-source-score-bonus-scale", type=float, default=0.0, help="Scale for adding a train-source reward-score prior to each retrieved residual " "candidate before field selection. Score is progress plus terminal success.", ) parser.add_argument( "--retrieval-residual-source-advantage-bonus-scale", type=float, default=0.0, help="Scale for adding a train-source reward-score advantage prior relative to the " "source residual anchor action.", ) parser.add_argument( "--retrieval-residual-composite-l2-penalty-scale", type=float, default=0.0, help="Penalty scale subtracted from composed residual candidates according to " "their mean squared residual energy.", ) parser.add_argument( "--retrieval-residual-action-l2-penalty", type=float, default=0.0, help="Penalty scale subtracted from residual candidate field scores according to " "mean squared action displacement from the policy mean.", ) parser.add_argument( "--retrieval-residual-scale", type=float, default=1.0, help="Scale applied to train-state residuals before adding them to the policy mean.", ) parser.add_argument( "--retrieval-residual-scales", default="", help="Optional comma-separated scale grid for counterfactual residual ray search. " "When set, this overrides --retrieval-residual-scale during candidate generation.", ) parser.add_argument( "--retrieval-residual-anchor", choices=("expert", "policy"), default="expert", help="Source action used to form train residuals: candidate-anchor.", ) parser.add_argument( "--retrieval-residual-direction", choices=("candidate_minus_anchor", "anchor_minus_candidate"), default="candidate_minus_anchor", help="Direction used to form train residuals. The default preserves existing " "candidate-anchor tangents; 'anchor_minus_candidate' builds corrective repair " "tangents from non-anchor actions back toward the anchor.", ) parser.add_argument( "--retrieval-residual-reduce", choices=( "none", "mean_by_type", "median_by_type", "kernel_mean_by_type", "compose_mean_by_type", "field_softmax", ), default="none", help="Optional consensus reduction over retrieved residuals with the same candidate type. " "'kernel_mean_by_type' weights source residuals by train-state retrieval distance; " "'compose_mean_by_type' also adds pairwise sums of type-consensus tangents; " "'field_softmax' forms a field-weighted tangent barycenter before final scoring.", ) parser.add_argument( "--retrieval-residual-challenger-types", default="", help="Optional comma-separated residual candidate types that may override the " "primary retrieval-residual selection in a second-stage challenger gate.", ) parser.add_argument( "--retrieval-residual-challenger-scales", default="", help="Optional comma-separated residual scales where challenger candidates may " "override the primary retrieval-residual selection. Empty allows all scales.", ) parser.add_argument( "--retrieval-residual-challenger-margin", type=float, default=0.0, help="Required field-potential margin for a challenger candidate to override the " "primary retrieval-residual selection.", ) parser.add_argument( "--retrieval-residual-challenger-type-margins", default="", help="Optional comma-separated candidate_type=margin overrides for challenger " "families. Types without an entry use --retrieval-residual-challenger-margin.", ) parser.add_argument( "--retrieval-residual-challenger-tasks", default="", help="Optional comma-separated task IDs where challenger candidates may override " "the primary retrieval-residual selection. Empty enables challengers for all tasks.", ) parser.add_argument( "--retrieval-residual-challenger-type-tasks", default="", help="Optional semicolon-separated candidate_type=TaskA|TaskB limits for individual " "challenger types. Types without an entry keep the global challenger task behavior.", ) parser.add_argument( "--lattice-exclude-types", default="", help="Comma-separated candidate_type values to mask in selection-mode=lattice, " "for example 'expert'.", ) parser.add_argument( "--lattice-exclude-type-tasks", default="", help="Optional semicolon-separated candidate_type=TaskA|TaskB exclusions applied " "only on listed tasks, e.g. 'wrong_gripper=StackCube-v1'.", ) parser.add_argument( "--candidate-type-bonuses", default="", help="Comma-separated candidate_type=bonus priors added to field potentials before " "selection, e.g. 'residual_no_op=0.05'. Empty preserves previous behavior.", ) parser.add_argument( "--candidate-type-bonus-components", action="store_true", help="Let composite candidate types inherit the sum of configured component bonuses " "unless an exact composite bonus is configured.", ) parser.add_argument( "--candidate-oracle-rollouts", type=int, default=0, help="Diagnostic only: execute the selected candidate plus the top field-ranked " "residual candidates and report the best measured outcome in that prefix. " "Zero disables the diagnostic and preserves deployment evaluation.", ) parser.add_argument( "--candidate-oracle-unique-tolerance", type=float, default=1.0e-6, help="Max absolute action difference for treating two diagnostic oracle candidates " "as duplicates. Duplicate candidates are rolled only as invalid padding and cannot " "win the oracle choice.", ) args = parser.parse_args(argv) lattice_exclude_types = tuple( item.strip() for item in args.lattice_exclude_types.split(",") if item.strip() ) lattice_exclude_type_tasks: dict[str, tuple[str, ...]] = {} try: for item in args.lattice_exclude_type_tasks.split(";"): item = item.strip() if not item: continue key, value = item.split("=", 1) lattice_exclude_type_tasks[key.strip()] = tuple( task.strip() for task in value.split("|") if task.strip() ) except ValueError as exc: raise SystemExit( "--lattice-exclude-type-tasks must contain semicolon-separated " "candidate_type=TaskA|TaskB entries" ) from exc retrieval_residual_challenger_types = tuple( item.strip() for item in args.retrieval_residual_challenger_types.split(",") if item.strip() ) retrieval_residual_challenger_tasks = tuple( item.strip() for item in args.retrieval_residual_challenger_tasks.split(",") if item.strip() ) retrieval_residual_challenger_type_margins: dict[str, float] = {} try: for item in args.retrieval_residual_challenger_type_margins.split(","): item = item.strip() if not item: continue key, value = item.split("=", 1) retrieval_residual_challenger_type_margins[key.strip()] = float( value.strip() ) except ValueError as exc: raise SystemExit( "--retrieval-residual-challenger-type-margins must contain " "comma-separated candidate_type=margin entries" ) from exc retrieval_residual_challenger_type_tasks: dict[str, tuple[str, ...]] = {} try: for item in args.retrieval_residual_challenger_type_tasks.split(";"): item = item.strip() if not item: continue key, value = item.split("=", 1) retrieval_residual_challenger_type_tasks[key.strip()] = tuple( task.strip() for task in value.split("|") if task.strip() ) except ValueError as exc: raise SystemExit( "--retrieval-residual-challenger-type-tasks must contain " "semicolon-separated candidate_type=TaskA|TaskB entries" ) from exc proposal_lattice_types = tuple( item.strip() for item in args.proposal_lattice_types.split(",") if item.strip() ) try: retrieval_residual_challenger_scales = tuple( float(item.strip()) for item in args.retrieval_residual_challenger_scales.split(",") if item.strip() ) except ValueError as exc: raise SystemExit( "--retrieval-residual-challenger-scales must contain comma-separated floats" ) from exc candidate_type_bonuses: dict[str, float] = {} try: for item in args.candidate_type_bonuses.split(","): item = item.strip() if not item: continue key, value = item.split("=", 1) candidate_type_bonuses[key.strip()] = float(value.strip()) except ValueError as exc: raise SystemExit( "--candidate-type-bonuses must contain comma-separated candidate_type=bonus pairs" ) from exc try: retrieval_residual_scales = tuple( float(item.strip()) for item in args.retrieval_residual_scales.split(",") if item.strip() ) except ValueError as exc: raise SystemExit("--retrieval-residual-scales must contain comma-separated floats") from exc result = evaluate_maniskill_policy_rollout( args.checkpoint, args.dataset, output_path=args.out, device=args.device, all_groups=args.all_groups, max_groups=args.max_groups, group_batch_size=args.group_batch_size, sim_backend=args.sim_backend, render_backend=args.render_backend, selection_mode=args.selection_mode, num_candidates=args.num_candidates, candidate_sigma=args.candidate_sigma, selection_seed=args.selection_seed, selection_margin=args.selection_margin, prepend_policy_candidate=args.prepend_policy_candidate, proposal_lattice_types=proposal_lattice_types, field_optim_steps=args.field_optim_steps, field_optim_step_size=args.field_optim_step_size, field_optim_trust_radius=args.field_optim_trust_radius, field_optim_l2_penalty=args.field_optim_l2_penalty, retrieval_neighbors=args.retrieval_neighbors, retrieval_metric=args.retrieval_metric, retrieval_type_min_success=args.retrieval_type_min_success, retrieval_type_success_bonus_scale=args.retrieval_type_success_bonus_scale, retrieval_residual_consensus_penalty_scale=( args.retrieval_residual_consensus_penalty_scale ), retrieval_residual_min_source_progress=args.retrieval_residual_min_source_progress, retrieval_residual_min_source_advantage=args.retrieval_residual_min_source_advantage, retrieval_residual_source_progress_bonus_scale=( args.retrieval_residual_source_progress_bonus_scale ), retrieval_residual_source_score_bonus_scale=( args.retrieval_residual_source_score_bonus_scale ), retrieval_residual_source_advantage_bonus_scale=( args.retrieval_residual_source_advantage_bonus_scale ), retrieval_residual_composite_l2_penalty_scale=( args.retrieval_residual_composite_l2_penalty_scale ), retrieval_residual_action_l2_penalty=args.retrieval_residual_action_l2_penalty, retrieval_residual_scale=args.retrieval_residual_scale, retrieval_residual_scales=retrieval_residual_scales, retrieval_residual_anchor=args.retrieval_residual_anchor, retrieval_residual_direction=args.retrieval_residual_direction, retrieval_residual_reduce=args.retrieval_residual_reduce, retrieval_residual_challenger_types=retrieval_residual_challenger_types, retrieval_residual_challenger_scales=retrieval_residual_challenger_scales, retrieval_residual_challenger_margin=args.retrieval_residual_challenger_margin, retrieval_residual_challenger_type_margins=( retrieval_residual_challenger_type_margins ), retrieval_residual_challenger_tasks=retrieval_residual_challenger_tasks, retrieval_residual_challenger_type_tasks=( retrieval_residual_challenger_type_tasks ), lattice_exclude_types=lattice_exclude_types, lattice_exclude_type_tasks=lattice_exclude_type_tasks, candidate_type_bonuses=candidate_type_bonuses, candidate_type_bonus_components=args.candidate_type_bonus_components, candidate_oracle_rollouts=args.candidate_oracle_rollouts, candidate_oracle_unique_tolerance=args.candidate_oracle_unique_tolerance, ) print(json.dumps({key: value for key, value in result.items() if key != "rows"}, indent=2)) return 0 if __name__ == "__main__": raise SystemExit(main())