mirror of
https://github.com/Tavish9/any4lerobot.git
synced 2026-05-11 12:09:41 +00:00
4dc21b9b70
* add libero2lerobot readme * use datatrove for libero2lerobot * update libero2lerobot readme * update README.md * Update libero2lerobot/README.md Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * Update libero2lerobot/README.md Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * fix * set upload_large_folder to false * use vectorized operations for faster transform --------- Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
37 lines
1.5 KiB
Python
37 lines
1.5 KiB
Python
from pathlib import Path
|
|
|
|
import numpy as np
|
|
from h5py import File
|
|
|
|
|
|
def load_local_episodes(input_h5: Path):
|
|
with File(input_h5, "r") as f:
|
|
for demo in f["data"].values():
|
|
demo_len = len(demo["obs/agentview_rgb"])
|
|
# (-1: open, 1: close) -> (0: close, 1: open)
|
|
action = np.array(demo["actions"])
|
|
action = np.concatenate(
|
|
[
|
|
action[:, :6],
|
|
(1 - np.clip(action[:, -1], 0, 1))[:, None],
|
|
],
|
|
axis=1,
|
|
)
|
|
state = np.concatenate(
|
|
[
|
|
np.array(demo["obs/ee_states"]),
|
|
np.array(demo["obs/gripper_states"]),
|
|
],
|
|
axis=1,
|
|
)
|
|
episode = {
|
|
"observation.images.image": np.array(demo["obs/agentview_rgb"]),
|
|
"observation.images.wrist_image": np.array(demo["obs/eye_in_hand_rgb"]),
|
|
"observation.state": np.array(state, dtype=np.float32),
|
|
"observation.states.ee_state": np.array(demo["obs/ee_states"], dtype=np.float32),
|
|
"observation.states.joint_state": np.array(demo["obs/joint_states"], dtype=np.float32),
|
|
"observation.states.gripper_state": np.array(demo["obs/gripper_states"], dtype=np.float32),
|
|
"action": np.array(action, dtype=np.float32),
|
|
}
|
|
yield [{**{k: v[i] for k, v in episode.items()}} for i in range(demo_len)]
|