File size: 1,362 Bytes
ee3f4b5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
{
    "_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": 10,
        "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
        }
    }
}