This commit is contained in:
Alexander Soare
2024-03-27 18:33:48 +00:00
parent 120f0aef5c
commit b7c9c33072
10 changed files with 20 additions and 33 deletions

View File

@@ -2,8 +2,9 @@ import pytest
import torch
from lerobot.common.datasets.factory import make_offline_buffer
from lerobot.common.utils import init_hydra_config
from .utils import DEVICE, init_config
from .utils import DEVICE, DEFAULT_CONFIG_PATH
@pytest.mark.parametrize(
@@ -18,7 +19,10 @@ from .utils import DEVICE, init_config
],
)
def test_factory(env_name, dataset_id):
cfg = init_config(overrides=[f"env={env_name}", f"env.task={dataset_id}", f"device={DEVICE}"])
cfg = init_hydra_config(
DEFAULT_CONFIG_PATH,
overrides=[f"env={env_name}", f"env.task={dataset_id}", f"device={DEVICE}"]
)
offline_buffer = make_offline_buffer(cfg)
for key in offline_buffer.image_keys:
img = offline_buffer[0].get(key)

View File

@@ -7,8 +7,9 @@ from lerobot.common.datasets.factory import make_offline_buffer
from lerobot.common.envs.factory import make_env
from lerobot.common.envs.pusht.env import PushtEnv
from lerobot.common.envs.simxarm.env import SimxarmEnv
from lerobot.common.utils import init_hydra_config
from .utils import DEVICE, init_config
from .utils import DEVICE, DEFAULT_CONFIG_PATH
def print_spec_rollout(env):
@@ -89,7 +90,10 @@ def test_pusht(from_pixels, pixels_only):
],
)
def test_factory(env_name):
cfg = init_config(overrides=[f"env={env_name}", f"device={DEVICE}"])
cfg = init_hydra_config(
DEFAULT_CONFIG_PATH,
overrides=[f"env={env_name}", f"device={DEVICE}"],
)
offline_buffer = make_offline_buffer(cfg)

View File

@@ -1,4 +1,3 @@
from omegaconf import open_dict
import pytest
from tensordict import TensorDict
from tensordict.nn import TensorDictModule
@@ -10,8 +9,8 @@ from lerobot.common.policies.factory import make_policy
from lerobot.common.envs.factory import make_env
from lerobot.common.datasets.factory import make_offline_buffer
from lerobot.common.policies.abstract import AbstractPolicy
from .utils import DEVICE, init_config
from lerobot.common.utils import init_hydra_config
from .utils import DEVICE, DEFAULT_CONFIG_PATH
@pytest.mark.parametrize(
"env_name,policy_name,extra_overrides",
@@ -34,7 +33,8 @@ def test_concrete_policy(env_name, policy_name, extra_overrides):
- Updating the policy.
- Using the policy to select actions at inference time.
"""
cfg = init_config(
cfg = init_hydra_config(
DEFAULT_CONFIG_PATH,
overrides=[
f"env={env_name}",
f"policy={policy_name}",

View File

@@ -1,13 +1,6 @@
import os
import hydra
from hydra import compose, initialize
CONFIG_PATH = "../lerobot/configs"
# Pass this as the first argument to init_hydra_config.
DEFAULT_CONFIG_PATH = "lerobot/configs/default.yaml"
DEVICE = os.environ.get('LEROBOT_TESTS_DEVICE', "cuda")
def init_config(config_name="default", overrides=None):
hydra.core.global_hydra.GlobalHydra.instance().clear()
initialize(config_path=CONFIG_PATH)
cfg = compose(config_name=config_name, overrides=overrides)
return cfg