File size: 1,463 Bytes
987ed1b | 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 | if __name__ == "__main__":
import sys
import os
import pathlib
ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent)
sys.path.append(ROOT_DIR)
import os
import click
import pathlib
import numpy as np
from diffusion_policy.common.replay_buffer import ReplayBuffer
@click.command()
@click.option('-i', '--input', required=True, help='input dir contains npy files')
@click.option('-o', '--output', required=True, help='output zarr path')
@click.option('--abs_action', is_flag=True, default=False)
def main(input, output, abs_action):
data_directory = pathlib.Path(input)
observations = np.load(
data_directory / "multimodal_push_observations.npy"
)
actions = np.load(data_directory / "multimodal_push_actions.npy")
masks = np.load(data_directory / "multimodal_push_masks.npy")
buffer = ReplayBuffer.create_empty_numpy()
for i in range(len(masks)):
eps_len = int(masks[i].sum())
obs = observations[i,:eps_len].astype(np.float32)
action = actions[i,:eps_len].astype(np.float32)
if abs_action:
prev_eef_target = obs[:,8:10]
next_eef_target = prev_eef_target + action
action = next_eef_target
data = {
'obs': obs,
'action': action
}
buffer.add_episode(data)
buffer.save_to_path(zarr_path=output, chunk_length=-1)
if __name__ == '__main__':
main()
|