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