| |
| """Decisive offline check: feed the model the CORRECT-ORDER state (from the |
| lerobot parquet) + the demo images, and compare predicted action to the recorded |
| ground-truth action (also lerobot order). If pred ~= gt, the model + preprocessing |
| + normalization are correct and the ONLY bug was joint ordering in our live scripts. |
| """ |
| import io, json, sys |
| import numpy as np |
| import go1_env |
| import types as _t |
| if "decord" not in sys.modules: |
| _f=_t.ModuleType("decord"); _f.VideoReader=type("VideoReader",(),{}); _f.cpu=lambda *a,**k:None |
| _f.bridge=_t.ModuleType("decord.bridge"); _f.bridge.set_bridge=lambda *a,**k:None |
| sys.modules["decord"]=_f; sys.modules["decord.bridge"]=_f.bridge |
| import torch |
| from PIL import Image |
|
|
| from go1_env import CKPT, STATS |
| PROMPT="What action should the robot take to reach into the bin, attach both suction cups to the part and lift it out level?" |
| CAM=["cam_head_color","cam_hand_l_color","cam_hand_r_color"] |
| DEV="cuda" |
| with open(STATS) as f: st=json.load(f) |
| sm=np.array(st["state"]["mean"],np.float32); ss=np.array(st["state"]["std"],np.float32) |
| am=np.array(st["action"]["mean"],np.float32); asd=np.array(st["action"]["std"],np.float32) |
|
|
| print("loading model...") |
| from go1.internvl.model.go1 import GO1Model, GO1ModelConfig |
| from go1.internvl.train.dataset import build_transform, dynamic_preprocess, preprocess_internvl2_5 |
| from transformers import AutoTokenizer |
| cfg=GO1ModelConfig.from_pretrained(CKPT,torch_dtype=torch.bfloat16,low_cpu_mem_usage=True) |
| model=GO1Model.from_pretrained(CKPT,config=cfg).to(torch.bfloat16).to(DEV).eval() |
| tok=AutoTokenizer.from_pretrained(CKPT,trust_remote_code=True,use_fast=False,add_eos_token=False) |
| tf=build_transform(is_train=False,input_size=cfg.force_image_size,pad2square=cfg.pad2square) |
| nit=int((cfg.force_image_size//cfg.vision_config.patch_size)**2*(cfg.downsample_ratio**2)) |
| torch.cuda.synchronize(); print("ready.\n") |
|
|
| def decode(a): return Image.open(io.BytesIO(bytes(a.tobytes()))).convert("RGB") |
|
|
| def predict(state, images): |
| pp,nt=[],[] |
| for im in images: |
| tiles=[im] if not cfg.dynamic_image_size else dynamic_preprocess(im,min_num=cfg.min_dynamic_patch, |
| max_num=cfg.max_dynamic_patch,image_size=cfg.force_image_size,use_thumbnail=cfg.use_thumbnail) |
| pp+=tiles; nt.append(len(tiles)) |
| pv=torch.stack([tf(x) for x in pp]); npatch=pv.size(0) |
| conv=[{"from":"human","value":("<image>"*len(images))+PROMPT},{"from":"gpt","value":""}] |
| rp=preprocess_internvl2_5("internvl2_5",[conv],tok,[nit*n for n in nt],num_image=len(images),group_by_length=True) |
| pos=rp["attention_mask"].long().cumsum(-1)-1; pos.masked_fill_(rp["attention_mask"]==0,1) |
| sn=(state-sm)/(ss+1e-6) |
| inp=dict(input_ids=rp["input_ids"][0].cuda().unsqueeze(0),attention_mask=rp["attention_mask"][0].cuda().unsqueeze(0), |
| position_ids=pos[0].cuda().unsqueeze(0),pixel_values=pv.to(torch.bfloat16).cuda(), |
| image_flags=torch.tensor([1]*npatch,dtype=torch.long).cuda(), |
| state=torch.from_numpy(sn).to(torch.bfloat16).cuda().unsqueeze(0).unsqueeze(0), |
| ctrl_freqs=torch.tensor([30.0],dtype=torch.bfloat16,device="cuda").unsqueeze(0)) |
| with torch.no_grad(): out=model(**inp) |
| torch.cuda.synchronize() |
| return out[1][0].float().cpu().numpy()*asd+am |
|
|
| d=np.load("/tmp/demo_frames2.npz") |
| LBL=["L1","L2","L3","L4","L5","L6","L7","R1","R2","R3","R4","R5","R6","R7","Lsuc","Rsuc","b1","b2","b3","b4","b5","hpit"] |
| print("%-6s %-10s %-11s %-10s %s"%("frame","|p0-gt0|","|p0-state|","chunkMAE","worst dim")) |
| for fr in d["frames"].tolist(): |
| state=d["state_%d"%fr].astype(np.float32) |
| images=[decode(d["img_%d_%s"%(fr,k)]) for k in CAM] |
| gt=d["act_%d"%fr].astype(np.float32) |
| pred=predict(state,images)[:gt.shape[0]] |
| d0=np.abs(pred[0]-gt[0]); d0s=np.abs(pred[0]-state) |
| mae=np.abs(pred-gt).mean(); wj=int(d0.argmax()) |
| print("%-6d %-10.3f %-11.3f %-10.3f %s(%.2f)"%(fr,d0.max(),d0s.max(),mae,LBL[wj],d0[wj])) |
| if fr==0: |
| print(" pred[0]=",[round(x,2) for x in pred[0].tolist()]) |
| print(" gt[0] =",[round(x,2) for x in gt[0].tolist()]) |
| print("\n|p0-gt0| small => model CORRECT with right order. |p0-state| small => action[0]~current (no jump).") |
|
|
