| { | |
| "_target_": "stable_worldmodel.wm.lewm.LeWM", | |
| "encoder": { | |
| "_target_": "stable_pretraining.backbone.utils.vit_hf", | |
| "size": "tiny", | |
| "patch_size": 14, | |
| "image_size": 224, | |
| "pretrained": false, | |
| "use_mask_token": false | |
| }, | |
| "predictor": { | |
| "_target_": "stable_worldmodel.wm.lewm.module.Predictor", | |
| "num_frames": 3, | |
| "input_dim": 192, | |
| "hidden_dim": 192, | |
| "output_dim": 192, | |
| "depth": 6, | |
| "heads": 16, | |
| "mlp_dim": 2048, | |
| "dim_head": 64, | |
| "dropout": 0.1, | |
| "emb_dropout": 0.0 | |
| }, | |
| "action_encoder": { | |
| "_target_": "stable_worldmodel.wm.lewm.module.Embedder", | |
| "input_dim": 25, | |
| "emb_dim": 192 | |
| }, | |
| "projector": { | |
| "_target_": "stable_worldmodel.wm.lewm.module.MLP", | |
| "input_dim": 192, | |
| "output_dim": 192, | |
| "hidden_dim": 2048, | |
| "norm_fn": { | |
| "_target_": "torch.nn.BatchNorm1d", | |
| "_partial_": true | |
| } | |
| }, | |
| "pred_proj": { | |
| "_target_": "stable_worldmodel.wm.lewm.module.MLP", | |
| "input_dim": 192, | |
| "output_dim": 192, | |
| "hidden_dim": 2048, | |
| "norm_fn": { | |
| "_target_": "torch.nn.BatchNorm1d", | |
| "_partial_": true | |
| } | |
| } | |
| } | |