File size: 1,120 Bytes
7935123
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
{
  "_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_": "module.ProjectorBN",
    "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
    }
  }
}