policy.json

{
  "type": "rho",
  "name": "rho",
  "device": "cuda",
  "dtype": "bfloat16",
  "lr_scheduler": {
    "type": "warmup_stable_decay",
    "num_warmup_steps": 1000,
    "num_decay_steps": 45000,
    "peak_lr": 0.0001,
    "decay_lr": 1e-05,
    "num_cycles": 1,
    "num_rewarm_steps": null
  },
  "optimizer": {
    "type": "adamw",
    "lr": 1e-06,
    "weight_decay": 0.1,
    "grad_clip_norm": 10.0,
    "betas": [
      0.9,
      0.95
    ],
    "eps": 1e-06
  },
  "pretrained_repo_id": "microsoft/rho-base",
  "resize_imgs_with_padding": [
    256,
    256
  ],
  "embed_dim": 2048,
  "num_heads": 16,
  "ff_dim": 4096,
  "norm": "adaptive",
  "hidden_state_idx": 14,
  "num_blocks": 12,
  "max_seq_len": 6144,
  "n_obs_steps": 1,
  "chunk_size": 50,
  "n_action_steps": 25,
  "max_state_dim": 32,
  "max_action_dim": 32,
  "num_steps": 10,
  "num_flow_samples": 1,
  "attention_implementation": "flash_attention_2",
  "attention_type": "cross",
  "adaln_mode": "shared",
  "adaln_lora_rank": 256,
  "gqa_groups": 4,
  "cross_attention_dim": null,
  "vlm_layer_select": "last",
  "shared_vlm_projection": false,
  "kv_pos_encoding": true,
  "freeze_vlm_backbone": false,
  "freeze_vision_encoder": false,
  "freeze_language_model": false,
  "train_expert_only": false,
  "enable_gradient_checkpointing": true,
  "optimizer_lr": 1e-06,
  "optimizer_lr_action_expert": 0.0001,
  "optimizer_betas": [
    0.9,
    0.95
  ],
  "optimizer_eps": 1e-06,
  "optimizer_weight_decay": 0.1,
  "scheduler_warmup_steps": 1000,
  "scheduler_decay_steps": 30000,
  "scheduler_decay_lr": 2.5e-06,
  "time_sampling_strategy": "beta",
  "dropout_p": 0.1,
  "pos_emb_method": "sinusoidal",
  "log_hidden_state_stats": true
}