File size: 6,001 Bytes
ed7442b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | #!/usr/bin/env python
from __future__ import annotations
import argparse
import json
import sys
from pathlib import Path
from typing import Any
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.data.datasets import CILDataset # noqa: E402
from dovla_cil.eval.lattice_eval import _validation_group_ids # noqa: E402
from dovla_cil.eval.maniskill_policy_rollout import _numeric_action_values # noqa: E402
from dovla_cil.models.dovla import ( # noqa: E402
DoVLAConfig,
DoVLAModel,
load_model_state,
vectorize_toy_observation,
)
def main(argv: list[str] | None = None) -> int:
parser = argparse.ArgumentParser(
description="Export policy BC targets chosen by a trained field on CIL action lattices."
)
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("--split", choices=("train", "val", "all"), default="train")
parser.add_argument("--batch-groups", type=int, default=32)
parser.add_argument(
"--exclude-types",
default="expert",
help="Comma-separated candidate_type values to exclude before field selection.",
)
parser.add_argument("--max-groups", type=int, default=None)
args = parser.parse_args(argv)
if args.batch_groups <= 0:
parser.error("--batch-groups must be positive")
if args.max_groups is not None and args.max_groups <= 0:
parser.error("--max-groups must be positive when provided")
try:
import torch
except ImportError as exc: # pragma: no cover
raise ImportError("export_field_selected_policy_targets.py requires torch") from exc
checkpoint = torch.load(
args.checkpoint,
map_location=_resolve_device(args.device),
weights_only=False,
)
model_config = DoVLAConfig(**checkpoint["model_config"])
if model_config.observation_mode != "state":
raise ValueError("field-selected target export currently supports state observations only")
device = _resolve_device(args.device)
model = DoVLAModel(model_config).to(device)
load_model_state(model, checkpoint)
model.eval()
dataset = CILDataset(args.dataset)
trainer_config = checkpoint.get("trainer_config", {})
val_ids = set(
_validation_group_ids(
dataset.group_ids,
val_fraction=float(trainer_config.get("val_fraction", 0.2)),
seed=int(trainer_config.get("seed", 0)),
)
)
if args.split == "train":
group_ids = [group_id for group_id in dataset.group_ids if group_id not in val_ids]
elif args.split == "val":
group_ids = [group_id for group_id in dataset.group_ids if group_id in val_ids]
else:
group_ids = list(dataset.group_ids)
if args.max_groups is not None:
group_ids = group_ids[: args.max_groups]
excluded = {item.strip() for item in args.exclude_types.split(",") if item.strip()}
targets: dict[str, dict[str, Any]] = {}
counts: dict[str, int] = {}
with torch.no_grad():
for start in range(0, len(group_ids), args.batch_groups):
for group_id in group_ids[start : start + args.batch_groups]:
records = [
record
for record in dataset.get_group(group_id)
if record.candidate_type not in excluded
]
if not records:
records = dataset.get_group(group_id)
if not records:
continue
obs = torch.tensor(
[
vectorize_toy_observation(
records[0].observation_inline or {},
obs_dim=model_config.obs_dim,
)
]
* len(records),
dtype=torch.float32,
device=device,
)
actions = torch.tensor(
[_numeric_action_values(record) for record in records],
dtype=torch.float32,
device=device,
)
field = model.forward_field(
obs,
[record.instruction for record in records],
actions,
)
best_idx = int(torch.argmax(field["potential"].reshape(len(records))).item())
best = records[best_idx]
counts[best.candidate_type] = counts.get(best.candidate_type, 0) + 1
targets[group_id] = {
"record_id": best.record_id,
"candidate_type": best.candidate_type,
"task_id": best.task_id,
"score": float(best.reward.score),
"rank_within_group": best.rank_within_group,
}
payload = {
"checkpoint": str(args.checkpoint),
"dataset": str(args.dataset),
"split": args.split,
"excluded_candidate_types": sorted(excluded),
"num_groups": len(group_ids),
"num_targets": len(targets),
"selected_candidate_type_counts": counts,
"targets": targets,
}
args.out.parent.mkdir(parents=True, exist_ok=True)
args.out.write_text(json.dumps(payload, indent=2) + "\n")
print(json.dumps({k: v for k, v in payload.items() if k != "targets"}, indent=2))
print(f"Wrote {args.out}")
return 0
def _resolve_device(device: str) -> str:
if device != "auto":
return device
try:
import torch
except ImportError: # pragma: no cover
return "cpu"
return "cuda" if torch.cuda.is_available() else "cpu"
if __name__ == "__main__":
raise SystemExit(main())
|