memorax.networks

Neural network components for building RL agents.

Core

Network - Main network class composing feature extractors, torsos, and heads.

FeatureExtractor - Extracts features from observations, actions, rewards, and done flags.

Identity - Identity module that passes input through unchanged.

Architectures

MLP - Multi-layer perceptron.

CNN - Convolutional neural network.

ViT - Vision Transformer.

Wrappers

SequenceModelWrapper - Wraps non-recurrent models for use as sequence models.