Hugging Face's logo

SaiResearch
/
sai_wm

Model card Files
xet
 Community
1
sai_wm
File size: 3,488 Bytes
c02d272
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
{
  "name": "choppercommand_iris_world_model",
  "env": "ChopperCommandNoFrameskip-v4",
  "model_type": "iris",
  "metadata": {
    "latent_dim": [1, 16],
    "num_tokens": 340
  },
  "util_folders":{
    "models": "../../src/models"
  },
  "requirements":{
    "-r": "requirements.txt"
  },
  "models": [
    {
      "name": "world_model",
      "framework": null,
      "format": "state_dict",
      "source": {
        "weights_path": "world_model.pt",
        "class_path": "../../src/world_model.py",
        "class_name": "WorldModel",
        "class_args": [
        {
          "vocab_size": 512,
          "act_vocab_size": 18,
          "tokens_per_block": 17,
          "max_blocks": 20,
          "attention": "causal",
          "num_layers": 10,
          "num_heads": 4,
          "embed_dim": 256,
          "embed_pdrop": 0.1,
          "resid_pdrop": 0.1,
          "attn_pdrop": 0.1
        }]
      },
      "signature": {
        "inputs": ["tokens", "past_keys_values"],
        "call_mode": "positional"
      },
      "sub_models":
      [
        {
          "name": "transformer",
          "sub_model_name": "transformer",
          "signature": 
          {
            "inputs": ["sequences", "past_keys_values"],
            "call_mode": "positional"
          }
        }
      ],
      "methods":
      [
        {
          "name": "generate_empty_keys_values",
          "method_name": "generate_empty_keys_values",
          "signature": 
          {
            "inputs": ["n"]
          }
        }
      ]
    },
    {
      "name": "tokenizer",
      "framework": null,
      "format": "state_dict",
      "source": {
        "weights_path": "tokenizer.pt",
        "class_path": "../../src/tokenizer.py",
        "class_name": "Tokenizer",
        "class_args": [{
          "vocab_size": 512,
          "embed_dim": 512,
          "encoder": {
            "resolution": 64,
            "in_channels": 3,
            "z_channels": 512,
            "ch": 64,
            "ch_mult": [1, 1, 1, 1, 1],
            "num_res_blocks": 2,
            "attn_resolutions": [8, 16],
            "out_ch": 3,
            "dropout": 0.0
          },
          "decoder": { 
            "resolution": 64,
            "in_channels": 3,
            "z_channels": 512,
            "ch": 64,
            "ch_mult": [1, 1, 1, 1, 1],
            "num_res_blocks": 2,
            "attn_resolutions": [8, 16],
            "out_ch": 3,
            "dropout": 0.0
          }
        }]
      },
      "signature": {
        "inputs": ["x", "should_preprocess", "should_postprocess"],
        "call_mode": "positional"
      },
      "sub_models":
      [
        {
          "name": "embedding",
          "sub_model_name": "embedding",
          "signature": 
          {
            "call_mode": "auto"
          }
        }
      ],
      "methods":
      [
        {
          "name": "decode",
          "method_name": "decode",
          "signature": 
          {
            "inputs": ["z", "should_postprocess"]
          }
        },
        {
          "name": "decode_obs_tokens",
          "method_name": "decode_obs_tokens",
          "signature": 
          {
            "inputs": ["obs_tokens", "num_observations_tokens"]
          }
        },
        {
          "name": "encode",
          "method_name": "encode",
          "signature": 
          {
            "inputs": ["observations", "should_preprocess"]
          }
        }
      ]
    }
  ]
}