Support real SO-101 arm training and fix some bugs

[Ghosty2003]

Add SO-101 real-arm support + port anti-collapse/bugfix patches from mpail-lerobot fork

Summary

Two commits, two different purposes:

1. b55aff0 — purely additive: a new mpail2/envs/real/so101/ environment stack for
   training on the real SO-101 arm. Touches no existing library code.
2. 39402f8 — ports several fixes/features from the mpail-lerobot fork into the shared
   core library (dynamics.py, learner.py, planner.py, reward.py, configs,
   obs_normalizer.py). These affect all envs (sim + real), not just SO-101.

---

1. SO-101 real-arm environment (b55aff0)

New package mpail2/envs/real/so101/, mirroring the existing Franka/Kinova real-env layout:

• env_factory.py, so101_env.py, wrappers.py, ik_utils.py, robot_limits.py —
  the Gym-style env itself: joint-space IK, per-joint limits/tolerances (gripper handled
  separately since it's a slower/weaker actuator than the 5 arm joints).
• network/server.py + transport/*_pb2*.py — a gRPC server/client pair so the env can
  step a physically remote robot process (arm + cameras run in a separate lerobot conda env
  that doesn't have mpail2 installed; mpail2 talks to it over gRPC instead of importing it
  directly). transport/*_pb2*.py are generated protobuf code from
  transport/so101_robot.proto — not hand-written, please skip line-by-line review of those.
• lerobot_patch/async_inference/ — a local copy of lerobot's async inference stack
  (policy_server.py, robot_client.py, robot_server.py, rl_client.py,
  teleop_with_planner.py, etc.) adapted to serve MPAIL2 policies instead of LeRobot's own.
• training/ — SO-101-specific entry points: train_so101_local.py (drives the robot
  in-process via the standard MPAIL2Runner.learn() rollout loop, blocking on real
  env.step()/env.reset() so no observation is ever dropped — as opposed to
  demo_recording_server.py, which reacts to inbound RPCs from a separately-running
  lerobot robot_client), plus convert.py / convert_lerobot.py (demo format conversion),
  replay_demo.py, check_encoder_collapse.py.
• soa.urdf — SO-101 URDF for IK.
• README.md — full setup instructions (separate lerobot conda env + clone).
• pyproject.toml — new so101 extra (opencv-python, pyrealsense2, grpcio,
  grpcio-tools, ikpy).
• docs/INSTALL.md — new row/bullet pointing at the so101 extra and the README above.

No existing files are modified beyond the two-line pyproject.toml/INSTALL.md additions.

---

2. Core library changes ported from mpail-lerobot fork (39402f8)

SIGReg anti-collapse regularizer

• mpail2/dynamics.py — new SIGReg(nn.Module). Projects a batch of latents onto
  num_proj random 1D directions and penalizes deviation of each projection's empirical
  characteristic function from a standard Gaussian's (quadrature-integrated over knots
  points). Pushes latent variance to spread across directions instead of collapsing onto a
  few — i.e. an encoder-collapse regularizer.
• mpail2/learner.py
  • MPAIL2Learner.__init__ — builds self._sigreg = SIGReg(knots=..., num_proj=...) when
    dynamics_learner_cfg.sigreg_coeff is set and > 0; None otherwise (opt-in, no-op by
    default since sigreg_coeff defaults to None in the base cfgs.py).
  • MPAIL2Learner.update_dynamics — adds sigreg_coeff * sigreg_loss to the JEP loss and
    logs Dyn/sigreg_loss.
• New config fields (configs/cfgs.py): sigreg_coeff, sigreg_knots, sigreg_num_proj
  on DynamicsLearnerCfg. Tuned defaults (configs/defs.py,
  DynamicsLearnerConfig): sigreg_coeff=0.02, sigreg_knots=17, sigreg_num_proj=1024
  — a conservative starting point (the fork used 0.1 on a sim pick-place task; started
  lower here so it doesn't swamp the JEP loss before being tuned on this codebase).

Encoder architecture: LayerNorm-ending instead of trailing SiLU

• configs/defs.py, CNNCoderConfig / MultiCoderConfig — model_kwargs switched from
  "override_last_layer_activation": True to "override_last_layer_norm": True. Rationale
  in-line: a trailing one-sided SiLU after the encoder's LayerNorm distorts the
  normalized/zero-centered geometry the LayerNorm just established.

Reward output clamping

• mpail2/reward.py, Reward.forward — clamps output to [-reward_clip, reward_clip]
  when cfg.reward_clip is not None (default None → disabled). Motivated by observed
  drift in the WGAN-style critic's unbounded output scale (mean_demo_reward climbing
  ~6→16, Value/mean_q_value swinging -187↔+164 across training iterations).
• New config field reward_clip on RewardCfg / RewardConfig (defaults to None,
  i.e. no behavior change unless explicitly set).

Runtime reward scaling hook

• mpail2/learner.py, update_value and _n_step_return_lambda — both now scale the
  reward used in the value target by getattr(self, '_reward_scale', 1.0). This isn't a
  config field; it's a plain instance attribute set at runtime — currently only from
  training/train_so101_local.py (runner.learner._reward_scale = float(args.reward_scale)).
  No-op (1.0) for every existing sim/Isaac training path.

Planner warm-start bug fix

• mpail2/planner.py, Planner.optimize — was calling self.sampling.reset_iter_state()
  with no arguments every single step, which resets _iter_mean to zero each time —
  discarding the warm-started _opt_controls carried over from the previous step and turning
  ~95% of each decision's candidates (the noise-sampled ones, not the small
  policy_proportion fraction) into a fresh zero-mean/max-std blind search with no memory of
  the previous decision's converged plan. Now passes
  reset_iter_state(prev_controls=self._opt_controls), matching reset_iter_state's own
  docstring contract.
• Also adds Planner.act_policy_only() — returns the policy network's own deterministic
  (tanh-squashed mean) action directly, bypassing CEM/MPPI entirely (no noise rollouts, no
  elite re-weighting). Useful for evaluating what the policy net alone has learned,
  independent of the planner's search.

Python 3.11+ dataclass mutable-default fix

• configs/defs.py — PlannerConfig, LearnerConfig: fields like
  reward_cfg: RewardConfig = RewardConfig() (a mutable default shared across all instances)
  changed to field(default_factory=RewardConfig). This pattern happened to work on whatever
  Python version was previously used but raises ValueError: mutable default <class '...'> for field ... is not allowed under Python 3.11+'s stricter dataclass checks — blocking anyone
  on a newer Python from importing the configs at all.

obs_normalizer.py fixes

• Dotted-key buffer names: FixedObsNormalizer registered buffers as f"{key}_mean" /
  f"{key}_std", which breaks (register_buffer rejects dots in names) for any obs key
  containing a dot (e.g. SO-101's cam.wrist-style keys). Added
  _buf_name() (key.replace(".", "_")) and routed all three read/write sites
  (_compute_statistics, forward, inverse) through it.
• Double-divide-by-255 bug: CamOnlyObsNormalizer.forward/inverse unconditionally did
  value / 255.0 - 0.5 / (value + 0.5) * 255.0. Some callers (so101_env.py, convert.py)
  already divide by 255 before calling this, so the unconditional divide silently
  double-divided and crushed the whole batch into a near-zero-variance sliver near -0.5
  regardless of image content. Now guards with the same value.max() > 1.0 check
  convert.py already uses, and inverse no longer re-multiplies by 255 (this codebase's
  actual convention is that callers always hand forward() already-[0,1] images, so
  scaling back to [0,255] was the wrong inverse).

Tuned hyperparameter defaults (configs/defs.py)

Param	Old	New
LR (shared)	3e-4	2e-4
RewardLearnerConfig.gp_coeff	0.1	5.0
PolicyLearnerConfig.target_entropy	-3.0 (-ACTION_DIM)	-2.0 (more exploration; -ACTION_DIM=-5 judged too conservative)
DynamicsLearnerConfig.enc_lr_scale	0.1	0.08

All four match the tuned values used in the mpail-lerobot fork.

---

Notes for reviewers

• 39402f8's changes are not SO-101-specific — they touch shared config defaults and
  loss functions used by every existing sim/Isaac training run. Worth a sanity training run
  on an existing sim task (e.g. Ant-v5 or the Franka pick-place) to confirm the LR/gp_coeff/
  target_entropy/enc_lr_scale changes and the encoder LayerNorm-ending switch don't regress
  current results, since none of the new features (sigreg_coeff, reward_clip) are on by
  default but these defaults changes are.
• transport/*_pb2*.py are generated protobuf code — regenerate from
  transport/so101_robot.proto rather than hand-editing if changes are needed there.
• _reward_scale is a runtime-set instance attribute rather than a config field (see above);
  flagging in case a config field was actually intended instead.