File size: 5,460 Bytes
2195144
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
{
    "task": {
        "dataset": {
            "_target_": "datasets.droid.make_droid_rlds_dataset",
            "dataset_id": "droid_rlds",
            "data_dir": "/opt/ml/input/data/training",
            "batch_size": 200,
            "shape_meta": {
                "obs": {
                    "image": {
                        "shape": [
                            180,
                            320,
                            3
                        ],
                        "type": "rgb"
                    },
                    "wrist_image": {
                        "shape": [
                            180,
                            320,
                            3
                        ],
                        "type": "rgb"
                    },
                    "joint_position": {
                        "shape": [
                            7
                        ],
                        "type": "low_dim"
                    },
                    "gripper_position": {
                        "shape": [
                            1
                        ],
                        "type": "low_dim"
                    }
                },
                "action": {
                    "shape": [
                        8
                    ]
                }
            },
            "obs_stack": 1,
            "action_chunk": 16,
            "shuffle_buffer_size": 100000,
            "filter_dict_path": "/opt/ml/input/data/training/droid_sample_ranges_v1_0_1.json",
            "normalize_action": true,
            "delta_action": true,
            "num_parallel_reads": 16
        }
    },
    "wandb": {
        "project": "latent-reasoning-icml",
        "entity": "tri",
        "job_type": "train",
        "group": "lmp-policy-ppo",
        "name": "alpha0.1_clipexpdiff0.05_ll0.01_s16v64_sig0.2-0.02_42",
        "resume": false
    },
    "optimizer": {
        "_target_": "common.optimization.get_optimizer",
        "name": "AdamW",
        "model": "???",
        "lr": 0.0001,
        "weight_decay": 1e-06,
        "betas": [
            0.9,
            0.999
        ],
        "eps": 1e-08,
        "obs_encoder_lr_scale": 0.1
    },
    "scheduler": {
        "_target_": "common.optimization.get_scheduler",
        "name": "constant",
        "optimizer": "???",
        "num_warmup_steps": 20
    },
    "algo": "lmp-policy-ppo",
    "exp_id": "alpha0.1_clipexpdiff0.05_ll0.01_s16v64_sig0.2-0.02",
    "seed": 42,
    "logdir": "/opt/ml/checkpoints/lmp-policy-ppo/droid_rlds/alpha0.1_clipexpdiff0.05_ll0.01_s16v64_sig0.2-0.02/42",
    "resume": false,
    "use_amp": true,
    "compile": false,
    "pretrain_checkpoint_path": null,
    "num_steps": 5000000,
    "train_every": 3200,
    "num_train_epochs": 2,
    "eval_every": 128000,
    "save_every": 128000,
    "eval_rollout_every": 1000000,
    "num_eval_rollouts": 20,
    "rollout_batch_size": 200,
    "train_batch_size": 64,
    "clip_grad_norm": 1.0,
    "free_nats_ratio": 0.05,
    "clip_range": 0.05,
    "alpha": 0.1,
    "kl_coef": 1.0,
    "ll_coef": 0.01,
    "model": {
        "_target_": "algorithms.latent_memory_palace.LatentMemoryPalacePolicy",
        "obs_encoder": {
            "_target_": "models.obs_encoders.SequenceObservationEncoder",
            "shape_meta": {
                "obs": {
                    "image": {
                        "shape": [
                            180,
                            320,
                            3
                        ],
                        "type": "rgb"
                    },
                    "wrist_image": {
                        "shape": [
                            180,
                            320,
                            3
                        ],
                        "type": "rgb"
                    },
                    "joint_position": {
                        "shape": [
                            7
                        ],
                        "type": "low_dim"
                    },
                    "gripper_position": {
                        "shape": [
                            1
                        ],
                        "type": "low_dim"
                    }
                },
                "action": {
                    "shape": [
                        8
                    ]
                }
            },
            "embed_dim": 768,
            "num_timesteps": 1,
            "resize_shape": [
                240,
                240
            ],
            "crop_shape": [
                224,
                224
            ],
            "random_crop": false,
            "color_jitter": null,
            "vision_backbone": "pi05",
            "use_low_dim": true,
            "use_language": true,
            "pool_output": true,
            "max_language_length": 64,
            "num_film_layers": 2
        },
        "action_dim": 8,
        "action_chunk": 16,
        "seq_len": 16,
        "vocab_size": 64,
        "enc_embed_dim": 768,
        "enc_depth": 12,
        "enc_num_heads": 12,
        "enc_mlp_ratio": 4.0,
        "enc_num_registers": 1,
        "dec_embed_dim": 768,
        "dec_depth": 10,
        "dec_num_heads": 12,
        "dec_mlp_ratio": 4.0,
        "sigma_max": 0.2,
        "sigma_min": 0.02,
        "sigma_type": "learned_grid",
        "reg_type": "none",
        "reg_coef": 0
    }
}