mstar.model.pi05.components.flow_matching

mstar.model.pi05.components.flow_matching#

Flow-matching utilities for Pi0.5: timestep embedding, Euler step, state discretization.

Functions

discretize_state(state[, num_bins, ...])

Map continuous robot state values in [value_min, value_max] to bin indices.

euler_step(x, velocity, dt)

One Euler integration step: x_{t+dt} = x_t + dt * v(x_t, t).

sincos_timestep_embedding(t, dim, fraction, ...)

Sinusoidal timestep embedding (matches the openpi reference).
mstar.model.pi05.components.flow_matching.discretize_state(state, num_bins=256, value_min=-1.0, value_max=1.0)[source]#

Map continuous robot state values in [value_min, value_max] to bin indices.

Used by Pi0.5 to tokenize proprioceptive state into language tokens that PaliGemma can embed alongside image and language tokens.

Parameters:

  • state (Tensor) – 1D tensor of state values.

  • num_bins (int) – number of discrete bins.

  • value_max (float) – assumed range of normalized state values.

  • value_min (float)

  • value_max

Returns:

1D torch.long tensor of bin indices in [0, num_bins).

Return type:

Tensor

mstar.model.pi05.components.flow_matching.euler_step(x, velocity, dt)[source]#

One Euler integration step: x_{t+dt} = x_t + dt * v(x_t, t).

Parameters:
Return type:

Tensor

mstar.model.pi05.components.flow_matching.sincos_timestep_embedding(t, dim, fraction, output_buffer, min_period=4e-3, max_period=4.0)[source]#

Sinusoidal timestep embedding (matches the openpi reference).

Parameters:

  • t (Tensor) – scalar or 1D tensor of timesteps in [0, 1].

  • dim (int) – embedding dimension. Must be even.

  • fraction (Tensor) – (dim // 2,) frequency ladder, caller-owned so it is allocated once rather than per step (see Pi05LLMSubmodule._get_timestep_emb_fraction).

  • output_buffer (Tensor) – (len(t), dim) destination, written in place and returned; preallocated so the action-gen step stays capturable.

  • max_period (float) – frequency range for the sinusoidal basis.

  • min_period (float)

  • max_period

Returns:

output_buffer, filled with the sin/cos embedding.

Return type:

Tensor