Add gym-aloha, rename simxarm -> xarm, refactor

This commit is contained in:
Simon Alibert
2024-04-08 16:18:53 +02:00
parent 5dff6d8339
commit 3f6dfa4916
15 changed files with 91 additions and 97 deletions

View File

@@ -19,10 +19,10 @@ def make_dataset(
normalize=True,
stats_path=None,
):
if cfg.env.name == "simxarm":
from lerobot.common.datasets.simxarm import SimxarmDataset
if cfg.env.name == "xarm":
from lerobot.common.datasets.xarm import XarmDataset
clsfunc = SimxarmDataset
clsfunc = XarmDataset
elif cfg.env.name == "pusht":
from lerobot.common.datasets.pusht import PushtDataset

View File

@@ -24,7 +24,7 @@ def download(raw_dir):
zip_path.unlink()
class SimxarmDataset(torch.utils.data.Dataset):
class XarmDataset(torch.utils.data.Dataset):
available_datasets = [
"xarm_lift_medium",
]