Skip to content

moe

Mixture-of-Experts MLP model for rsl-rl.

Mirrors lav2/runner/skrl/cfg/LAV2_base_moe.py but adapted to rsl-rl's MLPModel interface.

Uses loss-free load balancing: expert bias terms are adjusted after each PPO update based on expert utilisation, without an auxiliary loss term.

References
  • https://kexue.fm/archives/10699 (Math details)
  • https://kexue.fm/archives/10757 (Auxiliary-Loss-Free Load Balancing)
  • https://github.com/ambisinister/lossfreebalance (Loss-Free impl)

Classes:

Name Description
MoEMLPModel

MLP model whose trunk uses a mixture-of-experts layer.

MoEMLPModel 

MoEMLPModel(obs, obs_groups, obs_set, output_dim, hidden_dims=(64, 64), activation='elu', obs_normalization=False, distribution_cfg=None, num_experts: int = 4, k: int = 2)

Bases: MLPModel

MLP model whose trunk uses a mixture-of-experts layer.

The architecture mirrors the skrl counterpart: Linear → Act → MoELayer → Act → Linear(head).

Parameters:

Name Type Description Default

num_experts

int

Number of experts in the MoE layer.

 4

k

int

Number of experts routed to per token (top-k).

 2

Build the MoE-enhanced trunk and distribution head.