mstar.model.pi05.components.normalization

mstar.model.pi05.components.normalization#

Per-deployment normalization of the robot state and actions.

openpi’s policy server quantile-normalizes the raw robot state before it is binned into the prompt, and maps the model’s actions back to robot units (transforms.Normalize / Unnormalize with use_quantiles=True, which pi0.5 uses). A deployment that gives the same statistics in its serving config takes the raw state from clients and returns raw actions, like openpi does.

Functions

load_norm_stats(spec)

Parse the norm_stats model kwarg.

Classes

Pi05NormStats([state, actions])

QuantileStats(q01, q99)

Per-dimension 1st and 99th percentiles, which normalization maps to -1 and 1.
class mstar.model.pi05.components.normalization.Pi05NormStats(state: 'QuantileStats | None' = None, actions: 'QuantileStats | None' = None)[source]#

Bases: object

Parameters:
actions: QuantileStats | None = None#
normalize_state(state)[source]#

Quantile-normalize a raw state. Like openpi, the stats are cut to the state’s length, so a robot with fewer dims than the stats is fine.

Parameters:

state (Tensor)

Return type:

Tensor

state: QuantileStats | None = None#
unnormalize_actions(actions)[source]#

Map the dims the stats cover back to robot units. The rest, the padding up to action_dim, pass through as openpi leaves them.

Parameters:

actions (Tensor)

Return type:

Tensor

class mstar.model.pi05.components.normalization.QuantileStats(q01, q99)[source]#

Bases: NamedTuple

Per-dimension 1st and 99th percentiles, which normalization maps to -1 and 1.

Parameters:
q01: Tensor#

Alias for field number 0

q99: Tensor#

Alias for field number 1

mstar.model.pi05.components.normalization.load_norm_stats(spec)[source]#

Parse the norm_stats model kwarg.

spec is either a mapping laid out like openpi’s norm_stats.json (state and/or actions, each with q01 and q99 lists, other fields such as mean / std ignored, optionally under a top-level norm_stats key) or the path of such a JSON file. A file’s trailing dims with q01 == q99 are openpi’s padding and are dropped. Inline stats are taken as written.

Parameters:

spec (dict | str | None)

Return type:

Pi05NormStats | None