protomotions.agents.common.mlp module

Multi-Layer Perceptron (MLP) network implementations.

This module provides MLP architectures used throughout the codebase. All MLPs support optional observation normalization and operate on TensorDict inputs.

These are the building blocks for actor and critic networks in RL agents.

Key Classes:

• MLP: Feedforward network with optional observation normalization

• MultiHeadedMLP: MLP with multiple parallel input heads and a trunk

Functions:

• build_mlp: Factory function to construct MLP from configuration

protomotions.agents.common.mlp.build_mlp(config)

Build a multi-layer perceptron from configuration using LazyLinear.

Uses LazyLinear for automatic input size inference. The first forward pass will materialize the layers with the correct input dimensions.

Parameters:

config (MLPWithConcatConfig) – MLP configuration specifying layers, activations, and output dimensions.

Returns:

Sequential neural network module with lazy initialization.

class protomotions.agents.common.mlp.MLPWithConcat(*args, **kwargs)

Bases: TensorDictModuleBase

Multi-layer perceptron network with optional observation normalization.

Feedforward network that processes observations through multiple fully-connected layers with configurable activations. Optionally normalizes inputs using running mean/std statistics.

REQUIRES explicit obs_key and out_key to prevent key collisions. Always operates on TensorDict for clean, traceable data flow.

Parameters:

config (MLPWithConcatConfig) – Configuration specifying input/output dimensions, hidden layers, and normalization settings. Both obs_key and out_key must be explicitly set.

norm

NormObsBase module for optional normalization (plain nn.Module).

mlp

Sequential network of linear layers and activations.

in_keys

List of input keys (always non-empty).

out_keys

List of output keys (always non-empty).

__init__(config)

config: MLPWithConcatConfig

forward(tensordict)

Forward pass with optional normalization.

Parameters:

• tensordict (MockTensorDict) – TensorDict containing observations.

• return_norm_obs – If True, stores normalized obs in tensordict.

Returns:

TensorDict with processed outputs.

Return type:

MockTensorDict