scripts/eval_maniskill_policy_rollout.py

#!/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())