| """Apply the released adapter to one matched offline record on CPU. |
| |
| Install the BAS-VLA code repository and its dependencies first. The input action |
| contract is recorded in adapter/adapter_metadata.json. Use that exact contract |
| when supplying new actions from a carrier. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
|
|
| from bas_vla.breaking.residual_adapter import ResidualAdapter, encode_instruction_bow |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--artifact-dir", type=Path, default=Path(__file__).resolve().parent) |
| parser.add_argument("--record-index", type=int, default=160, help="Default: first held-out test record.") |
| args = parser.parse_args() |
| root = args.artifact_dir |
| meta = json.loads((root / "adapter/adapter_metadata.json").read_text()) |
| vocab = json.loads((root / "adapter/vocab.json").read_text()) |
| records = json.loads((root / "data/triplets.json").read_text()) |
| if not 0 <= args.record_index < len(records): |
| parser.error("record-index is outside the released records") |
| record = records[args.record_index] |
| model = ResidualAdapter( |
| input_dim=meta["input_dim"], hidden_dim=meta["hidden_dim"], |
| output_dim=meta["output_dim"], delta_scale=meta["delta_scale"], |
| ) |
| weights = torch.load(root / "adapter/residual_adapter.pt", map_location="cpu", weights_only=True) |
| model.load_state_dict(weights, strict=True) |
| model.eval() |
| outputs = {} |
| with torch.no_grad(): |
| for variant in ("clean", "control", "break"): |
| action = np.asarray(record[f"{variant}_base_action"], dtype=np.float32) |
| features = encode_instruction_bow(record[f"{variant}_instruction"], vocab) |
| calibrated = model(torch.from_numpy(action[None]), torch.from_numpy(features[None]))[0] |
| outputs[variant] = { |
| "instruction": record[f"{variant}_instruction"], |
| "base_action": action.tolist(), "calibrated_action": calibrated.tolist(), |
| } |
| print(json.dumps({ |
| "record_id": record["record_id"], "split": record["split"], |
| "action_space": meta["action_space"], "outputs": outputs, |
| }, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
