admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-21 17:01:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

load from parquet

Browse Source
This commit is contained in:
Pepijn
2026-02-21 07:51:57 +01:00
parent 8abc9037a3
commit f906270ec4
1 changed files with 9 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/scripts/lerobot_train.py
+9 -5
View File
@@ -254,16 +254,20 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None):
from lerobot.processor.delta_action_processor import to_delta_actions from lerobot.processor.delta_action_processor import to_delta_actions
max_samples = min(100_000, len(dataset)) max_samples = min(100_000, len(dataset))
indices = np.random.choice(len(dataset), max_samples, replace=False) indices = np.random.choice(len(dataset), max_samples, replace=False).tolist()
logging.info( logging.info(
f"use_delta_actions is enabled — computing delta action stats from {max_samples} dataset chunks" f"use_delta_actions is enabled — computing delta action stats from {max_samples} dataset chunks"
) )
# Read only action and state from parquet (no video decoding)
hf = dataset.hf_dataset
actions_raw = hf.select(indices)["action"]
states_raw = hf.select(indices)["observation.state"]
all_delta_actions = [] all_delta_actions = []
for i in indices: for action, state in zip(actions_raw, states_raw):
item = dataset[int(i)] action = torch.as_tensor(action).float()
action = item["action"] state = torch.as_tensor(state).float()
state = item["observation.state"]
if action.ndim == 1: if action.ndim == 1:
action = action.unsqueeze(0) action = action.unsqueeze(0)
mask = [True] * action.shape[-1] mask = [True] * action.shape[-1]