| { |
| "_target_": "lewm_model.LeWorldModel", |
| "encoder": { |
| "image_size": 224, |
| "patch_size": 14, |
| "embed_dim": 192, |
| "num_layers": 12, |
| "num_heads": 3, |
| "intermediate_size": 768 |
| }, |
| "predictor": { |
| "num_frames": 3, |
| "input_dim": 192, |
| "hidden_dim": 192, |
| "output_dim": 192, |
| "depth": 6, |
| "heads": 16, |
| "mlp_dim": 2048, |
| "dim_head": 64, |
| "dropout": 0.1, |
| "emb_dropout": 0.0 |
| }, |
| "action_encoder": { |
| "input_dim": 10, |
| "emb_dim": 192, |
| "mlp_scale": 4 |
| }, |
| "projector": { |
| "input_dim": 192, |
| "output_dim": 192, |
| "hidden_dim": 2048, |
| "norm_fn": "torch.nn.BatchNorm1d" |
| }, |
| "pred_proj": { |
| "input_dim": 192, |
| "output_dim": 192, |
| "hidden_dim": 2048, |
| "norm_fn": "torch.nn.BatchNorm1d" |
| }, |
| "training": { |
| "batch_size": 128, |
| "epochs": 10, |
| "frameskip": 5, |
| "history_size": 3, |
| "seq_len": 4, |
| "lambda": 0.1, |
| "lr": 0.001, |
| "weight_decay": 0.05, |
| "grad_clip": 1.0, |
| "optimizer": "AdamW", |
| "scheduler": "LinearWarmupCosineAnnealingLR" |
| }, |
| "planning": { |
| "cem_samples": 300, |
| "cem_iterations": 30, |
| "cem_elites": 30, |
| "cem_init_var": 1.0, |
| "planning_horizon": 5 |
| } |
| } |
|
|
