| """Original released DINO-WM PushT model with the shared LeWM evaluator. |
| |
| Keep common CEM coordinates; translate inputs to the checkpoint's native |
| normalization. Chunk candidates only for memory, without changing their count. |
| """ |
| import sys,types |
| from pathlib import Path |
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
| ROOT=Path(__file__).resolve().parents[1] |
|
|
| def sdpa_forward(self,x): |
| b,t,_=x.shape |
| q,k,v=[a.reshape(b,t,self.heads,-1).transpose(1,2) for a in self.to_qkv(self.norm(x)).chunk(3,dim=-1)] |
| mask=self.bias[:,:,:t,:t].to(device=x.device,dtype=torch.bool) |
| y=F.scaled_dot_product_attention(q,k,v,attn_mask=mask,dropout_p=0.0) |
| return self.to_out(y.transpose(1,2).reshape(b,t,-1)) |
|
|
| class DinoAdapter(nn.Module): |
| def __init__(self,path,process,chunk_size=300): |
| super().__init__() |
| sys.path.insert(0,str(ROOT/'repos/dino_wm')) |
| from models.dino import DinoV2Encoder |
| from models.visual_world_model import VWorldModel |
| from models.vit import Attention |
| torch.hub.set_dir(str(ROOT/'checkpoints/dino-hub')) |
| payload=torch.load(path,map_location='cpu',weights_only=False) |
| print('DINO checkpoint keys:',list(payload),flush=True) |
| encoder=DinoV2Encoder('dinov2_vits14','x_norm_patchtokens') |
| self.wm=VWorldModel(image_size=224,num_hist=3,num_pred=1,encoder=encoder, |
| proprio_encoder=payload['proprio_encoder'],action_encoder=payload['action_encoder'], |
| decoder=None,predictor=payload['predictor'],proprio_dim=10,action_dim=10, |
| concat_dim=1,num_action_repeat=1,num_proprio_repeat=1, |
| train_encoder=False,train_predictor=False,train_decoder=False) |
| self.wm.eval();self.wm.cuda();self.wm.requires_grad_(False) |
| |
| for layer in self.wm.predictor.modules(): |
| if isinstance(layer,Attention): |
| layer.bias=layer.bias.cuda() |
| if not hasattr(self,'attention_max_error'): |
| with torch.no_grad(): |
| x=torch.randn(1,32,layer.norm.normalized_shape[0],device='cuda') |
| reference=layer(x);replacement=sdpa_forward(layer,x) |
| torch.testing.assert_close(reference,replacement,rtol=2e-4,atol=2e-5) |
| self.attention_max_error=float((reference-replacement).abs().max()) |
| layer.forward=types.MethodType(sdpa_forward,layer) |
| for k in ['action','proprio']: |
| self.register_buffer(k+'_mean',torch.tensor(process[k].mean_,dtype=torch.float32,device='cuda')) |
| self.register_buffer(k+'_std',torch.tensor(process[k].scale_,dtype=torch.float32,device='cuda')) |
| self.register_buffer('native_action_mean',torch.tensor([-.0087,.0068],device='cuda')) |
| self.register_buffer('native_action_std',torch.tensor([.2019,.2002],device='cuda')) |
| self.register_buffer('native_proprio_mean',torch.tensor([236.6155,264.5674,-2.93032027,2.54307914],device='cuda')) |
| self.register_buffer('native_proprio_std',torch.tensor([101.1202,87.0112,74.84556075,74.14009094],device='cuda')) |
| self.chunk_size=chunk_size;self._info=None |
|
|
| def _obs(self,info,goal=False): |
| prop=info['goal_proprio' if goal else 'proprio'][:,0].cuda() |
| prop=(prop*self.proprio_std+self.proprio_mean-self.native_proprio_mean)/self.native_proprio_std |
| return {'visual':info['goal' if goal else 'pixels'][:,0].cuda(),'proprio':prop} |
|
|
| def latent_rollout(self,encoded,actions): |
| b,t,_=actions.shape |
| visual=encoded['visual'].expand(b,-1,-1,-1) |
| proprio=encoded['proprio'].expand(b,-1,-1) |
| ae=self.wm.encode_act(actions[:,:1]) |
| patches=visual.shape[2] |
| z=torch.cat([visual,proprio[:,:,None].expand(-1,-1,patches,-1),ae[:,:,None].expand(-1,-1,patches,-1)],dim=-1) |
| for step in range(1,t): |
| new=self.wm.predict(z[:,-3:])[:,-1:] |
| new=self.wm.replace_actions_from_z(new,actions[:,step:step+1]) |
| z=torch.cat([z,new],dim=1) |
| pred=self.wm.predict(z[:,-3:])[:,-1:] |
| return self.wm.separate_emb(pred)[0] |
|
|
| @torch.inference_mode() |
| def get_cost(self,info,candidates): |
| assert candidates.shape[0]==1,'The shared solver must use batch_size=1' |
| if self._info is not info: |
| self._info=info |
| self._start=self.wm.encode_obs(self._obs(info)) |
| self._goal=self.wm.encode_obs(self._obs(info,True)) |
| actions=candidates[0].reshape(-1,5,5,2) |
| actions=(actions*self.action_std+self.action_mean-self.native_action_mean)/self.native_action_std |
| actions=actions.flatten(-2) |
| costs=[] |
| for chunk in actions.split(self.chunk_size): |
| pred=self.latent_rollout(self._start,chunk) |
| costs.append(sum((pred[k]-self._goal[k]).square().flatten(1).mean(1) for k in ['visual','proprio'])) |
| return torch.cat(costs)[None] |
|
|
