config.json

{
    "task": {
        "dataset": {
            "_target_": "datasets.make_dataset",
            "seq_len": 1,
            "pad_mode": "none",
            "val_ratio": 0.05,
            "obs_stack": null,
            "action_chunk": 16,
            "to_memory": true,
            "dataset_id": "robomimic-all_actions-ph",
            "data_paths": [
                "/home/ubuntu/robomimic/datasets/lift/ph/image_v15.hdf5",
                "/home/ubuntu/robomimic/datasets/can/ph/image_v15.hdf5",
                "/home/ubuntu/robomimic/datasets/square/ph/image_v15.hdf5",
                "/home/ubuntu/robomimic/datasets/tool_hang/ph/image_v15.hdf5"
            ],
            "buffer_dir": "buffers/robomimic/all_actions/ph/",
            "shape_meta": {
                "obs": {},
                "action": {
                    "shape": [
                        7
                    ]
                }
            },
            "normalize_low_dim": true
        }
    },
    "optimizer": {
        "_target_": "common.optimization.get_optimizer",
        "name": "AdamW",
        "model": "???",
        "lr": 0.0001,
        "weight_decay": 1e-06,
        "betas": [
            0.9,
            0.999
        ],
        "eps": 1e-08
    },
    "scheduler": {
        "_target_": "common.optimization.get_scheduler",
        "name": "constant",
        "optimizer": "???"
    },
    "algo": "lmp-vae",
    "exp_id": "benchmark",
    "seed": 42,
    "logdir": "/home/ubuntu/logdir/lmp-vae/robomimic_all_actions/benchmark/42",
    "resume": false,
    "use_amp": false,
    "compile": true,
    "num_steps": 25000000,
    "train_every": 25000,
    "num_train_epochs": 2,
    "eval_every": 1000000,
    "save_every": 1000000,
    "eval_rollout_every": 1000000,
    "num_eval_rollouts": 20,
    "rollout_batch_size": 500,
    "train_batch_size": 500,
    "clip_grad_norm": 1.0,
    "free_nats_ratio": 0.2,
    "clip_range": 0.01,
    "kl_coef": 0.1,
    "ll_coef": 1.0,
    "model": {
        "_target_": "algorithms.latent_memory_palace.LatentMemoryPalaceVAE",
        "action_dim": 7,
        "action_chunk": 16,
        "seq_len": 8,
        "vocab_size": 64,
        "enc_embed_dim": 384,
        "enc_depth": 8,
        "enc_num_heads": 8,
        "enc_mlp_ratio": 4.0,
        "enc_num_registers": 1,
        "dec_embed_dim": 384,
        "dec_depth": 6,
        "dec_num_heads": 8,
        "dec_mlp_ratio": 4.0,
        "sigma_max": 1.0,
        "sigma_min": 1.0,
        "sigma_type": "learned_scalar"
    }
}