protomotions.agents.common.config module#

class protomotions.agents.common.config.NormObsBaseConfig( normalize_obs=False, norm_clamp_value=5.0, norm_ema_decay=0.999, )[source]#

Bases: object

Base configuration for modules that support optional observation normalization.

With LazyLinear, only num_out is needed - input sizes are inferred automatically. This is purely about normalization settings and output dimensions. Individual TensorDictModules add their own obs_key/out_key fields as needed.

Attributes:: normalize_obs: Whether to normalize observations using running statistics. norm_clamp_value: Clamp normalized values to [-value, value] to prevent extreme outliers. norm_ema_decay: EMA decay for obs normalization (None = Welford). Default 0.999 tracks a ~1000-sample moving window.

normalize_obs: bool = False#

norm_clamp_value: float = 5.0#

norm_ema_decay: float | None = 0.999#

__init__( normalize_obs=False, norm_clamp_value=5.0, norm_ema_decay=0.999, )#

class protomotions.agents.common.config.ModuleOperationConfig[source]#

Bases: object

Configuration for module operations.

__init__()#

class protomotions.agents.common.config.ModuleOperationForwardConfig[source]#

Bases: ModuleOperationConfig

Configuration for module operation forward.

__init__()#

class protomotions.agents.common.config.ModuleOperationPermuteConfig(new_order)[source]#

Bases: ModuleOperationConfig

Configuration for module operation permute.

Attributes:: new_order: New dimension order, e.g. [0, 2, 1] swaps dims 1 and 2.

new_order: List[int]#

__init__(new_order)#

class protomotions.agents.common.config.ModuleOperationReshapeConfig(new_shape)[source]#

Bases: ModuleOperationConfig

Configuration for module operation reshape.

Attributes:: new_shape: New shape. Use ‘batch_size’ for dynamic batch dim.

new_shape: List[int | str]#

__init__(new_shape)#

class protomotions.agents.common.config.ModuleOperationSqueezeConfig(squeeze_dim)[source]#

Bases: ModuleOperationConfig

Configuration for module operation squeeze.

Attributes:: squeeze_dim: Dimension to squeeze (remove if size is 1).

squeeze_dim: int#

__init__(squeeze_dim)#

class protomotions.agents.common.config.ModuleOperationUnsqueezeConfig(unsqueeze_dim)[source]#

Bases: ModuleOperationConfig

Configuration for module operation unsqueeze.

Attributes:: unsqueeze_dim: Position where to insert new dimension of size 1.

unsqueeze_dim: int#

__init__(unsqueeze_dim)#

class protomotions.agents.common.config.ModuleOperationExpandConfig(expand_shape)[source]#

Bases: ModuleOperationConfig

Configuration for module operation expand.

Attributes:: expand_shape: Target shape to expand to. Use -1 to keep original size.

expand_shape: List[int]#

__init__(expand_shape)#

class protomotions.agents.common.config.ModuleOperationSphereProjectionConfig[source]#

Bases: ModuleOperationConfig

Configuration for sphere projection operation (L2 normalization to unit sphere).

__init__()#

class protomotions.agents.common.config.ObsProcessorConfig( normalize_obs=False, norm_clamp_value=5.0, norm_ema_decay=0.999, _target_='protomotions.agents.common.common.ObsProcessor', in_keys=<factory>, out_keys=<factory>, module_operations=<factory>, )[source]#

Bases: NormObsBaseConfig

General observation processor - applies operations and normalization.

Supports all module_operations. ForwardConfig applies normalization but skips the forward model (no MLP). Useful for reshaping, normalizing, and other tensor manipulations.

Attributes:: normalize_obs: Whether to normalize observations using running statistics. norm_clamp_value: Clamp normalized values to [-value, value] to prevent extreme outliers. norm_ema_decay: EMA decay for obs normalization (None = Welford). Default 0.999 tracks a ~1000-sample moving window. in_keys: Input tensor keys to read from TensorDict. out_keys: Output tensor keys to write to TensorDict. module_operations: Sequence of operations to apply (reshape, permute, etc).

in_keys: List[str]#

out_keys: List[str]#

module_operations: List[ModuleOperationConfig]#

__init__( normalize_obs=False, norm_clamp_value=5.0, norm_ema_decay=0.999, _target_='protomotions.agents.common.common.ObsProcessor', in_keys=<factory>, out_keys=<factory>, module_operations=<factory>, )#

class protomotions.agents.common.config.MLPLayerConfig(units=512, activation='relu', use_layer_norm=False)[source]#

Bases: object

Configuration for a single MLP layer.

Attributes:: units: Number of neurons in this layer. activation: Activation function for this layer. use_layer_norm: Whether to apply layer normalization after activation.

units: int = 512#

activation: str = 'relu'#

use_layer_norm: bool = False#

__init__( units=512, activation='relu', use_layer_norm=False, )#

class protomotions.agents.common.config.MLPWithConcatConfig( normalize_obs=False, norm_clamp_value=5.0, norm_ema_decay=0.999, num_out=None, layers=<factory>, _target_='protomotions.agents.common.mlp.MLPWithConcat', in_keys=<factory>, out_keys=<factory>, output_activation=None, module_operations=<factory>, )[source]#

Bases: NormObsBaseConfig

Configuration for Multi-Layer Perceptron with optional normalization.

Unified MLP configuration that supports optional input normalization. Set normalize_obs=False if you don’t want normalization (default is False).

Attributes:: normalize_obs: Whether to normalize observations using running statistics. norm_clamp_value: Clamp normalized values to [-value, value] to prevent extreme outliers. norm_ema_decay: EMA decay for obs normalization (None = Welford). Default 0.999 tracks a ~1000-sample moving window. num_out: Output dimension of the MLP. Required. layers: List of layer configurations defining the MLP architecture. in_keys: Input tensor keys to read and concatenate from TensorDict. out_keys: Output tensor keys to write to TensorDict. output_activation: Activation function for the output layer (None for linear output). module_operations: Sequence of operations including forward pass and reshapes.

num_out: int = None#

layers: List[MLPLayerConfig]#

in_keys: List[str]#

out_keys: List[str]#

output_activation: str | None = None#

module_operations: List[ModuleOperationConfig]#

__init__( normalize_obs=False, norm_clamp_value=5.0, norm_ema_decay=0.999, num_out=None, layers=<factory>, _target_='protomotions.agents.common.mlp.MLPWithConcat', in_keys=<factory>, out_keys=<factory>, output_activation=None, module_operations=<factory>, )#

class protomotions.agents.common.config.ModuleContainerConfig( models=<factory>, _target_='protomotions.agents.common.common.ModuleContainer', in_keys=<factory>, out_keys=<factory>, )[source]#

Bases: object

Configuration for a container of modules that are executed sequentially.

Modules are processed in order, with each module’s outputs available to subsequent modules. Input keys are passed through, and all specified output keys must be produced by internal modules.

Attributes:: models: List of module configurations to execute sequentially. in_keys: Input tensor keys required by this container. out_keys: Output tensor keys produced by this container.

models: List[Any]#

in_keys: List[str]#

out_keys: List[str]#

__init__( models=<factory>, _target_='protomotions.agents.common.common.ModuleContainer', in_keys=<factory>, out_keys=<factory>, )#

class protomotions.agents.common.config.PretrainedModelConfig( checkpoint_path='', module_path='', module_config_path=None, module_config=None, config_path='agent.model', state_dict_key='model', strict=True, freeze=True, )[source]#

Bases: object

Configuration for loading a model from a checkpoint and returning a submodule.

Attributes:: checkpoint_path: Checkpoint containing state dict. module_path: Dotted module path to return from the loaded model, e.g. ‘actor.mu’ or ‘actor.mu.decoder’. Empty string returns the whole loaded model. module_config_path: Optional dotted config path for the returned module inside the pretrained checkpoint’s resolved configs. module_config: Optional embedded config for the returned module. Inference bundles use this to instantiate frozen modules whose weights are saved in the owning checkpoint. config_path: Dotted config path inside resolved_configs.pt. state_dict_key: Top-level checkpoint key containing model weights. strict: Use strict state-dict loading for the model. freeze: Freeze returned module parameters after loading.

checkpoint_path: str = ''#

module_path: str = ''#

module_config_path: str | None = None#

module_config: Any | None = None#

config_path: str = 'agent.model'#

state_dict_key: str = 'model'#

strict: bool = True#

freeze: bool = True#

__init__( checkpoint_path='', module_path='', module_config_path=None, module_config=None, config_path='agent.model', state_dict_key='model', strict=True, freeze=True, )#

class protomotions.agents.common.config.TransformerConfig( _target_='protomotions.agents.common.transformer.Transformer', in_keys=<factory>, out_keys=<factory>, input_and_mask_mapping=None, transformer_token_size=512, latent_dim=512, num_heads=4, ff_size=1024, num_layers=4, dropout=0.0, activation='relu', output_activation=None, )[source]#

Bases: object

Configuration for Transformer encoder.

Multi-head self-attention transformer that processes tokenized inputs. Supports optional masking for variable-length sequences.

Attributes:: in_keys: Input tensor keys (tokens and optional masks). out_keys: Output tensor key for transformer output (exactly one). input_and_mask_mapping: Maps input token keys to their mask keys for attention masking. transformer_token_size: Expected input token dimension size. latent_dim: Internal/output dimension of transformer. num_heads: Number of attention heads. Must divide latent_dim evenly. ff_size: Feed-forward network hidden dimension. num_layers: Number of transformer encoder layers. dropout: Dropout probability. Default 0 since RL has enough noise. activation: Activation function for feed-forward layers. output_activation: Optional activation for transformer output.

in_keys: List[str]#

out_keys: List[str]#

input_and_mask_mapping: Dict[str, str] | None = None#

transformer_token_size: int = 512#

latent_dim: int = 512#

num_heads: int = 4#

ff_size: int = 1024#

num_layers: int = 4#

dropout: float = 0.0#

activation: str = 'relu'#

output_activation: str | None = None#

__init__( _target_='protomotions.agents.common.transformer.Transformer', in_keys=<factory>, out_keys=<factory>, input_and_mask_mapping=None, transformer_token_size=512, latent_dim=512, num_heads=4, ff_size=1024, num_layers=4, dropout=0.0, activation='relu', output_activation=None, )#

class protomotions.agents.common.config.DiscreteAutoregressiveTransformerConfig( _target_='protomotions.agents.common.autoregressive.DiscreteAutoregressiveTransformer', in_keys=<factory>, out_keys=<factory>, context_key='context', token_key='tokens', logits_key='logits', generated_tokens_key=None, logprob_key=None, context_embedding_key='_ar_context_embedding', token_embedding_key='_ar_token_embedding', hidden_key='_ar_hidden', context_encoder=<factory>, token_encoder=<factory>, output_head=<factory>, d_model=1024, num_heads=4, num_layers=6, ff_size=4096, dropout=0.1, activation='gelu', num_tokens=1, vocab_size=2, max_seq_len=0, )[source]#

Bases: object

Configuration for categorical autoregressive transformer modules.

Attributes:: in_keys: TensorDict input keys. Defaults to context_key and token_key. out_keys: TensorDict output keys. Defaults to logits_key. context_key: TensorDict key containing context observations. token_key: TensorDict key containing teacher-forced tokens. logits_key: TensorDict key for token logits. generated_tokens_key: TensorDict key for generated token indices. logprob_key: Optional TensorDict key for generated-token log probabilities. context_embedding_key: Internal TensorDict key for context embeddings. token_embedding_key: Internal TensorDict key for token embeddings. hidden_key: Internal TensorDict key for transformer hidden states. context_encoder: Context projection container. token_encoder: Teacher-token projection container. output_head: Logit projection container. d_model: Transformer hidden size. num_heads: Number of attention heads. num_layers: Number of transformer layers. ff_size: Feed-forward hidden size. dropout: Dropout probability. activation: Transformer activation. num_tokens: Number of generated tokens. Use 0 when an owning module derives the token shape before instantiating the transformer. vocab_size: Categorical vocabulary size. Use 0 when an owning module derives the token shape before instantiating the transformer. max_seq_len: Maximum sequence length including context. 0 uses num_tokens + 1.

in_keys: List[str]#

out_keys: List[str]#

context_key: str = 'context'#

token_key: str = 'tokens'#

logits_key: str = 'logits'#

generated_tokens_key: str | None = None#

logprob_key: str | None = None#

context_embedding_key: str = '_ar_context_embedding'#

token_embedding_key: str = '_ar_token_embedding'#

hidden_key: str = '_ar_hidden'#

context_encoder: ModuleContainerConfig#

token_encoder: ModuleContainerConfig#

output_head: ModuleContainerConfig#

d_model: int = 1024#

num_heads: int = 4#

num_layers: int = 6#

ff_size: int = 4096#

dropout: float = 0.1#

activation: str = 'gelu'#

num_tokens: int = 1#

vocab_size: int = 2#

max_seq_len: int = 0#

__init__( _target_='protomotions.agents.common.autoregressive.DiscreteAutoregressiveTransformer', in_keys=<factory>, out_keys=<factory>, context_key='context', token_key='tokens', logits_key='logits', generated_tokens_key=None, logprob_key=None, context_embedding_key='_ar_context_embedding', token_embedding_key='_ar_token_embedding', hidden_key='_ar_hidden', context_encoder=<factory>, token_encoder=<factory>, output_head=<factory>, d_model=1024, num_heads=4, num_layers=6, ff_size=4096, dropout=0.1, activation='gelu', num_tokens=1, vocab_size=2, max_seq_len=0, )#