File size: 4,117 Bytes
9375a59 | 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 57 58 59 60 61 62 63 64 65 66 67 68 69 | """Model definitions for locally trained, explicitly labeled baseline checkpoints."""
import sys
from pathlib import Path
from types import SimpleNamespace
import torch
from torch import nn
ROOT=Path(__file__).resolve().parents[1]
sys.path.insert(0,str(ROOT/'runtime'))
for source in (ROOT/'checkpoints/dino-hub').glob('facebookresearch_dinov2*'):sys.path.insert(0,str(source))
class DinoBackbone(nn.Module):
def __init__(self):
super().__init__()
hub=ROOT/'checkpoints/dino-hub'
torch.hub.set_dir(str(hub))
source=next(hub.glob('facebookresearch_dinov2*'))
self.model=torch.hub.load(str(source),'dinov2_vits14',source='local',pretrained=True)
self.requires_grad_(False)
def forward(self,x,**kwargs):
with torch.no_grad():z=self.model.forward_features(x)['x_norm_patchtokens']
return SimpleNamespace(last_hidden_state=torch.cat([z.new_zeros(z.shape[0],1,z.shape[-1]),z],1))
def build_model(method,action_dim):
from stable_worldmodel.wm.pldm.pldm import PLDM
from stable_worldmodel.wm.pldm.module import Predictor,Embedder,MLP
if method in ['pldm','lejepa']:
from stable_pretraining.backbone.utils import vit_hf
return PLDM(encoder=vit_hf(size='tiny',patch_size=14,image_size=224,pretrained=False,use_mask_token=False),predictor=Predictor(num_frames=3,input_dim=192,hidden_dim=192,output_dim=192,depth=6,heads=16,mlp_dim=2048,dim_head=64,dropout=.1,emb_dropout=0),action_encoder=Embedder(input_dim=5*action_dim,emb_dim=192),projector=MLP(input_dim=192,output_dim=192,hidden_dim=2048,norm_fn=nn.BatchNorm1d),pred_proj=MLP(input_dim=192,output_dim=192,hidden_dim=2048,norm_fn=nn.BatchNorm1d))
from stable_worldmodel.wm.prejepa.prejepa import PreJEPA
from stable_worldmodel.wm.prejepa.module import CausalPredictor,Embedder as PatchEmbedder
return PreJEPA(encoder=DinoBackbone(),predictor=CausalPredictor(num_patches=256,num_frames=3,dim=394,depth=6,heads=16,mlp_dim=2048,dim_head=64,dropout=.1,emb_dropout=0),extra_encoders=nn.ModuleDict({'action':PatchEmbedder(in_chans=action_dim*5,emb_dim=10)}),history_size=3,num_pred=1)
class TrainedCost(nn.Module):
"""Cache visual encoding, then use the trained predictor at every plan step."""
def __init__(self,model,method):
super().__init__();self.model=model;self.method=method;self._info=None
@torch.inference_mode()
def get_cost(self,info,candidates):
assert candidates.shape[0]==1
if self._info is not info:
self._info=info
if self.method=='dinowm':
self.initial=self.model._encode_image(info['pixels'][:,0])
self.goal=self.model._encode_image(info['goal'][:,0])[:,-1:]
else:
self.initial=self.model.encode({'pixels':info['pixels'][:,0]})['emb']
self.goal=self.model.encode({'pixels':info['goal'][:,0]})['emb'][:,-1:]
costs=[]
for actions in candidates[0].split(100 if self.method=='dinowm' else 300):
z=self.initial.expand(actions.shape[0],*self.initial.shape[1:])
if self.method=='dinowm':
# Single initial frame in the common evaluator; each action is a five-step block.
ae=self.model.extra_encoders['action'](actions)
z=torch.cat([z,ae[:,:1,None,:].expand(-1,-1,z.shape[2],-1)],-1)
for t in range(actions.shape[1]):
pred=self.model.predict(z[:,-3:])[:,-1:]
if t+1<actions.shape[1]:
pred=torch.cat([pred[...,:384],ae[:,t+1:t+2,None,:].expand(-1,-1,z.shape[2],-1)],-1)
z=torch.cat([z,pred],1)
error=pred[...,:384]-self.goal
else:
ae=self.model.action_encoder(actions)
for t in range(actions.shape[1]):
pred=self.model.predict(z[:,-3:],ae[:,max(0,t-2):t+1])[:,-1:]
z=torch.cat([z,pred],1)
error=pred-self.goal
costs.append(error.square().flatten(1).mean(1))
return torch.cat(costs)[None]
|