tangger/lerobot_piper
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Address comments

Browse Source
This commit is contained in:
Cadene
2024-04-16 17:14:40 +00:00
parent b241ea46dd
commit 36d9e885ef
24 changed files with 100 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
lerobot/scripts/eval.py
View File

@@ -215,8 +215,8 @@ def eval_policy(
"timestamp": torch.arange(0, num_frames, 1) / fps,
"next.done": dones[ep_id, :num_frames],
"next.reward": rewards[ep_id, :num_frames].type(torch.float32),
"episode_data_id_from": torch.tensor([idx_from] * num_frames),
"episode_data_id_to": torch.tensor([idx_from + num_frames - 1] * num_frames),
"episode_data_index_from": torch.tensor([idx_from] * num_frames),
"episode_data_index_to": torch.tensor([idx_from + num_frames] * num_frames),
}
for key in observations:
ep_dict[key] = observations[key][ep_id][:num_frames]

8
lerobot/scripts/train.py
View File

@@ -141,15 +141,15 @@ def add_episodes_inplace(data_dict, online_dataset, concat_dataset, sampler, pc_
online_dataset.data_dict = data_dict
else:
# find episode index and data frame indices according to previous episode in online_dataset
start_episode = online_dataset.data_dict["episode_id"][-1].item() + 1
start_index = online_dataset.data_dict["index"][-1].item() + 1
start_episode = online_dataset.select_columns("episode_id")[-1]["episode_id"].item() + 1
start_index = online_dataset.select_columns("index")[-1]["index"].item() + 1
def shift_indices(example):
# note: we dont shift "frame_id" since it represents the index of the frame in the episode it belongs to
example["episode_id"] += start_episode
example["index"] += start_index
example["episode_data_id_from"] += start_index
example["episode_data_id_to"] += start_index
example["episode_data_index_from"] += start_index
example["episode_data_index_to"] += start_index
return example
disable_progress_bar() # map has a tqdm progress bar

2
lerobot/scripts/visualize_dataset.py
View File

@@ -77,7 +77,7 @@ def render_dataset(dataset, out_dir, max_num_episodes):
# add current frame to list of frames to render
frames[im_key].append(item[im_key])
end_of_episode = item["index"].item() == item["episode_data_id_to"].item()
end_of_episode = item["index"].item() == item["episode_data_index_to"].item()
out_dir.mkdir(parents=True, exist_ok=True)
for im_key in dataset.image_keys: