9d002032d1
WIP Aloha env tests pass Rendering works (fps look fast tho? TODO action bounding is too wide [-1,1]) Update README Copy past from act repo Remove download.py add a WIP for Simxarm Remove download.py add a WIP for Simxarm Add act yaml (TODO: try train.py) Training can runs (TODO: eval) Add tasks without end_effector that are compatible with dataset, Eval can run (TODO: training and pretrained model) Add AbstractEnv, Refactor AlohaEnv, Add rendering_hook in env, Minor modifications, (TODO: Refactor Pusht and Simxarm) poetry lock fix bug in compute_stats for action normalization fix more bugs in normalization fix training fix import PushtEnv inheriates AbstractEnv, Improve factory Normalization Add _make_env to EnvAbstract Add call_rendering_hooks to pusht env SimxarmEnv inherites from AbstractEnv (NOT TESTED) Add aloha tests artifacts + update pusht stats fix image normalization: before env was in [0,1] but dataset in [0,255], and now both in [0,255] Small fix on simxarm Add next to obs Add top camera to Aloha env (TODO: make it compatible with set of cameras) Add top camera to Aloha env (TODO: make it compatible with set of cameras)
129 lines
4.9 KiB
Python
129 lines
4.9 KiB
Python
import logging
|
|
import os
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
from torchrl.data.replay_buffers import PrioritizedSliceSampler, SliceSampler
|
|
|
|
from lerobot.common.envs.transforms import NormalizeTransform, Prod
|
|
|
|
DATA_DIR = Path(os.environ.get("DATA_DIR", "data"))
|
|
|
|
|
|
def make_offline_buffer(
|
|
cfg, overwrite_sampler=None, normalize=True, overwrite_batch_size=None, overwrite_prefetch=None
|
|
):
|
|
if cfg.policy.balanced_sampling:
|
|
assert cfg.online_steps > 0
|
|
batch_size = None
|
|
pin_memory = False
|
|
prefetch = None
|
|
else:
|
|
assert cfg.online_steps == 0
|
|
num_slices = cfg.policy.batch_size
|
|
batch_size = cfg.policy.horizon * num_slices
|
|
pin_memory = cfg.device == "cuda"
|
|
prefetch = cfg.prefetch
|
|
|
|
if overwrite_batch_size is not None:
|
|
batch_size = overwrite_batch_size
|
|
|
|
if overwrite_prefetch is not None:
|
|
prefetch = overwrite_prefetch
|
|
|
|
if overwrite_sampler is None:
|
|
# TODO(rcadene): move batch_size outside
|
|
num_traj_per_batch = cfg.policy.batch_size # // cfg.horizon
|
|
# TODO(rcadene): Sampler outputs a batch_size <= cfg.batch_size.
|
|
# We would need to add a transform to pad the tensordict to ensure batch_size == cfg.batch_size.
|
|
|
|
if cfg.offline_prioritized_sampler:
|
|
logging.info("use prioritized sampler for offline dataset")
|
|
sampler = PrioritizedSliceSampler(
|
|
max_capacity=100_000,
|
|
alpha=cfg.policy.per_alpha,
|
|
beta=cfg.policy.per_beta,
|
|
num_slices=num_traj_per_batch,
|
|
strict_length=False,
|
|
)
|
|
else:
|
|
logging.info("use simple sampler for offline dataset")
|
|
sampler = SliceSampler(
|
|
num_slices=num_traj_per_batch,
|
|
strict_length=False,
|
|
)
|
|
else:
|
|
sampler = overwrite_sampler
|
|
|
|
if cfg.env.name == "simxarm":
|
|
from lerobot.common.datasets.simxarm import SimxarmExperienceReplay
|
|
|
|
clsfunc = SimxarmExperienceReplay
|
|
dataset_id = f"xarm_{cfg.env.task}_medium"
|
|
|
|
elif cfg.env.name == "pusht":
|
|
from lerobot.common.datasets.pusht import PushtExperienceReplay
|
|
|
|
clsfunc = PushtExperienceReplay
|
|
dataset_id = "pusht"
|
|
|
|
elif cfg.env.name == "aloha":
|
|
from lerobot.common.datasets.aloha import AlohaExperienceReplay
|
|
|
|
clsfunc = AlohaExperienceReplay
|
|
dataset_id = f"aloha_{cfg.env.task}"
|
|
else:
|
|
raise ValueError(cfg.env.name)
|
|
|
|
offline_buffer = clsfunc(
|
|
dataset_id=dataset_id,
|
|
root=DATA_DIR,
|
|
sampler=sampler,
|
|
batch_size=batch_size,
|
|
pin_memory=pin_memory,
|
|
prefetch=prefetch if isinstance(prefetch, int) else None,
|
|
)
|
|
|
|
if cfg.policy.name == "tdmpc":
|
|
img_keys = []
|
|
for key in offline_buffer.image_keys:
|
|
img_keys.append(("next", *key))
|
|
img_keys += offline_buffer.image_keys
|
|
else:
|
|
img_keys = offline_buffer.image_keys
|
|
|
|
transforms = [Prod(in_keys=img_keys, prod=1 / 255)]
|
|
|
|
if normalize:
|
|
# TODO(rcadene): make normalization strategy configurable between mean_std, min_max, manual_min_max, min_max_from_spec
|
|
stats = offline_buffer.compute_or_load_stats()
|
|
|
|
# we only normalize the state and action, since the images are usually normalized inside the model for now (except for tdmpc: see the following)
|
|
in_keys = [("observation", "state"), ("action")]
|
|
|
|
if cfg.policy.name == "tdmpc":
|
|
# TODO(rcadene): we add img_keys to the keys to normalize for tdmpc only, since diffusion and act policies normalize the image inside the model for now
|
|
in_keys += img_keys
|
|
# TODO(racdene): since we use next observations in tdmpc, we also add them to the normalization. We are wasting a bit of compute on this for now.
|
|
in_keys += [("next", *key) for key in img_keys]
|
|
in_keys.append(("next", "observation", "state"))
|
|
|
|
if cfg.policy.name == "diffusion" and cfg.env.name == "pusht":
|
|
# TODO(rcadene): we overwrite stats to have the same as pretrained model, but we should remove this
|
|
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", "min"] = torch.tensor([12.0, 25.0], dtype=torch.float32)
|
|
stats["action", "max"] = torch.tensor([511.0, 511.0], dtype=torch.float32)
|
|
|
|
# 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"
|
|
transforms.append(NormalizeTransform(stats, in_keys, mode=normalization_mode))
|
|
|
|
offline_buffer.set_transform(transforms)
|
|
|
|
if not overwrite_sampler:
|
|
index = torch.arange(0, offline_buffer.num_samples, 1)
|
|
sampler.extend(index)
|
|
|
|
return offline_buffer
|