| { | |
| "_target_": "stable_worldmodel.wm.PreJEPA", | |
| "history_size": 3, | |
| "num_pred": 1, | |
| "interpolate_pos_encoding": true, | |
| "encoder": { | |
| "_target_": "stable_worldmodel.wm.prejepa.module.create_backbone", | |
| "name": "dinov2_small" | |
| }, | |
| "predictor": { | |
| "_target_": "stable_worldmodel.wm.prejepa.module.CausalPredictor", | |
| "num_patches": 256, | |
| "num_frames": 3, | |
| "dim": 404, | |
| "depth": 6, | |
| "heads": 16, | |
| "mlp_dim": 2048, | |
| "dim_head": 64, | |
| "dropout": 0.1, | |
| "emb_dropout": 0.0 | |
| }, | |
| "extra_encoders": { | |
| "_target_": "torch.nn.ModuleDict", | |
| "modules": { | |
| "proprio": { | |
| "_target_": "stable_worldmodel.wm.prejepa.module.Embedder", | |
| "in_chans": 4, | |
| "emb_dim": 10 | |
| }, | |
| "action": { | |
| "_target_": "stable_worldmodel.wm.prejepa.module.Embedder", | |
| "in_chans": 10, | |
| "emb_dim": 10 | |
| } | |
| } | |
| }, | |
| "pixel_token": "patch" | |
| } |