Initial pi0.7 minimal validation

src/pi07mini/infer.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+import argparse
+from pathlib import Path
+import matplotlib.pyplot as plt
+import numpy as np,torch
+from .data import Pi07SteerableDataset
+from .flow import FlowMatching
+from .model import Pi07Config,Pi07Mini,HighLevelPolicy,TinyGoalGenerator
+from .runtime import load_config,choose_device,write_json
+from .tokenizer import MiniTokenizer,EpisodeMetadata
+
+@torch.no_grad()
+def evaluate(checkpoint_path):
+ ck=torch.load(checkpoint_path,map_location='cpu',weights_only=False);cfg=ck['config'];root=Path(__file__).resolve().parents[2];dev=choose_device(cfg['device']);tok=MiniTokenizer();ds=Pi07SteerableDataset(root/'data/overfit/overfit_manifest.jsonl',image_size=cfg['data']['image_size'],history_frames=cfg['data']['history_frames'],action_horizon=cfg['data']['action_horizon'],action_statistics=ck['action_statistics'])
+ pcfg=Pi07Config(**ck['policy_config']);policy=Pi07Mini(pcfg);policy.load_state_dict(ck['policy']);policy.to(dev).eval();high=HighLevelPolicy(pcfg);high.load_state_dict(ck['high_level']);high.to(dev).eval();goal=TinyGoalGenerator(tok.vocab_size);goal.load_state_dict(ck['goal']);goal.to(dev).eval();flow=FlowMatching(cfg['flow']['integration_steps']);s=next(ds[i] for i in range(len(ds)) if ds[i]['policy_kind']=='expert');meta=EpisodeMetadata(**s['token_metadata']);image=s['image'][None].to(dev)
+ prompt,_=tok.encode_prompt('drive safely',None,meta,include_subtask=False);hl=high(image,prompt[None].to(dev),s['history_images'][None].to(dev),s['history_states'][None].to(dev),s['current_state'][None].to(dev));pred_subtask=tok.subtasks.decode(int(hl.logits.argmax(-1).cpu()))
+ full=tok.build_policy_sequence('drive safely',s['reasoning'],s['actions'],meta);ids=full.input_ids[:full.action_start_index][None].to(dev)
+ torch.manual_seed(cfg['seed']);g=torch.randn_like(s['subgoal_pixel_values'][None]).to(dev);grid=torch.linspace(0,1,6,device=dev)
+ for i in range(5): g=g+(grid[i+1]-grid[i])*goal(image,ids,g,grid[i].expand(1)).vector_field
+ generated_goal=g.cpu()[0]
+ def sample(name,include_meta=True,use_goal=True,use_history=True,cfg_beta=None,generated=False):
+  seq=tok.build_policy_sequence('drive safely',s['reasoning'],s['actions'],meta,include_episode_metadata=include_meta);xids=seq.input_ids[:seq.action_start_index][None].to(dev)
+  goal_pixels=(generated_goal if generated else s['subgoal_pixel_values'])[None].to(dev) if use_goal else None
+  common=dict(pixel_values=image,history_pixel_values=s['history_images'][None].to(dev) if use_history else None,history_states=s['history_states'][None].to(dev) if use_history else None,current_state=s['current_state'][None].to(dev),subgoal_pixel_values=goal_pixels)
+  def field(x,t):
+   if cfg_beta is None:return policy(input_ids=xids,noisy_actions=x,timesteps=t,**common).vector_field
+   un=tok.build_policy_sequence('drive safely',s['reasoning'],s['actions'],meta,include_episode_metadata=False);uids=un.input_ids[:un.action_start_index][None].to(dev)
+   return policy.metadata_cfg(xids,uids,guidance_scale=cfg_beta,noisy_actions=x,timesteps=t,**common)
+  torch.manual_seed(cfg['evaluation']['fixed_noise_seed']);a=flow.sample(field,(1,ds.action_horizon,2),dev,cfg['flow']['integration_steps'])[0].cpu();return ds.action_normalizer.decode(a,s['history_xy']).numpy()
+ curves={'full_oracle_goal':sample('full'),'generated_goal':sample('generated',generated=True),'no_metadata':sample('no_metadata',False),'no_goal':sample('no_goal',use_goal=False),'no_history':sample('no_history',use_history=False),'cfg_beta2':sample('cfg',cfg_beta=2.0)}
+ target=ds.action_normalizer.decode(s['rtc_target_actions'],s['history_xy']).numpy();valid=s['rtc_valid_mask'].numpy().astype(bool);metrics={k:float(np.mean((v[valid]-target[valid])**2)) for k,v in curves.items()}
+ out=Path(checkpoint_path).parent;fig,ax=plt.subplots(figsize=(9,6));ax.plot(target[:,0],target[:,1],'o-',label='expert target',linewidth=2)
+ for k,v in curves.items():ax.plot(v[:,0],v[:,1],'x--',label=f'{k} ({metrics[k]:.3f})')
+ ax.axis('equal');ax.grid(True,alpha=.3);ax.legend(fontsize=8);ax.set_xlabel('forward x (m)');ax.set_ylabel('left y (m)');ax.set_title('π0.7 context steering and ablations');fig.tight_layout();traj=out/'context_ablation.png';fig.savefig(traj,dpi=170);plt.close(fig)
+ fig,axes=plt.subplots(1,3,figsize=(10,3));axes[0].imshow(image[0].cpu().permute(1,2,0));axes[0].set_title('current observation');axes[1].imshow(s['subgoal_pixel_values'].permute(1,2,0));axes[1].set_title('endpoint goal target');axes[2].imshow(generated_goal.permute(1,2,0).clamp(0,1));axes[2].set_title('generated goal');
+ for a in axes:a.axis('off')
+ fig.tight_layout();goalpath=out/'goal_generation.png';fig.savefig(goalpath,dpi=170);plt.close(fig);write_json(out/'evaluation.json',{'predicted_subtask':pred_subtask,'trajectory_metrics':metrics});print({'predicted_subtask':pred_subtask,'trajectory_metrics':metrics});return traj,goalpath
+def main():
+ p=argparse.ArgumentParser();p.add_argument('--checkpoint',default='outputs/overfit/checkpoint.pt');evaluate(p.parse_args().checkpoint)
+if __name__=='__main__':main()