This commit is contained in:
Remi Cadene 2024-10-13 17:15:40 +02:00
parent 1f0bdecaf8
commit d02e204e10
2 changed files with 16 additions and 8 deletions

View File

@ -95,7 +95,7 @@ jobs:
- name: Test with pytest
run: |
pytest 'tests/test_control_robot.py::test_record_and_replay_and_policy[koch-True]' -v --cov=./lerobot --durations=0 \
pytest -sx 'tests/test_control_robot.py::test_record_and_replay_and_policy[koch-True]' -v --cov=./lerobot --durations=0 \
-W ignore::DeprecationWarning:imageio_ffmpeg._utils:7 \
-W ignore::UserWarning:torch.utils.data.dataloader:558 \
-W ignore::UserWarning:gymnasium.utils.env_checker:247 \

View File

@ -200,13 +200,13 @@ def to_hf_dataset(data_dict, video) -> Dataset:
features["next.done"] = Value(dtype="bool", id=None)
features["index"] = Value(dtype="int64", id=None)
# for key in data_dict:
# if isinstance(data_dict[key], list):
# print(key, len(data_dict[key]))
# elif isinstance(data_dict[key], torch.Tensor):
# print(key, data_dict[key].shape)
# else:
# print(key, data_dict[key])
for key in data_dict:
if isinstance(data_dict[key], list):
print(key, len(data_dict[key]))
elif isinstance(data_dict[key], torch.Tensor):
print(key, data_dict[key].shape)
else:
print(key, data_dict[key])
data_dict["episode_index"] = data_dict["episode_index"].tolist()
data_dict["frame_index"] = data_dict["frame_index"].tolist()
@ -214,6 +214,14 @@ def to_hf_dataset(data_dict, video) -> Dataset:
data_dict["next.done"] = data_dict["next.done"].tolist()
data_dict["index"] = data_dict["index"].tolist()
for key in data_dict:
if isinstance(data_dict[key], list):
print(key, len(data_dict[key]))
elif isinstance(data_dict[key], torch.Tensor):
print(key, data_dict[key].shape)
else:
print(key, data_dict[key])
hf_dataset = Dataset.from_dict(data_dict, features=Features(features))
hf_dataset.set_transform(hf_transform_to_torch)
return hf_dataset