Support real SO-101 arm training and fix some bugs

docs/INSTALL.md

@@ -11,6 +11,7 @@
 | 🧪 IsaacLab / Isaac Sim | [Install Isaac Sim / IsaacLab](#isaaclab--isaac-sim) |
 | 🦾 Franka real | Set up the [realtime kernel](https://frankarobotics.github.io/docs/libfranka/docs/real_time_kernel.html), then install `mpail2` with `".[franka]"` |
 | 🤖 Kinova real | Install [ROS 2](https://docs.ros.org/en/humble/Installation/Ubuntu-Install-Debs.html) and [`ros2_kortex`](https://github.com/Kinovarobotics/ros2_kortex) |
+| 🦿 SO-101 real | Install `mpail2` with `".[so101]"`, plus a separate [`lerobot`](https://github.com/huggingface/lerobot) clone — see [`mpail2/envs/real/so101/README.md`](../mpail2/envs/real/so101/README.md) |
 
 ## Default Install
 
@@ -31,6 +32,7 @@ IsaacLab and the real-robot environments require extra setup:
 - IsaacLab / Isaac Sim: install Isaac Sim / IsaacLab first, then install `mpail2`
 - Franka real robot: install `mpail2` with `".[franka]"`
 - Kinova real robot: install ROS 2 / `ros2_kortex`
+- SO-101 real robot: install `mpail2` with `".[so101]"` — see [`mpail2/envs/real/so101/README.md`](../mpail2/envs/real/so101/README.md) for the full setup (needs a separate `lerobot` clone + conda env)
 
 ## Gym / MuJoCo
 

mpail2/configs/cfgs.py

@@ -189,6 +189,13 @@ class RewardCfg:
 
     model_factory: Callable[[Any], torch.nn.Sequential] = mlp_factory
 
+    reward_clip: float = None
+    '''Clamp raw reward output to [-reward_clip, reward_clip]. None disables clamping.
+    The WGAN-style critic has no inherent output bound, and without one its scale can
+    drift arbitrarily over training (observed: mean_demo_reward climbing from ~6 to ~16
+    across iterations, Value/mean_q_value swinging from -187 to +164) — this bounds how
+    far that drift can propagate into the TD target / value loss.'''
+
 
 @dataclass(kw_only=True)
 class EnsembleValueCfg:
@@ -266,6 +273,15 @@ class DynamicsLearnerCfg:
     enc_lr_scale: float = None
     '''Scaling factor for the encoder learning rate relative to the dynamics model learning rate.'''
 
+    sigreg_coeff: float = None
+    '''Coefficient for SIGReg auxiliary regularization loss (see mpail2.dynamics.SIGReg). Disabled when None or <= 0.'''
+
+    sigreg_knots: int = None
+    '''Number of integration knots used by SIGReg.'''
+
+    sigreg_num_proj: int = None
+    '''Number of random projection vectors used by SIGReg.'''
+
 
 @dataclass(kw_only=True)
 class PolicyLearnerCfg:

mpail2/configs/defs.py

@@ -14,7 +14,7 @@
 
 #### SHARED CONSTANTS ####
 OPT = "adam"
-LR = 3e-4
+LR = 2e-4
 HORIZON = 7
 OPT_ITERS = 5
 GAMMA = 0.99
@@ -54,6 +54,7 @@ class RewardConfig(cfgs.RewardCfg):
         "use_layer_norm": False,
         "disable_output_bias": True,
     })
+    reward_clip: float = None
 
 @dataclass(kw_only=True)
 class EnsembleValueConfig(cfgs.EnsembleValueCfg):
@@ -83,7 +84,10 @@ class CNNCoderConfig(cfgs.CNNCoderCfg):
     model_kwargs: dict = field(default_factory=lambda: {
         **CNN_MODEL_KWARGS,
         ## OVERRIDES ##
-        "override_last_layer_activation": True,
+        # Camera embedding ends in LayerNorm (normalized, zero-centered), not a
+        # trailing one-sided SiLU which distorts the geometry LayerNorm just
+        # established (ported from mpail-research fix-main #18).
+        "override_last_layer_norm": True,
     })
 
 @dataclass(kw_only=True)
@@ -95,9 +99,9 @@ class ProprioCoderConfig(cfgs.MLPCoderCfg):
             **MODEL_KWARGS,
             ## OVERRIDES ##
             "hidden_dims": [256],
-            "use_layer_norm": False,
+            # Hidden LN on (inherited from MODEL_KWARGS), no trailing SiLU on the
+            # latent slice — same rationale as CNNCoderConfig above.
             "override_last_layer_norm": True,
-            "override_last_layer_activation": True,
         })
 
     model_kwargs:dict = field(default_factory=lambda: {
@@ -175,10 +179,10 @@ class PlannerConfig(cfgs.PlannerCfg):
     temperature: float = 2.0
     opt_iters: int = OPT_ITERS
 
-    reward_cfg: RewardConfig = RewardConfig()
-    value_cfg: EnsembleValueConfig = EnsembleValueConfig()
-    sampling_cfg: PolicySamplingConfig = PolicySamplingConfig()
-    dynamics_cfg: DynamicsConfig = DynamicsConfig()
+    reward_cfg: RewardConfig = field(default_factory=RewardConfig)
+    value_cfg: EnsembleValueConfig = field(default_factory=EnsembleValueConfig)
+    sampling_cfg: PolicySamplingConfig = field(default_factory=PolicySamplingConfig)
+    dynamics_cfg: DynamicsConfig = field(default_factory=DynamicsConfig)
 
     seed: int = 42
     u_per_command: int = 1
@@ -210,7 +214,7 @@ class RewardLearnerConfig(cfgs.RewardLearnerCfg):
 
     opt_params: dict = field(default_factory=lambda: OPT_PARAMS)
 
-    gp_coeff: float = 0.1
+    gp_coeff: float = 5.0
 
     gp_target_gradient: float = 1.0
 
@@ -223,7 +227,7 @@ class PolicyLearnerConfig(cfgs.PolicyLearnerCfg):
 
     max_grad_norm: float = 1.0
 
-    target_entropy: float = -3.0  # -ACTION_DIM
+    target_entropy: float = -2.0  # encourage exploration; -ACTION_DIM=-5 is too conservative
 
     alpha_lr: float = LR
 
@@ -236,12 +240,19 @@ class DynamicsLearnerConfig(cfgs.DynamicsLearnerCfg):
 
     opt_params: dict = field(default_factory=lambda: OPT_PARAMS)
 
-    enc_lr_scale: float = 0.1
+    enc_lr_scale: float = 0.08
 
     rho: float = 0.95
 
     recon_coeff: float = 1.0 # Exact coefficient doesn't matter if not using recon loss
 
+    # Conservative starting point (ported from mpail-research fix-main #18, which used
+    # 0.1 on a sim pick-place task) — start smaller so it doesn't swamp the JEP loss
+    # before it's been tuned here.
+    sigreg_coeff: float = 0.02
+    sigreg_knots: int = 17
+    sigreg_num_proj: int = 1024
+
 @dataclass(kw_only=True)
 class LearnerConfig(cfgs.MPAIL2LearnerCfg):
 
@@ -259,13 +270,13 @@ class LearnerConfig(cfgs.MPAIL2LearnerCfg):
     loss_horizon: int = HORIZON  # Match num_timesteps in sampling_cfg
 
     # Dynamics - ENABLED for training
-    dynamics_learner_cfg: cfgs.DynamicsLearnerCfg = DynamicsLearnerConfig()
+    dynamics_learner_cfg: cfgs.DynamicsLearnerCfg = field(default_factory=DynamicsLearnerConfig)
 
     # Reward (adversarial training)
-    reward_learner_cfg: cfgs.RewardLearnerCfg = RewardLearnerConfig()
+    reward_learner_cfg: cfgs.RewardLearnerCfg = field(default_factory=RewardLearnerConfig)
 
     # Value
-    value_learner_cfg: cfgs.ValueLearnerCfg = ValueLearnerConfig()
+    value_learner_cfg: cfgs.ValueLearnerCfg = field(default_factory=ValueLearnerConfig)
 
     # Policy Learner
-    policy_learner_cfg: cfgs.PolicyLearnerCfg = PolicyLearnerConfig()
+    policy_learner_cfg: cfgs.PolicyLearnerCfg = field(default_factory=PolicyLearnerConfig)

mpail2/dynamics.py

@@ -9,6 +9,46 @@
     from .configs.cfgs import DynamicsCfg
 
 
+class SIGReg(torch.nn.Module):
+    """Sketch Isotropic Gaussian Regularizer.
+
+    Projects a batch of latents onto random 1D directions and penalizes deviation
+    of each projection's empirical characteristic function from a standard
+    Gaussian's (via a differentiable normality test, quadrature-integrated over
+    `knots`). Minimizing this pushes variance to spread across many random
+    directions instead of concentrating in a few (representation/dimensional
+    collapse) — ported from mpail-research fix-main #18.
+    """
+
+    def __init__(self, knots: int = 17, num_proj: int = 1024):
+        super().__init__()
+        self.num_proj = num_proj
+
+        t = torch.linspace(0.0, 3.0, knots, dtype=torch.float32)
+        dt = 3.0 / (knots - 1)
+        weights = torch.full((knots,), 2.0 * dt, dtype=torch.float32)
+        weights[[0, -1]] = dt
+        window = torch.exp(-t.square() / 2.0)
+
+        self.register_buffer("t", t)
+        self.register_buffer("phi", window)
+        self.register_buffer("weights", weights * window)
+
+    def forward(self, proj: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            proj: Tensor of shape (T, B, D).
+        """
+        A = torch.randn(proj.size(-1), self.num_proj, device=proj.device, dtype=proj.dtype)
+        A = A.div_(A.norm(p=2, dim=0))
+
+        x_t = (proj @ A).unsqueeze(-1) * self.t.to(dtype=proj.dtype)
+        err = (x_t.cos().mean(-3) - self.phi.to(dtype=proj.dtype)).square()
+        err = err + x_t.sin().mean(-3).square()
+        statistic = (err @ self.weights.to(dtype=proj.dtype)) * proj.size(-2)
+        return statistic.mean()
+
+
 class Dynamics(torch.nn.Module):
     """
     Latent dynamics model. Propagates latent states through a feedforward network.

mpail2/envs/real/__init__.py

@@ -64,3 +64,38 @@
     )
 except ImportError:
     pass
+
+try:
+    from .so101 import (
+        ACTION_DIM as SO101_ACTION_DIM,
+        SO101RealWrapper,
+        SO101RealEnvArgs,
+        STATE_DIM as SO101_STATE_DIM,
+        EE_PROPRIO_DIM as SO101_EE_PROPRIO_DIM,
+        HOME_POSITION_DEG as SO101_HOME_POSITION_DEG,
+        EE_LOWER_M as SO101_EE_LOWER_M,
+        EE_UPPER_M as SO101_EE_UPPER_M,
+        JOINT_LOWER_DEG as SO101_JOINT_LOWER_DEG,
+        JOINT_UPPER_DEG as SO101_JOINT_UPPER_DEG,
+        MAX_EPISODE_STEPS as SO101_MAX_EPISODE_STEPS,
+        make_so101_env,
+    )
+
+    __all__.extend(
+        [
+            "SO101RealWrapper",
+            "SO101RealEnvArgs",
+            "SO101_STATE_DIM",
+            "SO101_ACTION_DIM",
+            "SO101_EE_PROPRIO_DIM",
+            "SO101_HOME_POSITION_DEG",
+            "SO101_EE_LOWER_M",
+            "SO101_EE_UPPER_M",
+            "SO101_JOINT_LOWER_DEG",
+            "SO101_JOINT_UPPER_DEG",
+            "SO101_MAX_EPISODE_STEPS",
+            "make_so101_env",
+        ]
+    )
+except ImportError:
+    pass