Add a strictly equivariant model for symmetry-aware RL

[vanillaturtlechips]

Closes #237.

Adds EquivariantMLPModel, a model whose network is equivariant by construction under a reflection of the robot, so that π(Ms) = Mπ(s) and V(Ms) = V(s) hold exactly (up to floating-point precision) throughout training, instead of being encouraged through data augmentation or a mirror loss.

As discussed in the issue, this PR targets the extras branch.

Changes

The change is purely additive. PPO, the existing models, and the symmetry extension are untouched.

+  rsl_rl/modules/equivariant.py             SignedPermutation, EquivariantLinear, EquivariantMLP,
                                             SymmetricEmpiricalNormalization, EquivariantGaussianDistribution
+  rsl_rl/models/mlp_model_equivariant.py    EquivariantMLPModel(MLPModel)
+  tests/modules/test_equivariant.py
+  tests/models/test_mlp_model_equivariant.py
M  rsl_rl/modules/__init__.py, rsl_rl/models/__init__.py   exports only
M  docs/guide/overview.rst, docs/guide/configuration.rst, docs/index.rst, docs/api/models.rst, docs/api/modules.rst
M  CONTRIBUTORS.md

The __init__.py exports only enable the short class_name. Without them, the model can also be used through a qualified name ("rsl_rl.models.mlp_model_equivariant:EquivariantMLPModel") via resolve_callable.

Design

• EquivariantLinear projects its weight onto the equivariant subspace on every forward pass, W_eq = ½(W + M_out W M_in) and b_eq = ½(b + M_out b). Gradients flow to the unconstrained parameter. The raw weight is rescaled at init so the projected weight starts at the nn.Linear scale.
• Hidden layers carry the regular representation of the order-two group (a pure permutation), so pointwise activations commute with the symmetry.
• EquivariantMLPModel subclasses MLPModel and replaces only self.mlp, the same way CNNModel and RNNModel extend it. Omitting equivariance_cfg["output"] gives the trivial output representation, i.e. an invariant critic.
• EquivariantGaussianDistribution averages the std over mirrored action pairs when it's used, so the sampled action distribution is equivariant and not only the mean. The stored parameter and optimizer state are left unchanged.
• SymmetricEmpiricalNormalization is used when obs_normalization=True. Running statistics of real data are generally not symmetric, which would break equivariance, so it learns each batch together with its mirror image and projects the statistics onto the symmetric subspace after every update. This makes normalization commute with the symmetry exactly.
• fold() turns a trained layer or MLP into plain nn.Linear/nn.Sequential modules computing the same function. as_jit()/as_onnx() use it, so the exported policy is a plain MLP without the per-forward projection.

Usage

The configuration is documented in the configuration guide.

actor = {
    "class_name": "EquivariantMLPModel",
    "hidden_dims": [512, 256, 128],
    "equivariance_cfg": {
        "obs":    {"perm": OBS_PERM, "sign": OBS_SIGN},
        "output": {"perm": ACT_PERM, "sign": ACT_SIGN},
    },
    "distribution_cfg": {"class_name": "EquivariantGaussianDistribution", "perm": ACT_PERM},
}
critic = {
    "class_name": "EquivariantMLPModel",
    "hidden_dims": [512, 256, 128],
    "equivariance_cfg": {"obs": {"perm": OBS_PERM, "sign": OBS_SIGN}},
}

Testing

• 21 new tests (27 cases with parametrization) check the structural property directly: representation validation, the per-layer constraint, actor equivariance, critic invariance, equivariance after optimizer steps, symmetric action noise, fold correctness, init scale, symmetric normalization under biased data, folded TorchScript/ONNX export, and config errors.
• The full suite passes (214 tests), and pre-commit is clean.
• Experimental results on a small humanoid in Isaac Lab (3 matched seeds) are in Add a strictly equivariant model for symmetry-aware RL #237.

Known limitations / possible follow-ups

• Only Gaussian distributions with a state-independent std are supported. Heteroscedastic Gaussian and Beta raise NotImplementedError.
• Only the order-two (reflection) group is supported. Z2 × Z2 (e.g. left-right × front-back for quadrupeds) would generalize the projection from an order-two average to a group average.

🤖 Generated with Claude Code

[ClemensSchwarke]

Hi @vanillaturtlechips,
I created the extras branch now. I quickly went over the PR and it looks quite good already. I think it would be nice to rename rsl_rl/models/equivariant_mlp_model.py rsl_rl/models/mlp_model_equivariant.py, just to have a nicer order in the file list. What do you think? Also, since you already have many classes in equivariant.py, it might be cleaner to move the distribution there as well. Lastly, if you could use the same augmentation function used by the rest that would make it much easier for users I guess, but I am not sure if that can be done in a clean way. Please check out the extras branch, it should give you some info on how and where to document your feature.

Thanks!

[vanillaturtlechips]

@ClemensSchwarke

Thanks for creating the extras branch and for the quick review! I rebased onto extras and retargeted the PR.

• Renamed the model to rsl_rl/models/mlp_model_equivariant.py (and the test file accordingly).
• Moved EquivariantGaussianDistribution into rsl_rl/modules/equivariant.py.
• Documented the model and the distribution in the overview guide, added the feature to the list on the main page, and added the API reference entries.

Regarding the augmentation function: since it is a signed permutation, the perm/sign can be derived exactly by passing the basis vectors through it. I added symmetry_cfg_from_augmentation() for that, and checked it against a replica of Isaac Lab's ANYmal compute_symmetric_states (including the height scan and all three augmentations). It also rejects functions that are not a signed involution, and groups the function leaves unchanged, since Isaac Lab only mirrors the policy group, so a separate critic group would otherwise silently yield an unconstrained model.

The one part I couldn't make fully clean is that these functions need env (e.g. env.unwrapped.observation_manager), and models don't receive it at construction. So for now the helper is called once from the training script:

obs = env.get_observations()
agent_cfg["actor"]["symmetry_cfg"] = symmetry_cfg_from_augmentation(
    compute_symmetric_states, env, obs, agent_cfg["obs_groups"]["actor"], env.num_actions
)

If you're OK with a small core change, a one-line env injection in construct_algorithm (similar to resolve_symmetry_config) would let users just put the augmentation function in the model config. Happy to add that, or to keep it as is. Let me know what you prefer!

[ClemensSchwarke]

Actually, I am not sure anymore if supporting the augmentation function makes sense. I saw that you just convert it to your symmetry config structure, that might make things more confusing than easier. Since you don't really interact with the symmetry extension and already call your classes equivariant, I would rename your configs to equivariance_cfg or something similar. Currently, there are two completely different symmetry_cfg dicts in the repo.

[ClemensSchwarke]

Dropping the connection to the symmetry extension should make things more clear, it would be great if you could also document the configuration of your classes in the configuration documentation.

[ClemensSchwarke]

Comments from my agent:
3. Unused code:

• SignedPermutation.to() is never called.
• The group_order argument of regular() accepts only 2.
• fold() on the layer and on the MLP is only used in tests. The export path deepcopies the unfolded MLP, so fold() never removes the per-forward index operations it was written to remove. Either use it in as_jit/as_onnx or remove it.

Questions for the author
4. CONTRIBUTORS.md has "Lee Myong Ii". Should that be "Il"?

[vanillaturtlechips]

Thanks for the review! I agree, keeping it separate from the symmetry extension is clearer.

• Dropped symmetry_cfg_from_augmentation() and renamed the model's symmetry_cfg to equivariance_cfg, so there is only one symmetry_cfg in the repo again.
• Documented EquivariantMLPModel and EquivariantGaussianDistribution in the configuration guide, including a YAML example.
• Unused code:
  • Removed SignedPermutation.to().
  • Removed the group_order argument of regular().
  • as_jit()/as_onnx() now use fold(), so the exported policy is a plain MLP without the per-forward index operations. The export tests check that no EquivariantLinear is left and that the outputs match.
• Fixed the Python 3.9 failure (missing from __future__ import annotations in the test files).
• Yes, it should be "Il", thanks for catching that!