-
Notifications
You must be signed in to change notification settings - Fork 12
/
Copy pathdiffusion-policy.patch
216 lines (207 loc) · 7.45 KB
/
diffusion-policy.patch
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
diff --git a/diffusion_policy/common/normalize_util.py b/diffusion_policy/common/normalize_util.py
index d6fc7b6..0fcc7f0 100644
--- a/diffusion_policy/common/normalize_util.py
+++ b/diffusion_policy/common/normalize_util.py
@@ -110,8 +110,8 @@ def robomimic_abs_action_normalizer_from_stat(stat, rotation_transformer):
def robomimic_abs_action_only_normalizer_from_stat(stat):
result = dict_apply_split(
stat, lambda x: {
- 'pos': x[...,:3],
- 'other': x[...,3:]
+ 'pos': x[...,:6],
+ 'other': x[...,6:]
})
def get_pos_param_info(stat, output_max=1, output_min=-1, range_eps=1e-7):
diff --git a/diffusion_policy/config/task/square_image_abs.yaml b/diffusion_policy/config/task/square_image_abs.yaml
index d27a916..53e735d 100644
--- a/diffusion_policy/config/task/square_image_abs.yaml
+++ b/diffusion_policy/config/task/square_image_abs.yaml
@@ -3,50 +3,29 @@ name: square_image
shape_meta: &shape_meta
# acceptable types: rgb, low_dim
obs:
- agentview_image:
+ base_image:
shape: [3, 84, 84]
type: rgb
- robot0_eye_in_hand_image:
+ wrist_image:
shape: [3, 84, 84]
type: rgb
- robot0_eef_pos:
+ base_pose:
shape: [3]
# type default: low_dim
- robot0_eef_quat:
+ arm_pos:
+ shape: [3]
+ arm_quat:
shape: [4]
- robot0_gripper_qpos:
- shape: [2]
+ gripper_pos:
+ shape: [1]
action:
- shape: [10]
+ shape: [13]
-task_name: &task_name square
-dataset_type: &dataset_type ph
-dataset_path: &dataset_path data/robomimic/datasets/${task.task_name}/${task.dataset_type}/image_abs.hdf5
+dataset_path: &dataset_path data/${task.name}.hdf5
abs_action: &abs_action True
env_runner:
- _target_: diffusion_policy.env_runner.robomimic_image_runner.RobomimicImageRunner
- dataset_path: *dataset_path
- shape_meta: *shape_meta
- # costs 1GB per env
- n_train: 6
- n_train_vis: 2
- train_start_idx: 0
- n_test: 50
- n_test_vis: 4
- test_start_seed: 100000
- # use python's eval function as resolver, single-quoted string as argument
- max_steps: ${eval:'500 if "${task.dataset_type}" == "mh" else 400'}
- n_obs_steps: ${n_obs_steps}
- n_action_steps: ${n_action_steps}
- render_obs_key: 'agentview_image'
- fps: 10
- crf: 22
- past_action: ${past_action_visible}
- abs_action: *abs_action
- tqdm_interval_sec: 1.0
- n_envs: 28
-# evaluation at this config requires a 16 core 64GB instance.
+ _target_: diffusion_policy.env_runner.real_pusht_image_runner.RealPushTImageRunner
dataset:
_target_: diffusion_policy.dataset.robomimic_replay_image_dataset.RobomimicReplayImageDataset
@@ -61,4 +40,4 @@ dataset:
use_legacy_normalizer: False
use_cache: True
seed: 42
- val_ratio: 0.02
+ val_ratio: 0.0
diff --git a/diffusion_policy/config/train_diffusion_unet_real_hybrid_workspace.yaml b/diffusion_policy/config/train_diffusion_unet_real_hybrid_workspace.yaml
index 5b32366..82c95b9 100644
--- a/diffusion_policy/config/train_diffusion_unet_real_hybrid_workspace.yaml
+++ b/diffusion_policy/config/train_diffusion_unet_real_hybrid_workspace.yaml
@@ -1,6 +1,6 @@
defaults:
- _self_
- - task: real_pusht_image
+ - task: square_image_abs
name: train_diffusion_unet_hybrid
_target_: diffusion_policy.workspace.train_diffusion_unet_hybrid_workspace.TrainDiffusionUnetHybridWorkspace
@@ -39,10 +39,9 @@ policy:
horizon: ${horizon}
n_action_steps: ${eval:'${n_action_steps}+${n_latency_steps}'}
n_obs_steps: ${n_obs_steps}
- num_inference_steps: 8
+ num_inference_steps: 16
obs_as_global_cond: ${obs_as_global_cond}
- # crop_shape: [76, 76] # 84x84 90%
- crop_shape: [216, 288] # ch, cw 320x240 90%
+ crop_shape: [76, 76]
# crop_shape: null
diffusion_step_embed_dim: 128
down_dims: [256,512,1024]
@@ -68,14 +67,14 @@ dataloader:
num_workers: 8
shuffle: True
pin_memory: True
- persistent_workers: True
+ persistent_workers: False
val_dataloader:
batch_size: 64
num_workers: 8
shuffle: False
pin_memory: True
- persistent_workers: True
+ persistent_workers: False
optimizer:
_target_: torch.optim.AdamW
@@ -92,7 +91,7 @@ training:
# optimization
lr_scheduler: cosine
lr_warmup_steps: 500
- num_epochs: 600
+ num_epochs: 601
gradient_accumulate_every: 1
# EMA destroys performance when used with BatchNorm
# replace BatchNorm with GroupNorm.
@@ -100,7 +99,7 @@ training:
# training loop control
# in epochs
rollout_every: 50
- checkpoint_every: 50
+ checkpoint_every: 100
val_every: 1
sample_every: 5
# steps per epoch
@@ -122,7 +121,7 @@ checkpoint:
topk:
monitor_key: train_loss
mode: min
- k: 5
+ k: 10
format_str: 'epoch={epoch:04d}-train_loss={train_loss:.3f}.ckpt'
save_last_ckpt: True
save_last_snapshot: False
diff --git a/diffusion_policy/dataset/robomimic_replay_image_dataset.py b/diffusion_policy/dataset/robomimic_replay_image_dataset.py
index 2728e9e..a277d2d 100644
--- a/diffusion_policy/dataset/robomimic_replay_image_dataset.py
+++ b/diffusion_policy/dataset/robomimic_replay_image_dataset.py
@@ -23,7 +23,6 @@ from diffusion_policy.common.replay_buffer import ReplayBuffer
from diffusion_policy.common.sampler import SequenceSampler, get_val_mask
from diffusion_policy.common.normalize_util import (
robomimic_abs_action_only_normalizer_from_stat,
- robomimic_abs_action_only_dual_arm_normalizer_from_stat,
get_range_normalizer_from_stat,
get_image_range_normalizer,
get_identity_normalizer_from_stat,
@@ -151,11 +150,7 @@ class RobomimicReplayImageDataset(BaseImageDataset):
# action
stat = array_to_stats(self.replay_buffer['action'])
if self.abs_action:
- if stat['mean'].shape[-1] > 10:
- # dual arm
- this_normalizer = robomimic_abs_action_only_dual_arm_normalizer_from_stat(stat)
- else:
- this_normalizer = robomimic_abs_action_only_normalizer_from_stat(stat)
+ this_normalizer = robomimic_abs_action_only_normalizer_from_stat(stat)
if self.use_legacy_normalizer:
this_normalizer = normalizer_from_stat(stat)
@@ -175,6 +170,8 @@ class RobomimicReplayImageDataset(BaseImageDataset):
this_normalizer = get_identity_normalizer_from_stat(stat)
elif key.endswith('qpos'):
this_normalizer = get_range_normalizer_from_stat(stat)
+ elif key == 'base_pose':
+ this_normalizer = get_range_normalizer_from_stat(stat)
else:
raise RuntimeError('unsupported')
normalizer[key] = this_normalizer
@@ -223,22 +220,13 @@ class RobomimicReplayImageDataset(BaseImageDataset):
def _convert_actions(raw_actions, abs_action, rotation_transformer):
actions = raw_actions
if abs_action:
- is_dual_arm = False
- if raw_actions.shape[-1] == 14:
- # dual arm
- raw_actions = raw_actions.reshape(-1,2,7)
- is_dual_arm = True
-
- pos = raw_actions[...,:3]
- rot = raw_actions[...,3:6]
- gripper = raw_actions[...,6:]
+ pos = raw_actions[...,:6]
+ rot = raw_actions[...,6:9]
+ gripper = raw_actions[...,9:]
rot = rotation_transformer.forward(rot)
raw_actions = np.concatenate([
pos, rot, gripper
], axis=-1).astype(np.float32)
-
- if is_dual_arm:
- raw_actions = raw_actions.reshape(-1,20)
actions = raw_actions
return actions