opentau.policies.value.modeling_value

Value Function Model using SIGLIP and Gemma 3 270M

A value function model that estimates state values for reinforcement learning. Uses SIGLIP for vision encoding and Gemma 3 270M for language processing.

Functions

`make_att_2d_masks`(pad_masks, att_masks)	Creates a 2-D attention mask given padding and 1-D attention masks.
`resize_with_pad`(img, width, height[, pad_value])	Resizes an image while preserving aspect ratio and padding to target dimensions.

Classes

`ValueFunction`(config[, dataset_stats, ...])	Wrapper class around Value Function model to train and run inference within OpenTau.
`ValueModel`(config)	Value Function Model using SIGLIP and Gemma 3 270M

class opentau.policies.value.modeling_value.ValueFunction(config: ValueConfig, dataset_stats: dict[str, dict[str, Tensor]] | None = None, per_dataset_stats: list[dict[str, dict[str, Tensor]]] | None = None, dataset_names: list[str] | None = None)[source]

Bases: PreTrainedPolicy

Wrapper class around Value Function model to train and run inference within OpenTau.

__init__(config: ValueConfig, dataset_stats: dict[str, dict[str, Tensor]] | None = None, per_dataset_stats: list[dict[str, dict[str, Tensor]]] | None = None, dataset_names: list[str] | None = None)[source]

Initializes the ValueFunction policy.

Parameters:

config – Value Function configuration class instance.
dataset_stats – Legacy single-dataset stats dict. Wrapped into a singleton list internally. Mutually exclusive with per_dataset_stats.
per_dataset_stats – Ordered list of per-dataset stat dicts. Accepted for compatibility with make_policy’s new plumbing — value is single-dataset by design, so only the first entry is used and a warning fires for longer lists.
dataset_names – Accepted for compatibility with make_policy; only the first entry is consumed.

calculate_value(logits: Tensor) → Tensor[source]

Calculates the expected value from the logits distribution.

Parameters:: logits – Tensor containing the logits for value bins.
Returns:: The expected value.
Return type:: Tensor

config_class: alias of ValueConfig

forward(batch: dict[str, Tensor], return_per_sample: bool = False) → dict[str, Tensor | PerSampleLoss][source]

Do a full training forward pass to compute the value loss.

Parameters:

batch – Dictionary containing observations and target values
return_per_sample – When True, also returns per-sample MSE_per_sample/CE_per_sample (PerSampleLoss) for the validation per-(dataset, control_mode) breakdown. MSE is a zero stub here, so MSE_per_sample carries zero sum and count; the CE pools the value-bin and response-token terms per sample.

Returns:

Dict with “MSE”/”CE”/”L1”/”Accuracy” (plus per-sample CE/MSE entries when return_per_sample is True).

get_optim_params() → dict[source]

Returns the parameters to be optimized.

Returns:: Dictionary of parameters to optimize.
Return type:: dict

name: None = 'value': The name of the policy. Must be defined in subclasses.

predict_value(batch: dict[str, Tensor]) → Tensor[source]

Predict value estimates given environment observations.

Parameters:: batch – Dictionary containing observations (images, state, prompt)
Returns:: Tensor of shape [batch_size, 1] containing value estimates

prepare_discrete_state(batch: dict[str, Tensor]) → list[str][source]

Discretizes the state into bins and converts it to a string representation.

Each dimension of the state vector is discretized into 256 bins. The values of each dimension of the state are expected to be in the range [-1, 1]. The discretization bins are linearly spaced between -1 and 1. The index of the bin for each dimension is then concatenated into a space-separated string.

Parameters:: batch – Batch of data containing the “state” tensor.
Returns:: A list of strings, where each string is a space-separated list of discretized state values.
Raises:: ValueError – If the state values are not normalized between -1 and 1.

prepare_images(batch)[source]

Preprocesses images for the model.

Resizes images to 224x224, pads to keep aspect ratio, and converts pixel range from [0.0, 1.0] to [-1.0, 1.0] as requested by SigLIP. Also handles missing images by creating empty placeholders.

Parameters:: batch – Dictionary containing batch data.
Returns:: A tuple containing a list of image tensors and a list of mask tensors.
Return type:: tuple
Raises:: ValueError – If all image features are missing from the batch.

prepare_language(batch) → tuple[Tensor, Tensor][source]

Tokenizes the text input for the model.

Parameters:: batch – Dictionary containing batch data, including “prompt”.
Returns:: A tuple containing language token tensors and attention mask tensors.
Return type:: tuple

prepare_response(batch: dict[str, Tensor]) → tuple[Tensor, Tensor][source]

Tokenize the response input.

Parameters:

batch – Batch of data containing the key “response”.

Returns:

response_tokens: Tensor of response language tokens.
response_masks: Tensor of response language attention masks.

Return type:

A tuple containing

reset()[source]

Resets the internal state of the policy.

This method is called at the beginning of each episode. For value functions, there is no internal state to reset.

sample_actions(batch: dict[str, Tensor], noise: Tensor | None = None)[source]

Samples actions from the policy.

Parameters:

batch – Dictionary containing observation data.
noise – Optional noise tensor.

Raises:

NotImplementedError – Value functions do not sample actions.

select_action(batch: dict[str, Tensor]) → Tensor[source]

Selects an action based on the current policy.

Parameters:: batch – Dictionary containing observation data.
Returns:: The selected action.
Return type:: Tensor
Raises:: NotImplementedError – Value functions do not select actions.

class opentau.policies.value.modeling_value.ValueModel(config)[source]

Bases: Module

Value Function Model using SIGLIP and Gemma 3 270M

Estimates state values for reinforcement learning by processing: - Images through SIGLIP vision encoder - Language tokens through Gemma 3 270M - Optional state information

┌──────────────────────────────┐ │ value │ │ ▲ │ │ ┌┴─────┐ │ │ │Gemma │ │ │ │3 270M│ │ │ │ │ │ │ ┌──────────┐ └▲──▲──┘ │ │ │ │ │ │ │ │ │ SIGLIP ├──┘ │ │ │ │ │ language │ │ └────▲─────┘ │ │ │ │ │ image(s) │ │ │ └──────────────────────────────┘

__init__(config)[source]

Initializes the ValueModel.

Parameters:: config – Configuration object for the model.

embed_sequence(images, img_masks, lang_tokens, lang_masks, response_tokens: Tensor | None = None, response_masks: Tensor | None = None) → tuple[Tensor, Tensor, Tensor][source]

Embeds sequence of images and language tokens.

Prepares embeddings for SiglipGemmaValueModel transformer processing.

Parameters:

images – List of image tensors.
img_masks – List of image mask tensors.
lang_tokens – Language token tensor.
lang_masks – Language mask tensor.
state – State tensor.

Returns:

A tuple containing embeddings, padding masks, and attention masks.

Return type:

tuple

forward(images: list[Tensor], img_masks: list[Tensor], lang_tokens: Tensor, lang_masks: Tensor, response_tokens: Tensor | None = None, response_masks: Tensor | None = None) → Tensor[source]

Predict value estimates given observations.

Parameters:

images – List of image tensors
img_masks – List of image masks
lang_tokens – Language token IDs
lang_masks – Language attention masks
state – Optional state tensor

Returns:

Tensor of shape [batch_size, 1] containing value estimates

get_value(images: list[Tensor], img_masks: list[Tensor], lang_tokens: Tensor, lang_masks: Tensor) → Tensor[source]

Predict value estimates given observations.

Parameters:

images – List of image tensors
img_masks – List of image masks
lang_tokens – Language token IDs
lang_masks – Language attention masks
state – Optional state tensor

Returns:

Tensor of shape [batch_size, 1] containing value estimates

opentau.policies.value.modeling_value.make_att_2d_masks(pad_masks, att_masks)[source]

Creates a 2-D attention mask given padding and 1-D attention masks.

Tokens can attend to valid inputs tokens which have a cumulative mask_ar smaller or equal to theirs. This way mask_ar int[B, N] can be used to setup several types of attention, for example:

[[1 1 1 1 1 1]]: pure causal attention.

[[0 0 0 1 1 1]]: prefix-lm attention. The first 3 tokens can attend between
themselves and the last 3 tokens have a causal attention. The first entry could also be a 1 without changing behaviour.

[[1 0 1 0 1 0 0 1 0 0]]: causal attention between 4 blocks. Tokens of a
block can attend all previous blocks and all tokens on the same block.

Parameters:

pad_masks – bool[B, N] true if its part of the input, false if padding.
att_masks – int32[B, N] mask that’s 1 where previous tokens cannot depend on it and 0 where it shares the same attention mask as the previous token.

Returns:

The 2D attention masks.

Return type:

torch.Tensor

Raises:

ValueError – If input masks are not 2D.

opentau.policies.value.modeling_value.resize_with_pad(img, width, height, pad_value=-1)[source]

Resizes an image while preserving aspect ratio and padding to target dimensions.

Parameters:

img – Input image tensor of shape (B, C, H, W).
width – Target width.
height – Target height.
pad_value – Value to use for padding. Defaults to -1.

Returns:

Resized and padded image tensor.

Return type:

torch.Tensor

Raises:

ValueError – If image dimensions are not 4D (B, C, H, W).