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])

Bin a normalized robot state the way openpi's PaligemmaTokenizer does.

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)[source]#

Bin a normalized robot state the way openpi’s PaligemmaTokenizer does.

That is np.digitize(state, np.linspace(-1, 1, num_bins + 1)[:-1]) - 1: equal bins over [-1, 1], with values below -1 giving -1 and values at or above 1 giving num_bins - 1. lerobot’s pi0.5 processor bins the same way.

Parameters:
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