Merge pull request #44 from Mingrui-Yu/main

APPO works in AllegroInhandRotation

Author: TATP-233

pyproject.toml

@@ -25,6 +25,7 @@ dependencies = [
     "rich",
     "typing-extensions",
     "mlx ; sys_platform == 'darwin'",
+    "lark>=1.3.1",
 ]
 
 [project.optional-dependencies]

scripts/train_appo.py

@@ -15,31 +15,21 @@
 
 
 def ensure_registries():
-    try:
-        import unilab.envs.locomotion
-
-        package = unilab.envs.locomotion
-        if hasattr(package, "__path__"):
-            for _, name, ispkg in pkgutil.walk_packages(package.__path__, package.__name__ + "."):
-                try:
-                    importlib.import_module(name)
-                except Exception:
-                    pass
-    except ImportError:
-        pass
-
-    try:
-        import unilab.envs.locomotion.walking
-
-        package = unilab.envs.locomotion.walking
-        if hasattr(package, "__path__"):
-            for _, name, ispkg in pkgutil.walk_packages(package.__path__, package.__name__ + "."):
-                try:
-                    importlib.import_module(name)
-                except Exception:
-                    pass
-    except ImportError:
-        pass
+    for pkg_name in (
+        "unilab.envs.locomotion",
+        "unilab.envs.locomotion.walking",
+        "unilab.envs.manipulation",
+    ):
+        try:
+            package = importlib.import_module(pkg_name)
+            if hasattr(package, "__path__"):
+                for _, name, _ in pkgutil.walk_packages(package.__path__, package.__name__ + "."):
+                    try:
+                        importlib.import_module(name)
+                    except Exception:
+                        pass
+        except ImportError:
+            pass
 
 
 def play_appo(args, rl_cfg):
@@ -219,9 +209,12 @@ def main():
 
     ensure_registries()
 
-    from unilab.config.locomotion_params import appo_config
+    from unilab.config import locomotion_params, manipulation_params
 
-    rl_cfg = appo_config(args.task)
+    if args.task in manipulation_params.KNOWN_TASKS:
+        rl_cfg = manipulation_params.appo_config(args.task)
+    else:
+        rl_cfg = locomotion_params.appo_config(args.task)
 
     # Fall back to config values for args not explicitly provided on the CLI.
     if args.max_iterations is None:

src/unilab/algos/torch/appo/runner.py

@@ -179,7 +179,7 @@ def learn(
 
         weight_param_shapes = {name: p.shape for name, p in learner.actor.state_dict().items()}
 
-        metrics_queue: mp.Queue = mp.Queue(maxsize=100)
+        metrics_queue: mp.Queue = mp.get_context("spawn").Queue(maxsize=100)
 
         # Start collector
         collector_kwargs = {

src/unilab/config/locomotion_params.py

@@ -115,6 +115,7 @@ def appo_config(env_name: str) -> config_dict.ConfigDict:
         steps_per_env=24,
         max_iterations=150,
         save_interval=50,
+        empirical_normalization=False,
     )
 
     if env_name == "Go1JoystickFlatTerrain":
@@ -133,13 +134,15 @@ def appo_config(env_name: str) -> config_dict.ConfigDict:
         steps_per_env=cfg.steps_per_env,
         max_iterations=cfg.max_iterations,
         save_interval=cfg.save_interval,
+        empirical_normalization=cfg.empirical_normalization,
         obs_groups=config_dict.create(
             actor=config_dict.create(policy=0),
         ),
         actor=config_dict.create(
             class_name="rsl_rl.models.MLPModel",
             hidden_dims=[512, 256, 128],
             activation="elu",
+            obs_normalization=cfg.empirical_normalization,
             distribution_cfg=config_dict.create(
                 class_name="rsl_rl.modules.distribution.GaussianDistribution",
                 init_std=1.0,
@@ -150,6 +153,7 @@ def appo_config(env_name: str) -> config_dict.ConfigDict:
             class_name="rsl_rl.models.MLPModel",
             hidden_dims=[512, 256, 128],
             activation="elu",
+            obs_normalization=cfg.empirical_normalization,
         ),
         algorithm=config_dict.create(
             num_learning_epochs=5,

src/unilab/config/manipulation_params.py

@@ -90,6 +90,86 @@ def ppo_config(env_name: str) -> config_dict.ConfigDict:
     )
 
 
+def appo_config(env_name: str) -> config_dict.ConfigDict:
+    """Return APPO config.
+
+    Mirrors the structure of ppo_config/offpolicy_config so training
+    infrastructure is reusable.  Algorithm hyperparams live under the
+    ``algorithm`` sub-dict; network architecture under ``actor``/``critic``.
+    """
+    cfg = config_dict.create(
+        seed=1,
+        num_envs=2048,
+        steps_per_env=24,
+        max_iterations=150,
+        save_interval=50,
+        value_loss_coef=1.0,
+        desired_kl=0.01,
+        empirical_normalization=False,
+    )
+
+    if env_name == "AllegroInhandRotation":
+        cfg.num_envs = 16384
+        cfg.max_iterations = 501
+        cfg.steps_per_env = 8
+        cfg.value_loss_coef = 4.0
+        cfg.desired_kl = 0.02
+        cfg.save_interval = 100
+        cfg.empirical_normalization = True
+    elif env_name == "AllegroInhandRotationSac":
+        raise NotImplementedError()
+
+    return config_dict.create(
+        algo="appo",
+        algo_log_name="appo",
+        seed=cfg.seed,
+        num_envs=cfg.num_envs,
+        steps_per_env=cfg.steps_per_env,
+        max_iterations=cfg.max_iterations,
+        save_interval=cfg.save_interval,
+        empirical_normalization=cfg.empirical_normalization,
+        obs_groups=config_dict.create(
+            actor=config_dict.create(policy=0),
+        ),
+        actor=config_dict.create(
+            class_name="rsl_rl.models.MLPModel",
+            hidden_dims=[512, 256, 128],
+            activation="elu",
+            obs_normalization=cfg.empirical_normalization,
+            distribution_cfg=config_dict.create(
+                class_name="rsl_rl.modules.distribution.GaussianDistribution",
+                init_std=1.0,
+                std_type="scalar",
+            ),
+        ),
+        critic=config_dict.create(
+            class_name="rsl_rl.models.MLPModel",
+            hidden_dims=[512, 256, 128],
+            activation="elu",
+            obs_normalization=cfg.empirical_normalization,
+        ),
+        algorithm=config_dict.create(
+            num_learning_epochs=5,
+            num_mini_batches=4,
+            clip_param=0.2,
+            gamma=0.99,
+            lam=0.95,
+            value_loss_coef=cfg.value_loss_coef,
+            entropy_coef=0.01,
+            learning_rate=1e-3,
+            max_grad_norm=1.0,
+            use_clipped_value_loss=True,
+            schedule="adaptive",
+            desired_kl=cfg.desired_kl,
+            optimizer="adam",
+            tau=1.0,
+            target_update_freq=1,
+            vtrace_clip_rho=1.0,
+            vtrace_clip_c=1.0,
+        ),
+    )
+
+
 def offpolicy_config(algo: str, env_name: str) -> config_dict.ConfigDict:
     """Return a unified off-policy config schema for SAC/TD3.
 

src/unilab/envs/manipulation/inhand_rot_allegro/rotation.py

@@ -36,6 +36,17 @@ class RewardConfig:
     reset_z_threshold: float = 0.125  # metres
 
 
+@dataclass
+class obs_cfg:
+    obs_dict = {
+        "dof_pos_norm": 16,
+        "targets": 16,
+        "ball_pos": 3,
+    }  # 'obs_name': dim
+    actor_obs = ["dof_pos_norm", "targets", "ball_pos"]
+    num_lag_steps: int = 3
+
+
 @dataclass
 class DomainRandConfig:
     # Small random perturbation added to each joint at reset (rad).
@@ -53,6 +64,7 @@ class AllegroRotationCfg(AllegroBaseCfg):
     max_episode_seconds: float = 20.0  # same as HORA
     reward_config: RewardConfig = field(default_factory=RewardConfig)
     domain_rand: DomainRandConfig = field(default_factory=DomainRandConfig)
+    obs_config: obs_cfg = field(default_factory=obs_cfg)
 
     # World-frame unit vector around which to reward rotation.
     # Default (+z) = counterclockwise when viewed from above.
@@ -96,10 +108,27 @@ def __init__(self, cfg: AllegroRotationCfg, num_envs: int = 1, backend_type: str
     # ── Spaces ──────────────────────────────────────────────────────
 
     def _init_obs_space(self):
-        # Last 3 timesteps of [dof_pos_norm(16) + targets(16) + ball_pos(3)] = 3 * 35 = 105
+        obs_cfg = self._cfg.obs_config
+        num_obs_per_step = sum(obs_cfg.obs_dict.values())
+        num_obs = obs_cfg.num_lag_steps * num_obs_per_step
         self._observation_space = gym.spaces.Box(
-            low=-np.inf, high=np.inf, shape=(105,), dtype=float
+            low=-np.inf, high=np.inf, shape=(num_obs,), dtype=float
         )
+        self.actor_indices = self._get_actor_indices()
+
+    def _get_actor_indices(self):
+        s = 0
+        indices = []
+        for i in self._cfg.obs_config.actor_obs:
+            s = 0
+            for k, v in self._cfg.obs_config.obs_dict.items():
+                if k == i:
+                    for q in range(s, s + v):
+                        indices.append(q)
+                    s += v
+                    break
+                s += v
+        return indices
 
     @property
     def observation_space(self) -> gym.spaces.Box:
@@ -212,7 +241,15 @@ def update_state(self, state: NpEnvState) -> NpEnvState:
         current_obs = np.concatenate([dof_pos_norm, targets, ball_pos], axis=1)  # (N, 35)
 
         obs_lag_history = state.info.get(
-            "obs_lag_history", np.zeros((self._num_envs, 3, 35), dtype=self._np_dtype)
+            "obs_lag_history",
+            np.zeros(
+                (
+                    self._num_envs,
+                    self._cfg.obs_config.num_lag_steps,
+                    sum(self._cfg.obs_config.obs_dict.values()),
+                ),
+                dtype=self._np_dtype,
+            ),
         )
         obs_lag_history[:, :-1] = obs_lag_history[:, 1:]
         obs_lag_history[:, -1] = current_obs
@@ -249,8 +286,11 @@ def _compute_rewards(self, state: NpEnvState) -> tuple:
 
     def _get_obs(self, info: dict) -> np.ndarray:
         """Extract observation from info. Pure function - does not modify state."""
+        num_obs_per_step = sum(self._cfg.obs_config.obs_dict.values())
+        num_lag_steps = self._cfg.obs_config.num_lag_steps
         obs_lag_history = info.get(
-            "obs_lag_history", np.zeros((self._num_envs, 3, 35), dtype=self._np_dtype)
+            "obs_lag_history",
+            np.zeros((self._num_envs, num_lag_steps, num_obs_per_step), dtype=self._np_dtype),
         )
         num_envs = obs_lag_history.shape[0]
         return np.asarray(obs_lag_history.reshape(num_envs, -1))
@@ -323,8 +363,13 @@ def reset(self, env_indices: np.ndarray) -> Tuple[np.ndarray, np.ndarray, dict]:
             [dof_pos_norm, init_ctrl, ball_pos_f32], axis=1
         )  # (num_reset, 35)
         obs_lag_history = np.broadcast_to(
-            init_obs[:, None, :], (num_reset, 3, 35)
-        ).copy()  # (num_reset, 3, 35)
+            init_obs[:, None, :],
+            (
+                num_reset,
+                self._cfg.obs_config.num_lag_steps,
+                sum(self._cfg.obs_config.obs_dict.values()),
+            ),
+        ).copy()  # (num_reset, num_lag_steps, num_obs_per_step)
 
         info = {
             "current_actions": np.zeros((num_reset, self._num_action), dtype=self._np_dtype),

src/unilab/utils/render_many.py

@@ -28,6 +28,9 @@ def _resolve_gl_backend() -> str:
         import ctypes
 
         ctypes.CDLL("libEGL.so.1")
+        # NVIDIA EGL requires a device ID to initialize offscreen contexts in
+        # spawned subprocesses; default to GPU 0 if the user hasn't set it.
+        os.environ.setdefault("MUJOCO_EGL_DEVICE_ID", "0")
         return "egl"
     except OSError:
         return "glfw"