refactor: move env base modules to unilab/base

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

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

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

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

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
 

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)

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")

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()
 

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
 

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
 

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()

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")

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()
 

unilab/base/__init__.py

@@ -0,0 +1 @@
+"""Environment registry and base classes."""

unilab/envs/np_env.py → 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

unilab/envs/__init__.py

@@ -1 +1 @@
-"""Environment registry and base classes."""
+"""Environment"""

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

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
 
 

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

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

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
 

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

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

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 ────────────────────────────
 

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

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 ────────────────────────────