moe

Shared model with Mixture-of-Experts trunk.

Classes:

Name	Description
MoELayer	Mixture-of-experts layer with sigmoid top-k routing.
SharedMoE	Shared actor-critic with MoE trunk.

MoELayer

MoELayer(input_size: int, output_size: int, num_experts: int = 4, k: int = 2, bias_update_speed: float = 0.001)

Bases: Module

Mixture-of-experts layer with sigmoid top-k routing.

Initialize experts, gate network, and load-balancing bias.

Methods:

Name	Description
forward	Compute routed expert outputs.

forward

forward(x: Tensor) -> torch.Tensor

Compute routed expert outputs.

Returns:

Type	Description
Tensor	torch.Tensor: Weighted combination of top-k expert outputs.

SharedMoE

SharedMoE(observation_space, state_space, action_space, device, clip_actions=False, clip_log_std=True, min_log_std=-20, max_log_std=2, reduction='sum')

Bases: GaussianMixin, DeterministicMixin, Model

Shared actor-critic with MoE trunk.

Initialize the shared MoE actor-critic model.

Methods:

Name	Description
act	Compute actions or values depending on role.
compute	Compute policy/value outputs depending on role.

act

act(inputs: Any, role: str) -> Any

Compute actions or values depending on role.

Returns:

Type	Description
Any	tuple | None: Action outputs for the given role, or None if role is unknown.

compute

compute(inputs: Any, role: str) -> Any

Compute policy/value outputs depending on role.

Returns:

Type	Description
Any	tuple | None: (output_tensor, extras_dict) for policy/value roles, or None.