Spaces:
Running
on
Zero
Running
on
Zero
File size: 7,502 Bytes
246c106 |
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 |
import h5py
import tqdm
import numpy as np
import torch
import cv2
import json
import imageio
from sim.simulator import GenieSimulator, ReplaySimulator
from sim.policy import ReplayPolicy
from diffusion_policy.util.pytorch_util import dict_apply
SYNTHETIC_DATA_RATIO = 0.25
DP_RES = 84
def load_demo(demo_idx: int):
with h5py.File('data/robomimic/datasets/lift/ph/image.hdf5', 'r') as f:
demo = f['data'][f'demo_{demo_idx}']
actions = demo['actions'][:].astype(np.float32)
frames = demo['obs']['agentview_image'][:].astype(np.uint8)
assert len(actions) == len(frames)
return frames, actions
def main():
demos = dict()
demos_len = []
start_idx = int(200 * (1 - SYNTHETIC_DATA_RATIO))
end_idx = 200
print(f"Generating {end_idx - start_idx} synthetic demos")
# copy actual data
for demo_idx in tqdm.tqdm(range(start_idx)):
with h5py.File('data/robomimic/datasets/lift/ph/image.hdf5', 'r') as f:
demo = f['data'][f'demo_{demo_idx}']
actions = demo['actions'][:].astype(np.float32)
frames = demo['obs']['agentview_image'][:].astype(np.uint8)
assert len(actions) == len(frames)
assert frames.shape[-3:] == (DP_RES, DP_RES, 3), frames.shape
demos[f"demo_{demo_idx}"] = {
"obs": {
"agentview_image": frames
},
"actions": actions
}
demos_len.append(len(actions))
print(f"Loaded {len(demos)} actual demos")
genie_simulator = GenieSimulator(
image_encoder_type='temporalvae',
image_encoder_ckpt='stabilityai/stable-video-diffusion-img2vid',
quantize=False,
backbone_type="stmar",
backbone_ckpt="data/mar_ckpt/robomimic_best",
prompt_horizon=11,
action_stride=1,
domain='robomimic',
)
# generate synthetic data
# synthetic data in training set only
for demo_idx in tqdm.tqdm(range(start_idx, end_idx)):
frames, actions = load_demo(demo_idx)
replay_policy = ReplayPolicy(actions, 1, 11)
replay_simulator = ReplaySimulator(frames, 11)
assert len(replay_policy) == len(replay_simulator), \
(len(replay_policy), len(replay_simulator))
# prompt genie
genie_simulator.set_initial_state((
replay_simulator.prompt(),
replay_policy.prompt()
))
image = genie_simulator.reset()
this_demo = {
"obs": {
"agentview_image": []
},
"actions": []
}
for _ in range(len(replay_policy)):
action = replay_policy.generate_action(None)
this_demo['obs']['agentview_image'].append(cv2.resize(image, (DP_RES, DP_RES)))
this_demo['actions'].append(action[0])
image = genie_simulator.step(action)['pred_next_frame']
this_demo = dict_apply(this_demo, lambda x: np.array(x))
demos[f"demo_{demo_idx}"] = this_demo
demos_len.append(len(this_demo['actions']))
with h5py.File(f'data/robomimic_synthetic/robomimic_synthetic{end_idx - start_idx}.hdf5', 'w') as f:
"""
saving format:
data (group)
total (attribute) - number of state-action samples in the dataset
env_args (attribute) - a json string that contains metadata on the environment and relevant arguments used for collecting data. Three keys: env_name, the name of the environment or task to create, env_type, one of robomimic’s supported environment types, and env_kwargs, a dictionary of keyword-arguments to be passed into the environment of type env_name.
demo_0 (group) - group for the first trajectory (every trajectory has a group)
num_samples (attribute) - the number of state-action samples in this trajectory
model_file (attribute) - the xml string corresponding to the MJCF MuJoCo model. Only present for robosuite datasets.
states (dataset) - flattened raw MuJoCo states, ordered by time. Shape (N, D) where N is the length of the trajectory, and D is the dimension of the state vector. Should be empty or have dummy values for non-robosuite datasets.
actions (dataset) - environment actions, ordered by time. Shape (N, A) where N is the length of the trajectory, and A is the action space dimension
rewards (dataset) - environment rewards, ordered by time. Shape (N,) where N is the length of the trajectory.
dones (dataset) - done signal, equal to 1 if playing the corresponding action in the state should terminate the episode. Shape (N,) where N is the length of the trajectory.
obs (group) - group for the observation keys. Each key is stored as a dataset.
<obs_key_1> (dataset) - the first observation key. Note that the name of this dataset and shape will vary. As an example, the name could be “agentview_image”, and the shape could be (N, 84, 84, 3).
...
next_obs (group) - group for the next observations.
<obs_key_1> (dataset) - the first observation key.
...
demo_1 (group) - group for the second trajectory
...
https://robomimic.github.io/docs/datasets/overview.html
"""
data_group = f.create_group('data')
data_group.attrs['total'] = sum(demos_len)
data_group.attrs['env_args'] = json.dumps({
'env_name': 'Lift',
'type': 1,
'env_kwargs': {
'has_renderer': False,
'has_offscreen_renderer': True,
'ignore_done': True,
'use_object_obs': False,
'use_camera_obs': True,
'control_freq': 20,
'controller_configs': {
'type': 'OSC_POSE',
'input_max': 1,
'input_min': -1,
'output_max': [0.05, 0.05, 0.05, 0.5, 0.5, 0.5],
'output_min': [-0.05, -0.05, -0.05, -0.5, -0.5, -0.5],
'kp': 150,
'damping': 1,
'impedance_mode': 'fixed',
'kp_limits': [0, 300],
'damping_limits': [0, 10],
'position_limits': None,
'orientation_limits': None,
'uncouple_pos_ori': True,
'control_delta': True,
'interpolation': None,
'ramp_ratio': 0.2
},
'robots': ['Panda'],
'camera_depths': False,
'camera_heights': 84,
'camera_widths': 84,
'reward_shaping': False,
'camera_names': ['agentview', 'robot0_eye_in_hand'],
'render_gpu_device_id': 0
},
'use_image_obs': True
})
for demo_idx, demo_data in demos.items():
demo_group = data_group.create_group(demo_idx)
demo_group.attrs['num_samples'] = len(demo_data['actions'])
demo_group.create_dataset('actions', data=demo_data['actions'])
obs_group = demo_group.create_group('obs')
for key, value in demo_data['obs'].items():
obs_group.create_dataset(key, data=value)
if __name__ == '__main__':
main()
|