opentau.policies.value.siglip_gemma

SigLip + Gemma Model for Value Function Estimation.

This module defines the configuration and model classes for a value function estimator that combines a SigLIP vision encoder and a Gemma language model.

Classes

`SiglipGemmaValueConfig`([siglip_config, ...])	Configuration class for SiglipGemmaValueModel.
`SiglipGemmaValueModel`(config)	SigLIP + Gemma Model for Value Function Estimation.

class opentau.policies.value.siglip_gemma.SiglipGemmaValueConfig(siglip_config: dict | None = None, gemma_config: dict | None = None, num_value_bins: int = 201, response_max_length: int = 52, **kwargs)[source]

Bases: PretrainedConfig

Configuration class for SiglipGemmaValueModel.

Parameters:

siglip_config – Configuration for the SigLIP vision model.
gemma_config – Configuration for the Gemma language model.
num_value_bins – Number of bins for value discretization. Defaults to 201.
**kwargs – Additional keyword arguments passed to PretrainedConfig.

__init__(siglip_config: dict | None = None, gemma_config: dict | None = None, num_value_bins: int = 201, response_max_length: int = 52, **kwargs)[source]

model_type: str = 'SiglipGemmaValueModel'

sub_configs: dict[str, type['PretrainedConfig']] = {'gemma_config': <class 'transformers.models.auto.configuration_auto.AutoConfig'>, 'siglip_config': <class 'transformers.models.auto.configuration_auto.AutoConfig'>}

class opentau.policies.value.siglip_gemma.SiglipGemmaValueModel(config: SiglipGemmaValueConfig)[source]

Bases: PreTrainedModel

SigLIP + Gemma Model for Value Function Estimation.

This model combines a SigLIP vision encoder and a Gemma language model to estimate state values. It projects the final hidden state to a set of discretized value bins.

Parameters:: config – Configuration object of type SiglipGemmaValueConfig.

__init__(config: SiglipGemmaValueConfig)[source]

Initializes the SiglipGemmaValueModel.

Parameters:: config – Configuration object of type SiglipGemmaValueConfig.

config_class: alias of SiglipGemmaValueConfig

embed_image(image: Tensor) → Tensor[source]

Embeds images using the SIGLIP vision encoder.

Parameters:: image – Tensor containing image pixel values.
Returns:: The embedded image features.
Return type:: torch.Tensor

embed_language_tokens(tokens: Tensor) → Tensor[source]

Embeds language tokens using the Gemma embedding layer.

Parameters:: tokens – Tensor containing language token IDs.
Returns:: The embedded language tokens.
Return type:: torch.Tensor

forward(inputs_embeds: FloatTensor, attention_mask: Tensor, position_ids: LongTensor) → Tensor[source]

Forward pass that processes vision and language inputs and outputs a value.

Parameters:

inputs_embeds – Tensor of shape [batch_size, sequence_length, embedding_dim] containing the combined embeddings of images and text.
attention_mask – Attention mask for the sequence.
position_ids – Position IDs for RoPE.

Returns:

Logits for discretized values of shape [batch_size, num_value_bins].

Return type:

torch.Tensor

get_value(inputs_embeds: FloatTensor, attention_mask: Tensor, position_ids: LongTensor) → Tensor[source]

Forward pass that processes vision and language inputs and outputs a value.

Parameters:

inputs_embeds – Tensor of shape [batch_size, sequence_length, embedding_dim] containing the combined embeddings of images and text.
attention_mask – Attention mask for the sequence.
position_ids – Position IDs for RoPE.

Returns:

Logits for discretized values of shape [batch_size, num_value_bins].

Return type:

torch.Tensor