181 lines
6.5 KiB
Python
181 lines
6.5 KiB
Python
import os
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
from torchvision.transforms import v2
|
|
|
|
from lerobot.common.datasets.utils import compute_or_load_stats
|
|
from lerobot.common.transforms import NormalizeTransform, Prod
|
|
|
|
# DATA_DIR specifies to location where datasets are loaded. By default, DATA_DIR is None and
|
|
# we load from `$HOME/.cache/huggingface/hub/datasets`. For our unit tests, we set `DATA_DIR=tests/data`
|
|
# to load a subset of our datasets for faster continuous integration.
|
|
DATA_DIR = Path(os.environ["DATA_DIR"]) if "DATA_DIR" in os.environ else None
|
|
|
|
|
|
def make_dataset(
|
|
cfg,
|
|
# set normalize=False to remove all transformations and keep images unnormalized in [0,255]
|
|
normalize=True,
|
|
stats_path=None,
|
|
):
|
|
if cfg.env.name == "simxarm":
|
|
from lerobot.common.datasets.simxarm import SimxarmDataset
|
|
|
|
clsfunc = SimxarmDataset
|
|
|
|
elif cfg.env.name == "pusht":
|
|
from lerobot.common.datasets.pusht import PushtDataset
|
|
|
|
clsfunc = PushtDataset
|
|
|
|
elif cfg.env.name == "aloha":
|
|
from lerobot.common.datasets.aloha import AlohaDataset
|
|
|
|
clsfunc = AlohaDataset
|
|
else:
|
|
raise ValueError(cfg.env.name)
|
|
|
|
transforms = None
|
|
if normalize:
|
|
# TODO(rcadene): make normalization strategy configurable between mean_std, min_max, manual_min_max,
|
|
# min_max_from_spec
|
|
# stats = dataset.compute_or_load_stats() if stats_path is None else torch.load(stats_path)
|
|
|
|
if cfg.policy.name == "diffusion" and cfg.env.name == "pusht":
|
|
stats = {}
|
|
# TODO(rcadene): we overwrite stats to have the same as pretrained model, but we should remove this
|
|
stats["observation.state"] = {}
|
|
stats["observation.state"]["min"] = torch.tensor([13.456424, 32.938293], dtype=torch.float32)
|
|
stats["observation.state"]["max"] = torch.tensor([496.14618, 510.9579], dtype=torch.float32)
|
|
stats["action"] = {}
|
|
stats["action"]["min"] = torch.tensor([12.0, 25.0], dtype=torch.float32)
|
|
stats["action"]["max"] = torch.tensor([511.0, 511.0], dtype=torch.float32)
|
|
else:
|
|
# instantiate a one frame dataset with light transform
|
|
stats_dataset = clsfunc(
|
|
dataset_id=cfg.dataset_id,
|
|
root=DATA_DIR,
|
|
transform=Prod(in_keys=clsfunc.image_keys, prod=1 / 255.0),
|
|
)
|
|
stats = compute_or_load_stats(stats_dataset)
|
|
# TODO(rcadene): remove this and put it in config. Ideally we want to reproduce SOTA results just with mean_std
|
|
normalization_mode = "mean_std" if cfg.env.name == "aloha" else "min_max"
|
|
|
|
# # TODO(now): These stats are needed to use their pretrained model for sim_transfer_cube_human.
|
|
# # (Pdb) stats['observation']['state']['mean']
|
|
# # tensor([-0.0071, -0.6293, 1.0351, -0.0517, -0.4642, -0.0754, 0.4751, -0.0373,
|
|
# # -0.3324, 0.9034, -0.2258, -0.3127, -0.2412, 0.6866])
|
|
# stats["observation", "state", "mean"] = torch.tensor(
|
|
# [
|
|
# -0.00740268,
|
|
# -0.63187766,
|
|
# 1.0356655,
|
|
# -0.05027218,
|
|
# -0.46199223,
|
|
# -0.07467502,
|
|
# 0.47467607,
|
|
# -0.03615446,
|
|
# -0.33203387,
|
|
# 0.9038929,
|
|
# -0.22060776,
|
|
# -0.31011587,
|
|
# -0.23484458,
|
|
# 0.6842416,
|
|
# ]
|
|
# )
|
|
# # (Pdb) stats['observation']['state']['std']
|
|
# # tensor([0.0022, 0.0520, 0.0291, 0.0092, 0.0267, 0.0145, 0.0563, 0.0179, 0.0494,
|
|
# # 0.0326, 0.0476, 0.0535, 0.0956, 0.0513])
|
|
# stats["observation", "state", "std"] = torch.tensor(
|
|
# [
|
|
# 0.01219023,
|
|
# 0.2975381,
|
|
# 0.16728032,
|
|
# 0.04733803,
|
|
# 0.1486037,
|
|
# 0.08788499,
|
|
# 0.31752336,
|
|
# 0.1049916,
|
|
# 0.27933604,
|
|
# 0.18094037,
|
|
# 0.26604933,
|
|
# 0.30466506,
|
|
# 0.5298686,
|
|
# 0.25505227,
|
|
# ]
|
|
# )
|
|
# # (Pdb) stats['action']['mean']
|
|
# # tensor([-0.0075, -0.6346, 1.0353, -0.0465, -0.4686, -0.0738, 0.3723, -0.0396,
|
|
# # -0.3184, 0.8991, -0.2065, -0.3182, -0.2338, 0.5593])
|
|
# stats["action"]["mean"] = torch.tensor(
|
|
# [
|
|
# -0.00756444,
|
|
# -0.6281845,
|
|
# 1.0312834,
|
|
# -0.04664314,
|
|
# -0.47211358,
|
|
# -0.074527,
|
|
# 0.37389806,
|
|
# -0.03718753,
|
|
# -0.3261143,
|
|
# 0.8997205,
|
|
# -0.21371077,
|
|
# -0.31840396,
|
|
# -0.23360962,
|
|
# 0.551947,
|
|
# ]
|
|
# )
|
|
# # (Pdb) stats['action']['std']
|
|
# # tensor([0.0023, 0.0514, 0.0290, 0.0086, 0.0263, 0.0143, 0.0593, 0.0185, 0.0510,
|
|
# # 0.0328, 0.0478, 0.0531, 0.0945, 0.0794])
|
|
# stats["action"]["std"] = torch.tensor(
|
|
# [
|
|
# 0.01252818,
|
|
# 0.2957442,
|
|
# 0.16701928,
|
|
# 0.04584508,
|
|
# 0.14833844,
|
|
# 0.08763024,
|
|
# 0.30665937,
|
|
# 0.10600077,
|
|
# 0.27572668,
|
|
# 0.1805853,
|
|
# 0.26304692,
|
|
# 0.30708534,
|
|
# 0.5305411,
|
|
# 0.38381037,
|
|
# ]
|
|
# )
|
|
# transforms.append(NormalizeTransform(stats, in_keys, mode=normalization_mode)) # noqa: F821
|
|
|
|
transforms = v2.Compose(
|
|
[
|
|
# TODO(rcadene): we need to do something about image_keys
|
|
Prod(in_keys=clsfunc.image_keys, prod=1 / 255.0),
|
|
NormalizeTransform(
|
|
stats,
|
|
in_keys=[
|
|
"observation.state",
|
|
"action",
|
|
],
|
|
mode=normalization_mode,
|
|
),
|
|
]
|
|
)
|
|
|
|
delta_timestamps = cfg.policy.get("delta_timestamps")
|
|
if delta_timestamps is not None:
|
|
for key in delta_timestamps:
|
|
if isinstance(delta_timestamps[key], str):
|
|
delta_timestamps[key] = eval(delta_timestamps[key])
|
|
|
|
dataset = clsfunc(
|
|
dataset_id=cfg.dataset_id,
|
|
root=DATA_DIR,
|
|
delta_timestamps=delta_timestamps,
|
|
transform=transforms,
|
|
)
|
|
|
|
return dataset
|