protomotions.agents.common.fsq module#

Finite scalar quantization modules.

class protomotions.agents.common.fsq.FiniteScalarQuantizer(num_fsq_levels, num_fsq_scalars, eps=0.0001)[source]#

Bases: <Mock object at 0x7fd695670d90>[]

Finite scalar quantizer for continuous latent vectors.

This implements the fixed-codebook scalar quantizer from “Finite Scalar Quantization: VQ-VAE Made Simple” (Mentzer et al., 2023, arXiv:2309.15505). Each latent dimension is independently bounded and rounded to one of num_fsq_levels integer codes, so the implicit codebook is the Cartesian product of the per-dimension scalar levels.

The quantizer intentionally stays a plain nn.Module: it owns buffers and tensor transforms only. The TensorDict/model contract is provided by FSQAutoEncoder, which wraps this module inside an autoencoder bottleneck.

__init__(
num_fsq_levels,
num_fsq_scalars,
eps=0.0001,
)[source]#

Create a scalar quantizer with one shared level count per scalar.

Parameters:

  • num_fsq_levels (int) – Number of discrete scalar values available to each latent scalar. This implementation requires an odd value so zero is one of the quantization levels.

  • num_fsq_scalars (int) – Number of scalar code dimensions in the flattened latent.

  • eps (float) – Small shrink factor used by the tanh bounding transform to avoid saturating exactly at the outermost level before straight-through rounding.

bound(z)[source]#

Map unbounded latent values into the relaxed FSQ code range.

static round_ste(z)[source]#

Round with the straight-through estimator used for FSQ training.

quantize(z)[source]#

Bound and straight-through round latent values to scalar codes.

codes_to_indices(codes)[source]#

Convert centered scalar codes to non-negative FSQ indices.

indices_to_codes(indices)[source]#

Convert non-negative FSQ indices back to centered scalar codes.

calculate_perplexity(codes, skip=False)[source]#

Return mean per-FSQ-scalar perplexity for utilization logging.

class protomotions.agents.common.fsq.FSQAutoEncoder(*args, **kwargs)[source]#

Bases: AutoEncoder

Autoencoder with a finite scalar quantization bottleneck.

supports_log_internals = True#
config: FSQAutoEncoderConfig#
__init__(config)[source]#
property num_fsq_levels: int#
property num_fsq_scalars: int#
property L: MockTensor#
property half_width: MockTensor#
property half_L: MockTensor#
bound(z)[source]#
static round_ste(z)[source]#
codes_to_indices(codes)[source]#
indices_to_codes(indices)[source]#
quantize(z)[source]#
calculate_perplexity(codes, skip=False)[source]#
bottleneck(latent, tensordict)[source]#
internal_logs(latent, tensordict)[source]#