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.

Vision

ViT - Vision Transformer.

PatchEmbedding - Converts images to patch sequences via Conv2D.

Wrappers

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