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.

``models`` is execution order, not a dependency graph. Nothing sorts this list. A module that produces a key must sit before the module that consumes it; appending a producer to the end of the list does not work, however clearly its output names say what it feeds. find_key_flow_errors catches that specific mistake without building anything – see its docstring for why that matters.

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]#

find_key_flow_errors()[source]#

Return descriptions of producer-after-consumer ordering mistakes.

A module needs a key, no earlier module declares it as an output, and a LATER module does. That is unambiguous: the producer exists, it is simply in the wrong place, and no amount of runtime context can fix it.

Deliberately narrow. It does NOT report keys that are never produced at all, because a few modules derive their real keys at construction time rather than declaring them in config (DiscreteLatent takes its in_keys from its encoder, for instance), so “not declared anywhere” does not imply “missing”. Restricting the check to keys with a known later producer makes false positives impossible, which is what lets it run over every experiment config in the repo.

The value is that this needs no torch, no lazy-layer materialization, no GPU and no environment. ModuleContainer.__init__ performs the authoritative check, but only while building the modules – so a config that is wrong in this way otherwise gets discovered at model-construction time, which in a distributed job means one rank dying and the rest deadlocking at the next collective.

Returns:: One human-readable string per mistake found; empty when the ordering is sound.
Return type:: 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, )#