File size: 2,325 Bytes
7fd2dc5 | 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 | """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()
|