Hugging Face's logo

mzxuan
/
robopro_jax_30000

Safetensors
Model card Files
xet
 Community
mzxuan commited on
Commit
30e8b91
·
verified ·
1 Parent(s): a8b9cb4

Add PyTorch pi05 finetune (30000 steps)

Browse files
Files changed (4) hide show
  1. pytorch_30000/assets/roboreal_lerobot/norm_stats.json +136 -0
  2. pytorch_30000/metadata.pt +3 -0
  3. pytorch_30000/model.safetensors +3 -0
  4. pytorch_30000/train_config.py +60 -0
pytorch_30000/assets/roboreal_lerobot/norm_stats.json ADDED
@@ -0,0 +1,136 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "norm_stats": {
3
+ "state": {
4
+ "mean": [
5
+ -0.05335306376218796,
6
+ 0.7055787444114685,
7
+ 0.4741179049015045,
8
+ -0.1072346419095993,
9
+ 0.016872553154826164,
10
+ 0.048868030309677124,
11
+ 0.7768700122833252,
12
+ 0.06213012710213661,
13
+ 0.9724478721618652,
14
+ 0.6687734127044678,
15
+ -0.099928118288517,
16
+ 0.006652283016592264,
17
+ -0.04845580831170082,
18
+ 0.7147487998008728
19
+ ],
20
+ "std": [
21
+ 0.36716246604919434,
22
+ 0.9519953727722168,
23
+ 0.717629075050354,
24
+ 0.604361355304718,
25
+ 0.27203142642974854,
26
+ 0.5996415019035339,
27
+ 0.403693825006485,
28
+ 0.43324658274650574,
29
+ 0.9984343647956848,
30
+ 0.7902352809906006,
31
+ 0.738975465297699,
32
+ 0.3185270130634308,
33
+ 0.548449695110321,
34
+ 0.43587738275527954
35
+ ],
36
+ "q01": [
37
+ -1.1223593841552733,
38
+ -0.0032862133979797425,
39
+ -0.0046620978623628795,
40
+ -1.5200569910049437,
41
+ -0.9559813661575318,
42
+ -2.1879266061782836,
43
+ 0.0,
44
+ -1.0009263712883003,
45
+ -0.004651855564117402,
46
+ -0.003755022734403604,
47
+ -1.6461174297332763,
48
+ -1.145637583732605,
49
+ -1.8996453882217406,
50
+ 0.0
51
+ ],
52
+ "q99": [
53
+ 0.7493288616180411,
54
+ 2.656398411273956,
55
+ 2.521770501103997,
56
+ 1.4946733434677126,
57
+ 1.1651067465543745,
58
+ 2.0037146360397333,
59
+ 0.9998,
60
+ 1.0761545074462884,
61
+ 2.7071886105298995,
62
+ 2.7838811550661924,
63
+ 1.6799157733917238,
64
+ 1.018457032775879,
65
+ 1.2326969148635865,
66
+ 0.9998
67
+ ]
68
+ },
69
+ "actions": {
70
+ "mean": [
71
+ -0.010721348226070404,
72
+ 0.055250801146030426,
73
+ 0.037130311131477356,
74
+ -0.00516910757869482,
75
+ 0.0009717112989164889,
76
+ -0.0035654695238918066,
77
+ 0.7738659977912903,
78
+ 0.0017995729576796293,
79
+ 0.07148820906877518,
80
+ 0.053889963775873184,
81
+ -0.009835179895162582,
82
+ -0.002389969304203987,
83
+ -0.0005170264048501849,
84
+ 0.7156321406364441
85
+ ],
86
+ "std": [
87
+ 0.2341509759426117,
88
+ 0.42798152565956116,
89
+ 0.3539995849132538,
90
+ 0.3226429224014282,
91
+ 0.17103521525859833,
92
+ 0.3472360670566559,
93
+ 0.40591976046562195,
94
+ 0.29239627718925476,
95
+ 0.533614993095398,
96
+ 0.4418846368789673,
97
+ 0.40809792280197144,
98
+ 0.2093358039855957,
99
+ 0.3202282190322876,
100
+ 0.4358242452144623
101
+ ],
102
+ "q01": [
103
+ -0.8917047828674312,
104
+ -1.2250273670196534,
105
+ -1.0893902339935306,
106
+ -1.153321870803833,
107
+ -0.5912689544677736,
108
+ -1.3054964709281922,
109
+ 0.0,
110
+ -1.0367932884216309,
111
+ -1.6878083958625794,
112
+ -1.3210511589050293,
113
+ -1.3179429388046264,
114
+ -0.8158465300798414,
115
+ -1.2375638717651363,
116
+ 0.0
117
+ ],
118
+ "q99": [
119
+ 0.8131462455749521,
120
+ 1.9042484401702877,
121
+ 1.3641196742057797,
122
+ 1.1428229305744173,
123
+ 0.6477880647659298,
124
+ 1.2141571140289305,
125
+ 0.9998,
126
+ 0.9971832927703854,
127
+ 2.069108867883682,
128
+ 1.67551022567749,
129
+ 1.4281046690940862,
130
+ 0.7870594761371614,
131
+ 1.1032438911437996,
132
+ 0.9998
133
+ ]
134
+ }
135
+ }
136
+ }
pytorch_30000/metadata.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:88804fccb57e41311e7ee1671b2d8c59cda4cc442304fee62293110c831f3407
3
+ size 2852
pytorch_30000/model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:224bfbdf719c4897d6e7094a9a9f7c774e8cf5cc965b124f2aee8793ea850b88
3
+ size 7473091464
pytorch_30000/train_config.py ADDED
@@ -0,0 +1,60 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Training config for the PyTorch pi05 finetune on robopro_top_cam (roboreal_lerobot).
2
+ # Trainer: scripts/train_pytorch.py | env: pi05_pt | 3xH200, batch 256 (85/GPU), 30000 steps.
3
+ # Weights stored in bf16 (compute+master); see docs/pytorch_fp32master_ema_fix.md for the fp32-master fix.
4
+
5
+ train_config = {'name': 'pi05_robopro_top_cam_pt',
6
+ 'project_name': 'openpi',
7
+ 'exp_name': 'robopro_top_cam_pi05_pt',
8
+ 'model': {'action_dim': 32,
9
+ 'action_horizon': 50,
10
+ 'max_token_len': 200,
11
+ 'dtype': 'bfloat16',
12
+ 'paligemma_variant': 'gemma_2b',
13
+ 'action_expert_variant': 'gemma_300m',
14
+ 'pi05': True,
15
+ 'discrete_state_input': True},
16
+ 'weight_loader': {},
17
+ 'pytorch_weight_path': '/work/xuan/openpi_assets/pi05_base_pytorch',
18
+ 'pytorch_training_precision': 'bfloat16',
19
+ 'lr_schedule': {'warmup_steps': 1000, 'peak_lr': 2.5e-05, 'decay_steps': 30000, 'decay_lr': 2.5e-06},
20
+ 'optimizer': {'b1': 0.9, 'b2': 0.95, 'eps': 1e-08, 'weight_decay': 1e-10, 'clip_gradient_norm': 1.0},
21
+ 'ema_decay': 0.99,
22
+ 'freeze_filter': "'Nothing'",
23
+ 'data': {'repo_id': 'roboreal_lerobot',
24
+ 'assets': {'assets_dir': None, 'asset_id': None},
25
+ 'base_config': {'repo_id': None,
26
+ 'asset_id': None,
27
+ 'norm_stats': None,
28
+ 'repack_transforms': {'inputs': (), 'outputs': ()},
29
+ 'data_transforms': {'inputs': (), 'outputs': ()},
30
+ 'model_transforms': {'inputs': (), 'outputs': ()},
31
+ 'use_quantile_norm': False,
32
+ 'action_sequence_keys': ('actions',),
33
+ 'prompt_from_task': True,
34
+ 'rlds_data_dir': None,
35
+ 'action_space': None,
36
+ 'filter_dict_path': None},
37
+ 'use_delta_joint_actions': True,
38
+ 'default_prompt': None,
39
+ 'adapt_to_pi': True,
40
+ 'repack_transforms': {'inputs': [{'structure': {'images': {'cam_high': 'observation.images.countertop',
41
+ 'cam_left_wrist': 'observation.images.left',
42
+ 'cam_right_wrist': 'observation.images.right'},
43
+ 'state': 'observation.state',
44
+ 'actions': 'action',
45
+ 'prompt': 'prompt'}}],
46
+ 'outputs': ()},
47
+ 'action_sequence_keys': ('action',)},
48
+ 'assets_base_dir': './assets',
49
+ 'checkpoint_base_dir': './checkpoints',
50
+ 'seed': 42,
51
+ 'batch_size': 256,
52
+ 'num_workers': 16,
53
+ 'num_train_steps': 30000,
54
+ 'log_interval': 100,
55
+ 'save_interval': 1000,
56
+ 'keep_period': 5000,
57
+ 'overwrite': False,
58
+ 'resume': False,
59
+ 'wandb_enabled': True,
60
+ 'fsdp_devices': 1}