# Training config for the PyTorch pi05 finetune on robopro_top_cam (roboreal_lerobot).
# Trainer: scripts/train_pytorch.py  | env: pi05_pt  | 3xH200, batch 256 (85/GPU), 30000 steps.
# Weights stored in bf16 (compute+master); see docs/pytorch_fp32master_ema_fix.md for the fp32-master fix.

train_config = {'name': 'pi05_robopro_top_cam_pt',
 'project_name': 'openpi',
 'exp_name': 'robopro_top_cam_pi05_pt',
 'model': {'action_dim': 32,
           'action_horizon': 50,
           'max_token_len': 200,
           'dtype': 'bfloat16',
           'paligemma_variant': 'gemma_2b',
           'action_expert_variant': 'gemma_300m',
           'pi05': True,
           'discrete_state_input': True},
 'weight_loader': {},
 'pytorch_weight_path': '/work/xuan/openpi_assets/pi05_base_pytorch',
 'pytorch_training_precision': 'bfloat16',
 'lr_schedule': {'warmup_steps': 1000, 'peak_lr': 2.5e-05, 'decay_steps': 30000, 'decay_lr': 2.5e-06},
 'optimizer': {'b1': 0.9, 'b2': 0.95, 'eps': 1e-08, 'weight_decay': 1e-10, 'clip_gradient_norm': 1.0},
 'ema_decay': 0.99,
 'freeze_filter': "'Nothing'",
 'data': {'repo_id': 'roboreal_lerobot',
          'assets': {'assets_dir': None, 'asset_id': None},
          'base_config': {'repo_id': None,
                          'asset_id': None,
                          'norm_stats': None,
                          'repack_transforms': {'inputs': (), 'outputs': ()},
                          'data_transforms': {'inputs': (), 'outputs': ()},
                          'model_transforms': {'inputs': (), 'outputs': ()},
                          'use_quantile_norm': False,
                          'action_sequence_keys': ('actions',),
                          'prompt_from_task': True,
                          'rlds_data_dir': None,
                          'action_space': None,
                          'filter_dict_path': None},
          'use_delta_joint_actions': True,
          'default_prompt': None,
          'adapt_to_pi': True,
          'repack_transforms': {'inputs': [{'structure': {'images': {'cam_high': 'observation.images.countertop',
                                                                     'cam_left_wrist': 'observation.images.left',
                                                                     'cam_right_wrist': 'observation.images.right'},
                                                          'state': 'observation.state',
                                                          'actions': 'action',
                                                          'prompt': 'prompt'}}],
                                'outputs': ()},
          'action_sequence_keys': ('action',)},
 'assets_base_dir': './assets',
 'checkpoint_base_dir': './checkpoints',
 'seed': 42,
 'batch_size': 256,
 'num_workers': 16,
 'num_train_steps': 30000,
 'log_interval': 100,
 'save_interval': 1000,
 'keep_period': 5000,
 'overwrite': False,
 'resume': False,
 'wandb_enabled': True,
 'fsdp_devices': 1}