| |
| 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 ( |
| 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()) |
|
|
