DiscreteDecoder(decoder_input, observation_space, action_space, action_heads, device, observation_horizon, prediction_horizon, temperature, learnable_temperature, deterministic)
Bases: ActionDecoder
Base class for decoders trained on tokenized action targets.
Shape notation
B: batch size, A: target action-token length, D: token embedding
dimension, V: action-token vocabulary size.
Initialize common discrete-action decoder state.
Source code in src/versatil/models/decoding/decoders/discrete.py
| def __init__(
self,
decoder_input: DecoderInput,
observation_space: ObservationSpace,
action_space: ActionSpace,
action_heads: dict[str, BaseActionHead],
device: str,
observation_horizon: int,
prediction_horizon: int,
temperature: float,
learnable_temperature: bool,
deterministic: bool,
) -> None:
"""Initialize common discrete-action decoder state."""
super().__init__(
decoder_input=decoder_input,
observation_space=observation_space,
action_space=action_space,
action_heads=action_heads,
device=device,
observation_horizon=observation_horizon,
prediction_horizon=prediction_horizon,
)
self.deterministic: bool = deterministic
self.temperature: nn.Parameter = nn.Parameter(
torch.tensor(temperature, dtype=torch.float32),
requires_grad=learnable_temperature,
)
self.token_embedding: nn.Module | None = None
self.vocab_size: int | None = None
|
set_tokenizer(tokenizer=None)
Set tokenizer and bind a vocabulary action head when configured.
Source code in src/versatil/models/decoding/decoders/discrete.py
| def set_tokenizer(self, tokenizer: Tokenizer | None = None) -> None:
"""Set tokenizer and bind a vocabulary action head when configured."""
action_tokenizer = self._require_action_tokenizer(tokenizer=tokenizer)
self.tokenizer = action_tokenizer
if self.action_head_layout == ActionHeadLayout.VOCABULARY:
self._bind_vocabulary_action_tokenizer(action_tokenizer=action_tokenizer)
|