awm_diffusion_policy / config.json
anantk2006's picture
Upload policy weights, train config and readme
14a3199 verified
{
"type": "awm",
"n_obs_steps": 1,
"input_features": {
"observation.image": {
"type": "VISUAL",
"shape": [
3,
96,
96
]
},
"observation.state": {
"type": "STATE",
"shape": [
2
]
}
},
"output_features": {
"action": {
"type": "ACTION",
"shape": [
2
]
}
},
"device": "cuda",
"use_amp": false,
"use_peft": false,
"push_to_hub": true,
"repo_id": "anantk2006/awm_diffusion_policy",
"private": null,
"tags": null,
"license": null,
"pretrained_path": null,
"chunk_size": 10,
"n_action_steps": 5,
"normalization_mapping": {
"VISUAL": "MEAN_STD",
"STATE": "MEAN_STD",
"ACTION": "MIN_MAX"
},
"vision_backbone": "resnet18",
"pretrained_backbone_weights": "ResNet18_Weights.IMAGENET1K_V1",
"replace_final_stride_with_dilation": false,
"pre_norm": false,
"dim_model": 512,
"n_heads": 8,
"dim_feedforward": 3200,
"feedforward_activation": "relu",
"n_encoder_layers": 4,
"n_decoder_layers": 4,
"cross_attn_dim": 512,
"action_token_vocab_size": 128,
"action_ranges": null,
"dropout": 0.1,
"wm_loss_weight": 0.1,
"detach_encoder_from_wm": false,
"use_normalized_mse_wm_loss": false,
"wm_variance_loss_weight": 0.1,
"use_ema_target": false,
"ema_momentum": 0.996,
"ema_momentum_end": 0.999,
"ema_anneal_steps": 50000,
"n_wm_decoder_layers": 4,
"decoder_loss_weight": 0.1,
"n_image_viz_pairs": 12,
"optimizer_lr": 0.0002,
"optimizer_weight_decay": 0.0,
"optimizer_lr_backbone": 1e-05,
"optimizer_grad_clip_norm": 1.0,
"use_lr_schedule": false,
"scheduler_warmup_steps": 5000,
"scheduler_decay_steps": 100000,
"scheduler_decay_lr": 0.0
}