diff --git a/benchmark/benchmark_postprocess.py b/benchmark/benchmark_postprocess.py index bf20b9b03..a2768d03c 100644 --- a/benchmark/benchmark_postprocess.py +++ b/benchmark/benchmark_postprocess.py @@ -15,7 +15,7 @@ import torch import mlx.core as mx -from unilab.envs import registry +from unilab.base import registry try: from mujoco import mlx_step as mj_mlx_step diff --git a/scripts/play_interactive.py b/scripts/play_interactive.py index 2fda6a6c1..557e3ad6c 100644 --- a/scripts/play_interactive.py +++ b/scripts/play_interactive.py @@ -46,7 +46,7 @@ def ensure_registries(): ensure_registries() -from unilab.envs import registry +from unilab.base import registry from unilab.config import locomotion_params from unilab.utils.rsl_rl_compat import is_rsl_rl_v4, convert_config_v3_to_v4 from unilab.utils.run_utils import get_latest_run diff --git a/scripts/train_appo.py b/scripts/train_appo.py index 512b485ee..80f62922c 100644 --- a/scripts/train_appo.py +++ b/scripts/train_appo.py @@ -45,7 +45,7 @@ def play_appo(args, rl_cfg): import numpy as np import mediapy as media from tensordict import TensorDict - from unilab.envs import registry + from unilab.base import registry from unilab.utils import render_many from rsl_rl.utils import resolve_callable from unilab.utils.rsl_rl_compat import is_rsl_rl_v4, convert_config_v3_to_v4 diff --git a/scripts/train_mlx_ppo.py b/scripts/train_mlx_ppo.py index 0e45e2755..2388c0615 100644 --- a/scripts/train_mlx_ppo.py +++ b/scripts/train_mlx_ppo.py @@ -46,7 +46,7 @@ def ensure_registries() -> None: ensure_registries() from unilab.config import locomotion_params -from unilab.envs import registry +from unilab.base import registry from unilab.utils import render_many from unilab.utils.onpolicy_logger import OnPolicyLogger from unilab.algos.mlx.common import EmpiricalDiscountedVariationNormalization, RolloutBuffer @@ -159,7 +159,7 @@ def play_mlx_ppo(args, cfg, dtype, use_fp16, resolved_sim_backend, task_log_root """Play mode for MLX PPO.""" import mlx.core as mx import numpy as np - from unilab.envs import registry + from unilab.base import registry from unilab.utils import render_many play_model_dtype = mx.float32 if use_fp16 else dtype diff --git a/scripts/train_offpolicy.py b/scripts/train_offpolicy.py index 92e78c842..56c79730f 100644 --- a/scripts/train_offpolicy.py +++ b/scripts/train_offpolicy.py @@ -131,7 +131,7 @@ def play_offpolicy(algo_name: str, args, cfg) -> None: import mediapy as media import numpy as np import torch - from unilab.envs import registry + from unilab.base import registry from unilab.utils import render_many from unilab.utils.algo_utils import build_actor diff --git a/scripts/train_rsl_rl.py b/scripts/train_rsl_rl.py index ae80f4dbc..f770de0b4 100644 --- a/scripts/train_rsl_rl.py +++ b/scripts/train_rsl_rl.py @@ -31,7 +31,7 @@ def ensure_registries(): ensure_registries() -from unilab.envs import registry +from unilab.base import registry from unilab.config import locomotion_params, manipulation_params from unilab.utils import render_many from unilab.utils.torch_utils import to_torch, to_numpy @@ -193,7 +193,7 @@ def step(self, actions): def play_rsl_rl(args, cfg, device): """Play mode for RSL-RL.""" import torch - from unilab.envs import registry + from unilab.base import registry from unilab.utils.torch_utils import to_numpy env = registry.make(args.task, num_envs=args.play_env_num, sim_backend=args.sim_backend) diff --git a/scripts/visualization_env.py b/scripts/visualization_env.py index f04c5997b..43d887353 100644 --- a/scripts/visualization_env.py +++ b/scripts/visualization_env.py @@ -8,7 +8,7 @@ import unilab.envs.locomotion.g1 import unilab.envs.locomotion.go1 import unilab.envs.locomotion.go2 -from unilab.envs import registry +from unilab.base import registry parser = argparse.ArgumentParser(description="Visualize robot tasks") parser.add_argument("--task", type=str, default="Go1JoystickFlatTerrain", help="Task name") diff --git a/unilab/algos/torch/appo/runner.py b/unilab/algos/torch/appo/runner.py index 24935e98b..1294bbda6 100644 --- a/unilab/algos/torch/appo/runner.py +++ b/unilab/algos/torch/appo/runner.py @@ -61,7 +61,7 @@ def _resolve_dims(self): def _detect_dims(self): """Create a tiny env to read obs/action dims, then close it.""" - from unilab.envs import registry + from unilab.base import registry from unilab.utils.algo_utils import ensure_registries ensure_registries() diff --git a/unilab/algos/torch/appo/worker.py b/unilab/algos/torch/appo/worker.py index 145e022aa..5fd5e1d2d 100644 --- a/unilab/algos/torch/appo/worker.py +++ b/unilab/algos/torch/appo/worker.py @@ -46,7 +46,7 @@ def appo_collector_fn( Creates environment + policy, collects rollouts, writes to SharedOnPolicyStorage. """ from unilab.algos.torch.common.async_runner import SharedOnPolicyStorage, SharedWeightSync - from unilab.envs import registry + from unilab.base import registry from tensordict import TensorDict from unilab.utils.rsl_rl_compat import convert_config_v3_to_v4, is_rsl_rl_v4 diff --git a/unilab/algos/torch/fast_sac/runner.py b/unilab/algos/torch/fast_sac/runner.py index 1bb0c2ce1..4b1a8fce1 100644 --- a/unilab/algos/torch/fast_sac/runner.py +++ b/unilab/algos/torch/fast_sac/runner.py @@ -37,7 +37,7 @@ def __init__( use_gpu_buffer: bool = True, use_symmetry: bool = False, ): - from unilab.envs import registry + from unilab.base import registry from unilab.utils.algo_utils import ensure_registries import torch diff --git a/unilab/algos/torch/fast_td3/runner.py b/unilab/algos/torch/fast_td3/runner.py index a7993975c..33b8fd33b 100644 --- a/unilab/algos/torch/fast_td3/runner.py +++ b/unilab/algos/torch/fast_td3/runner.py @@ -101,7 +101,7 @@ def _default_device() -> str: @staticmethod def _detect_obs_action_dims(env_name: str, sim_backend: str = "mujoco") -> tuple[int, int]: - from unilab.envs import registry + from unilab.base import registry from unilab.utils.algo_utils import ensure_registries ensure_registries() diff --git a/unilab/algos/torch/offpolicy/runner.py b/unilab/algos/torch/offpolicy/runner.py index e7a4178e5..2351d7848 100644 --- a/unilab/algos/torch/offpolicy/runner.py +++ b/unilab/algos/torch/offpolicy/runner.py @@ -63,7 +63,7 @@ def __init__( self.obs_dim, self.action_dim = self._detect_dims() def _detect_dims(self): - from unilab.envs import registry + from unilab.base import registry from unilab.utils.algo_utils import ensure_registries ensure_registries() env = registry.make(self.env_name, num_envs=1, sim_backend="mujoco") diff --git a/unilab/algos/torch/offpolicy/worker.py b/unilab/algos/torch/offpolicy/worker.py index 5d62f6e10..88263a9e1 100644 --- a/unilab/algos/torch/offpolicy/worker.py +++ b/unilab/algos/torch/offpolicy/worker.py @@ -84,7 +84,7 @@ def _run_collector( obs_normalization, shared_obs_normalizer_stats, sim_backend ): from unilab.ipc import SharedReplayBuffer, SharedWeightSync - from unilab.envs import registry + from unilab.base import registry ensure_registries() diff --git a/unilab/base/__init__.py b/unilab/base/__init__.py new file mode 100644 index 000000000..998dc8c51 --- /dev/null +++ b/unilab/base/__init__.py @@ -0,0 +1 @@ +"""Environment registry and base classes.""" diff --git a/unilab/envs/backend/__init__.py b/unilab/base/backend/__init__.py similarity index 100% rename from unilab/envs/backend/__init__.py rename to unilab/base/backend/__init__.py diff --git a/unilab/envs/backend/base.py b/unilab/base/backend/base.py similarity index 100% rename from unilab/envs/backend/base.py rename to unilab/base/backend/base.py diff --git a/unilab/envs/backend/motrix_backend.py b/unilab/base/backend/motrix_backend.py similarity index 100% rename from unilab/envs/backend/motrix_backend.py rename to unilab/base/backend/motrix_backend.py diff --git a/unilab/envs/backend/motrix_numba_backend.py b/unilab/base/backend/motrix_numba_backend.py similarity index 100% rename from unilab/envs/backend/motrix_numba_backend.py rename to unilab/base/backend/motrix_numba_backend.py diff --git a/unilab/envs/backend/mujoco_backend.py b/unilab/base/backend/mujoco_backend.py similarity index 100% rename from unilab/envs/backend/mujoco_backend.py rename to unilab/base/backend/mujoco_backend.py diff --git a/unilab/envs/base.py b/unilab/base/base.py similarity index 100% rename from unilab/envs/base.py rename to unilab/base/base.py diff --git a/unilab/envs/curriculum.py b/unilab/base/curriculum.py similarity index 100% rename from unilab/envs/curriculum.py rename to unilab/base/curriculum.py diff --git a/unilab/envs/dtype_config.py b/unilab/base/dtype_config.py similarity index 100% rename from unilab/envs/dtype_config.py rename to unilab/base/dtype_config.py diff --git a/unilab/envs/np_env.py b/unilab/base/np_env.py similarity index 97% rename from unilab/envs/np_env.py rename to unilab/base/np_env.py index 24bf7a21b..dd67105b7 100644 --- a/unilab/envs/np_env.py +++ b/unilab/base/np_env.py @@ -6,9 +6,9 @@ import gymnasium as gym from typing import Optional -from unilab.envs.base import ABEnv, EnvCfg -from unilab.envs.backend import SimBackend -from unilab.envs.dtype_config import get_global_dtype +from unilab.base.base import ABEnv, EnvCfg +from unilab.base.backend import SimBackend +from unilab.base.dtype_config import get_global_dtype @dataclass diff --git a/unilab/envs/registry.py b/unilab/base/registry.py similarity index 100% rename from unilab/envs/registry.py rename to unilab/base/registry.py diff --git a/unilab/envs/__init__.py b/unilab/envs/__init__.py index 998dc8c51..85a4be3f6 100644 --- a/unilab/envs/__init__.py +++ b/unilab/envs/__init__.py @@ -1 +1 @@ -"""Environment registry and base classes.""" +"""Environment""" diff --git a/unilab/envs/locomotion/g1/base.py b/unilab/envs/locomotion/g1/base.py index abebdcf9d..0e81a2841 100644 --- a/unilab/envs/locomotion/g1/base.py +++ b/unilab/envs/locomotion/g1/base.py @@ -5,9 +5,9 @@ import numpy as np from dataclasses import dataclass, field -from unilab.envs.base import EnvCfg -from unilab.envs.np_env import NpEnv, NpEnvState -from unilab.envs.backend import SimBackend +from unilab.base.base import EnvCfg +from unilab.base.np_env import NpEnv, NpEnvState +from unilab.base.backend import SimBackend @dataclass diff --git a/unilab/envs/locomotion/g1/joystick.py b/unilab/envs/locomotion/g1/joystick.py index 909aea634..ea6bc019f 100644 --- a/unilab/envs/locomotion/g1/joystick.py +++ b/unilab/envs/locomotion/g1/joystick.py @@ -7,11 +7,11 @@ import math import numpy as np -from unilab.envs import registry -from unilab.envs.np_env import NpEnvState -from unilab.envs.backend import create_backend +from unilab.base import registry +from unilab.base.np_env import NpEnvState +from unilab.base.backend import create_backend from unilab.utils.math_utils import np_quat_mul, np_yaw_to_quat -from unilab.envs.dtype_config import get_global_dtype +from unilab.base.dtype_config import get_global_dtype from unilab.envs.locomotion.g1.base import G1BaseCfg, G1BaseEnv diff --git a/unilab/envs/locomotion/g1/joystick_sac.py b/unilab/envs/locomotion/g1/joystick_sac.py index 60aaa3126..25c34568a 100644 --- a/unilab/envs/locomotion/g1/joystick_sac.py +++ b/unilab/envs/locomotion/g1/joystick_sac.py @@ -5,12 +5,12 @@ from etils import epath import numpy as np -from unilab.envs import registry -from unilab.envs.backend import create_backend -from unilab.envs.dtype_config import get_global_dtype +from unilab.base import registry +from unilab.base.backend import create_backend +from unilab.base.dtype_config import get_global_dtype from unilab.envs.locomotion.g1.base import G1BaseCfg, G1BaseEnv from unilab.envs.locomotion.g1.joystick import G1JoystickPPO, InitState -from unilab.envs.curriculum import EpisodeLengthTracker, PenaltyCurriculum +from unilab.base.curriculum import EpisodeLengthTracker, PenaltyCurriculum @dataclass diff --git a/unilab/envs/locomotion/go1/base.py b/unilab/envs/locomotion/go1/base.py index 205e70081..3b8fbf03e 100644 --- a/unilab/envs/locomotion/go1/base.py +++ b/unilab/envs/locomotion/go1/base.py @@ -5,10 +5,10 @@ import numpy as np from dataclasses import dataclass, field -from unilab.envs.base import EnvCfg -from unilab.envs.np_env import NpEnv, NpEnvState -from unilab.envs.backend import SimBackend -from unilab.envs.dtype_config import get_global_dtype +from unilab.base.base import EnvCfg +from unilab.base.np_env import NpEnv, NpEnvState +from unilab.base.backend import SimBackend +from unilab.base.dtype_config import get_global_dtype @dataclass diff --git a/unilab/envs/locomotion/go1/joystick.py b/unilab/envs/locomotion/go1/joystick.py index 3b7ef3259..c2589fd64 100644 --- a/unilab/envs/locomotion/go1/joystick.py +++ b/unilab/envs/locomotion/go1/joystick.py @@ -5,11 +5,11 @@ import numpy as np from dataclasses import dataclass, field -from unilab.envs import registry -from unilab.envs.np_env import NpEnvState -from unilab.envs.backend import create_backend +from unilab.base import registry +from unilab.base.np_env import NpEnvState +from unilab.base.backend import create_backend from unilab.utils.math_utils import np_quat_mul, np_yaw_to_quat -from unilab.envs.dtype_config import get_global_dtype +from unilab.base.dtype_config import get_global_dtype from unilab.envs.locomotion.go1.base import Go1BaseEnv, Go1BaseCfg diff --git a/unilab/envs/locomotion/go2/base.py b/unilab/envs/locomotion/go2/base.py index 11e1a4e06..f6b926d60 100644 --- a/unilab/envs/locomotion/go2/base.py +++ b/unilab/envs/locomotion/go2/base.py @@ -5,9 +5,9 @@ import numpy as np from dataclasses import dataclass, field -from unilab.envs.base import EnvCfg -from unilab.envs.np_env import NpEnv, NpEnvState -from unilab.envs.backend import SimBackend +from unilab.base.base import EnvCfg +from unilab.base.np_env import NpEnv, NpEnvState +from unilab.base.backend import SimBackend @dataclass diff --git a/unilab/envs/locomotion/go2/joystick.py b/unilab/envs/locomotion/go2/joystick.py index 445b18aad..83ef3efdb 100644 --- a/unilab/envs/locomotion/go2/joystick.py +++ b/unilab/envs/locomotion/go2/joystick.py @@ -5,9 +5,9 @@ import numpy as np from dataclasses import dataclass, field -from unilab.envs import registry -from unilab.envs.np_env import NpEnvState -from unilab.envs.backend import create_backend +from unilab.base import registry +from unilab.base.np_env import NpEnvState +from unilab.base.backend import create_backend from unilab.utils.math_utils import np_quat_mul, np_yaw_to_quat from unilab.envs.locomotion.go2.base import Go2BaseEnv, Go2BaseCfg diff --git a/unilab/envs/manipulation/inhand_rot_allegro/base.py b/unilab/envs/manipulation/inhand_rot_allegro/base.py index f80aa1770..6aee2ed70 100644 --- a/unilab/envs/manipulation/inhand_rot_allegro/base.py +++ b/unilab/envs/manipulation/inhand_rot_allegro/base.py @@ -5,10 +5,10 @@ import numpy as np from dataclasses import dataclass, field -from unilab.envs.base import EnvCfg -from unilab.envs.np_env import NpEnv, NpEnvState -from unilab.envs.backend import SimBackend -from unilab.envs.dtype_config import get_global_dtype +from unilab.base.base import EnvCfg +from unilab.base.np_env import NpEnv, NpEnvState +from unilab.base.backend import SimBackend +from unilab.base.dtype_config import get_global_dtype # ─────────────────────────── Configuration ──────────────────────────── diff --git a/unilab/envs/manipulation/inhand_rot_allegro/gen_grasp.py b/unilab/envs/manipulation/inhand_rot_allegro/gen_grasp.py index 9f403fd4b..98a3cfc91 100644 --- a/unilab/envs/manipulation/inhand_rot_allegro/gen_grasp.py +++ b/unilab/envs/manipulation/inhand_rot_allegro/gen_grasp.py @@ -52,9 +52,9 @@ def ensure_registries(): ensure_registries() -from unilab.envs import registry # noqa: E402 (after sys.path setup) +from unilab.base import registry # noqa: E402 (after sys.path setup) from unilab.utils import render_many # noqa: E402 -from unilab.envs.dtype_config import get_global_dtype # noqa: E402 +from unilab.base.dtype_config import get_global_dtype # noqa: E402 # Explicit import to guarantee the @registry.env decorator runs, # since ensure_registries() silently swallows import errors. from unilab.envs.manipulation.inhand_rot_allegro import rotation as _rotation_register diff --git a/unilab/envs/manipulation/inhand_rot_allegro/rotation.py b/unilab/envs/manipulation/inhand_rot_allegro/rotation.py index 265de9c10..ca2ccca83 100644 --- a/unilab/envs/manipulation/inhand_rot_allegro/rotation.py +++ b/unilab/envs/manipulation/inhand_rot_allegro/rotation.py @@ -48,9 +48,9 @@ def _quat_to_axis_angle(q: np.ndarray) -> np.ndarray: ) return xyz / sin_half_over_angle # (N, 3) -from unilab.envs import registry -from unilab.envs.np_env import NpEnvState -from unilab.envs.backend import create_backend +from unilab.base import registry +from unilab.base.np_env import NpEnvState +from unilab.base.backend import create_backend from unilab.envs.manipulation.inhand_rot_allegro.base import AllegroBaseCfg, AllegroBaseMjEnv # ─────────────────────────── Configuration ────────────────────────────