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