| """Online evaluation of the fine-tuned OpenVLA policy at 224x224 in the eval env.""" |
| import argparse, os, json, time, statistics |
| import numpy as np, torch, gymnasium as gym, ogbench, mujoco |
| from PIL import Image |
| from transformers import AutoModelForVision2Seq, AutoProcessor |
| from peft import PeftModel |
| from common import wilson_ci |
| from vla_finetune import ActionHead, prompt_for |
|
|
| MID = "openvla/openvla-7b" |
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--run", required=True) |
| ap.add_argument("--env", default="visual-cube-single-v0") |
| ap.add_argument("--rollouts", type=int, default=20) |
| ap.add_argument("--tasks", type=int, default=5) |
| ap.add_argument("--base_seed", type=int, default=20260825) |
| a = ap.parse_args() |
| dev = "cuda:0" |
| proc = AutoProcessor.from_pretrained(MID, trust_remote_code=True) |
| vla = AutoModelForVision2Seq.from_pretrained( |
| MID, attn_implementation="sdpa", torch_dtype=torch.bfloat16, |
| low_cpu_mem_usage=True, trust_remote_code=True).to(dev) |
| vla = PeftModel.from_pretrained(vla, os.path.join(a.run, "lora")).eval() |
| ck = torch.load(os.path.join(a.run, "head.pt"), map_location=dev, weights_only=False) |
| head = ActionHead(act=5).to(dev).float(); head.load_state_dict(ck["head"]); head.eval() |
|
|
| env = gym.make(a.env); u = env.unwrapped |
| env.reset(seed=0, options=dict(task_id=1, render_goal=False)) |
| u.model.vis.global_.offwidth = 224; u.model.vis.global_.offheight = 224 |
| u._renderer = mujoco.Renderer(u.model, height=224, width=224) |
| u._render_height = 224; u._render_width = 224 |
|
|
| def act_fn(gx, gy): |
| img = Image.fromarray(np.asarray(u.get_pixel_observation())) |
| inp = proc(prompt_for(gx, gy), img).to(dev, dtype=torch.bfloat16) |
| with torch.no_grad(): |
| out = vla(input_ids=inp["input_ids"], attention_mask=inp["attention_mask"], |
| pixel_values=inp["pixel_values"], output_hidden_states=True) |
| h = out.hidden_states[-1][:, -1, :].float() |
| return head(h)[0].cpu().numpy() |
|
|
| per_task, all_k, all_n, steps_s = {}, 0, 0, [] |
| t0 = time.time() |
| for task_id in range(1, a.tasks + 1): |
| k = 0 |
| for r in range(a.rollouts): |
| ob, info = env.reset(seed=a.base_seed + 1000 * task_id + r, |
| options=dict(task_id=task_id, render_goal=False)) |
| gxyz = np.asarray(u.task_infos[task_id - 1]["goal_xyzs"])[0] |
| gx, gy = float(gxyz[0]), float(gxyz[1]) |
| succ, sstep = 0, None |
| for t in range(1000): |
| ob, rew, term, trunc, info = env.step(np.clip(act_fn(gx, gy), -1, 1)) |
| if info.get("success"): |
| succ = 1 |
| if sstep is None: sstep = t + 1 |
| if term or trunc: break |
| k += succ; all_n += 1 |
| if succ: steps_s.append(sstep) |
| p, lo, hi = wilson_ci(k, a.rollouts) |
| per_task[task_id] = dict(k=k, n=a.rollouts, rate=p, ci=[lo, hi]) |
| all_k += k |
| print(f" task{task_id}: {k}/{a.rollouts} = {p:.3f} [{time.time()-t0:.0f}s]", flush=True) |
| p, lo, hi = wilson_ci(all_k, all_n) |
| med = statistics.median(steps_s) if steps_s else None |
| res = dict(model="openvla-7b-oft-style", env=a.env, rollouts_per_task=a.rollouts, |
| overall=dict(k=all_k, n=all_n, rate=p, ci=[lo, hi]), per_task=per_task, |
| median_steps_to_success=med, timeout_rate=1 - all_k / all_n, |
| eval_seconds=time.time() - t0) |
| print("OVERALL: %d/%d = %.3f 95%% Wilson CI [%.3f, %.3f] | median steps=%s | %.0fs" |
| % (all_k, all_n, p, lo, hi, med, time.time() - t0)) |
| json.dump(res, open(os.path.join(a.run, "eval.json"), "w"), indent=2) |
|
|
| if __name__ == "__main__": |
| main() |
|
|
