{
    "_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
        }
    }
}