Embodied Data Interface#

The embodied data interface separates runtime communication from Actor-facing training data:

EnvOutput(obs, EnvTransition)
    -> PolicyInput (+ previous EnvPart)
    -> actions tensor

PolicyPart + completed EnvPart
    -> TrajectoryCollector
    -> Trajectory / episode shard / pipeline micro-batch

See Understand the Trajectory Collector for the lifecycle, data ownership, and execution-mode semantics.

Shape notation#

The class field comments use the following symbols:

  • B is the routed environment batch. Before decoupled merging, it is normally env.train.total_num_envs / env_world_size / rollout.pipeline_stage_num.

  • C is actor.model.num_action_chunks; eval-only runs use the Rollout model config instead.

  • A is actor.model.action_dim and D = C * A. With E = env.train.rollout_epoch, a full trajectory normally has T = E * env.train.max_steps_per_rollout_epoch / C chunks.

Environment and policy messages#

EnvOutput composes observations with one EnvTransition; it does not duplicate reward and boundary fields. PolicyInput carries observations and may piggyback the preceding EnvPart. Rollout sends only the action tensor back to Env. The full PolicyOutput stays on the trajectory path because it also contains log-probabilities, values, versions, and training inputs.

class rlinf.data.schema.embodied_types.EnvOutput#

Observation and transition returned by one environment chunk step.

property rewards: Tensor | None#

Return per-action environment rewards from the transition.

property dones: Tensor | None#

Return combined termination/truncation flags.

property terminations: Tensor | None#

Return natural terminal flags.

property truncations: Tensor | None#

Return time-limit or external truncation flags.

property intervene_actions: Tensor | None#

Return expert actions recorded by the environment.

property intervene_flags: Tensor | None#

Return which action slots used an expert action.

property rlt_switch_flags: Tensor | None#

Return per-action RLT route choices.

prepare_observations(obs)#

Select policy inputs and normalize task descriptions.

Parameters:

obs (dict[str, Any])

Return type:

dict[str, Any]

__init__(*, obs, transition=<factory>, final_obs=None, env_infos=None)#
Parameters:

  • obs (dict[str, Any])

  • transition (EnvTransition)

  • final_obs (dict[str, Any] | None)

  • env_infos (dict[str, Any] | None)

Return type:

None

class rlinf.data.schema.embodied_types.EnvTransition#

Environment-owned outcome of one executed action chunk.

with_trajectory_data(*, reward_model_output=None, reward_assign_lengths=None, episode_data=None)#

Return a copy enriched with collector-facing metadata.

Parameters:

  • reward_model_output (Tensor | None)

  • reward_assign_lengths (list[int] | None)

  • episode_data (dict[str, Any] | None)

Return type:

EnvTransition

to_dict()#

Expose transition fields for existing environment batch consumers.

Return type:

dict[str, Any]

split(split_sizes)#

Split this environment transition on its batch dimension.

Parameters:

split_sizes (list[int])

Return type:

list[EnvTransition]

classmethod merge(transitions)#

Merge source fragments in batch order.

Parameters:

transitions (list[EnvTransition])

Return type:

EnvTransition

compute_rewards(*, bootstrap_values, env_reward_weight, reward_weight, auto_reset, bootstrap_type, gamma)#

Build the reward tensor consumed by trajectory training.

Reward-model scores are weighted into environment rewards first. Auto-reset truncations then receive the configured bootstrap value.

The result has independent storage for later in-place reward updates.

Parameters:

  • bootstrap_values (Tensor | None)

  • env_reward_weight (float)

  • reward_weight (float)

  • auto_reset (bool)

  • bootstrap_type (str)

  • gamma (float)

Return type:

Tensor | None

assign_history_rewards(steps, *, reward_weight)#

Assign delayed model rewards to the requested preceding steps.

Parameters:
Return type:

None

__init__(*, rewards=None, dones=None, terminations=None, truncations=None, executed_mask=None, intervene_actions=None, intervene_flags=None, rlt_switch_flags=None, reward_model_output=None, reward_assign_lengths=None, episode_data=None)#
Parameters:

  • rewards (Tensor | None)

  • dones (Tensor | None)

  • terminations (Tensor | None)

  • truncations (Tensor | None)

  • executed_mask (Tensor | None)

  • intervene_actions (Tensor | None)

  • intervene_flags (Tensor | None)

  • rlt_switch_flags (Tensor | None)

  • reward_model_output (Tensor | None)

  • reward_assign_lengths (list[int] | None)

  • episode_data (dict[str, Any] | None)

Return type:

None

class rlinf.data.schema.embodied_types.PolicyInput#

One routed policy request, optionally carrying completed env data.

property requires_inference: bool#

Return whether Rollout must infer actions for this request.

split(split_sizes)#

Split this producer request on its batch dimension.

Parameters:

split_sizes (list[int])

Return type:

list[PolicyInput]

classmethod merge(policy_inputs)#

Merge routed policy requests in source order.

Parameters:

policy_inputs (list[PolicyInput])

Return type:

PolicyInput

__init__(*, obs, rlt_switch_flags=None, intervene_flags=None, sources=<factory>, env_parts=<factory>, request_sizes=<factory>, is_last=False, external_actions=None)#
Parameters:

  • obs (dict[str, Any])

  • rlt_switch_flags (Tensor | None)

  • intervene_flags (Tensor | None)

  • sources (list[TrajectorySource])

  • env_parts (list[EnvPart | None])

  • request_sizes (list[int])

  • is_last (bool)

  • external_actions (Tensor | None)

Return type:

None

class rlinf.data.schema.embodied_types.PolicyOutput#

Policy inference output retained until its environment transition arrives.

split(split_sizes)#

Split model outputs on their leading batch dimension.

Parameters:

split_sizes (list[int])

Return type:

list[PolicyOutput]

classmethod merge(outputs)#

Merge routed model-output fragments in batch order.

Parameters:

outputs (list[PolicyOutput])

Return type:

PolicyOutput

__init__(*, forward_inputs, prev_logprobs=None, prev_values=None, intervene_flags=None, versions=None)#
Parameters:

  • forward_inputs (dict[str, Any])

  • prev_logprobs (Tensor | None)

  • prev_values (Tensor | None)

  • intervene_flags (Tensor | None)

  • versions (Tensor | None)

Return type:

None

Trajectory parts#

PolicyPart and EnvPart are the only Actor channel input types. They are matched by TrajectoryKey. TrajectorySource preserves source size and offset when channel routing splits a batch. Env creates an incomplete EnvPart after stepping; Rollout adds optional terminal inference data and publishes the same type.

class rlinf.data.schema.embodied_types.TrajectoryKey#

Identity of one action chunk produced by a logical environment source.

__init__(step_id, epoch_id, env_rank, stage_id, chunk_id)#
Parameters:

  • step_id (int)

  • epoch_id (int)

  • env_rank (int)

  • stage_id (int)

  • chunk_id (int)

Return type:

None

class rlinf.data.schema.embodied_types.TrajectorySource#

Trajectory key and batch size carried by one routed source shard.

classmethod split(sources, split_sizes)#

Split routed source metadata along the batch dimension.

Parameters:
Return type:

list[list[TrajectorySource]]

classmethod merge(source_groups)#

Merge adjacent shards while preserving logical source boundaries.

Parameters:

source_groups (list[list[TrajectorySource]])

Return type:

list[TrajectorySource]

__init__(key, size, offset=0)#
Parameters:
Return type:

None

class rlinf.data.schema.embodied_types.PolicyPart#

Policy-owned data for one or more routed action chunks.

Exactly one of output and external_actions must be present. The latter represents an online intervention that skipped model inference.

property inferred: bool#

Return whether this part came from model inference.

split()#

Split this routed payload into one fragment per logical source.

Return type:

list[PolicyPart]

classmethod merge(fragments)#

Merge single-source fragments in their validated batch order.

Parameters:

fragments (list[PolicyPart])

Return type:

PolicyPart

__init__(*, sources, obs, output=None, external_actions=None)#
Parameters:
Return type:

None

class rlinf.data.schema.embodied_types.EnvPart#

Environment-owned half of one or more routed action chunks.

Env creates this object after stepping and attaches it to the next PolicyInput. Rollout adds terminal inference data with complete() before publishing the same type to the Actor channel.

complete(*, next_obs, next_rlt_obs, final_prev_values=None)#

Return the Actor-ready part after resolving rollout-owned fields.

Parameters:

  • next_obs (dict[str, Any] | None)

  • next_rlt_obs (dict[str, Any] | None)

  • final_prev_values (Tensor | None)

Return type:

EnvPart

split(split_sizes=None)#

Split by explicit route sizes or by the represented sources.

Parameters:

split_sizes (list[int] | None)

Return type:

list[EnvPart]

classmethod merge(fragments)#

Merge single-source fragments in their validated batch order.

Parameters:

fragments (list[EnvPart])

Return type:

EnvPart

compute_rewards(*, env_reward_weight, reward_weight, auto_reset, bootstrap_type, gamma)#

Return the training rewards represented by this environment part.

Parameters:

  • env_reward_weight (float)

  • reward_weight (float)

  • auto_reset (bool)

  • bootstrap_type (str)

  • gamma (float)

Return type:

Tensor | None

__init__(*, sources, transition, next_obs=None, requires_inference=False, next_rlt_obs=None, bootstrap_values=None, final_prev_values=None, initial_transition=None)#
Parameters:

  • sources (list[TrajectorySource])

  • transition (EnvTransition)

  • next_obs (dict[str, Any] | None)

  • requires_inference (bool)

  • next_rlt_obs (dict[str, Any] | None)

  • bootstrap_values (Tensor | None)

  • final_prev_values (Tensor | None)

  • initial_transition (EnvTransition | None)

Return type:

None

Collection#

TrajectoryPlan validates the configured output mode and derives rollout geometry. TrajectoryCollector is the one public channel collector for every embodied mode.

class rlinf.data.schema.embodied_trajectory.TrajectoryMode#

Actor-facing output mode for collected trajectory chunks.

__new__(value)#
class rlinf.data.schema.embodied_trajectory.TrajectoryPlan#

Validated trajectory mode and routing geometry for one run.

classmethod from_cfg(cfg)#

Build the single source of truth for trajectory collection.

Parameters:

cfg (DictConfig)

Return type:

TrajectoryPlan

static mode_from_cfg(cfg)#

Validate mode combinations and return the configured output mode.

Parameters:

cfg (DictConfig)

Return type:

TrajectoryMode

property dispatcher: str | None#

Return the channel dispatcher compatible with the output routing.

__init__(mode, geometry)#
Parameters:
Return type:

None

class rlinf.data.schema.embodied_trajectory.TrajectoryCollector#

Join trajectory parts and emit outputs selected by one validated plan.

setup(ctx)#

Initialize the join state and private output strategy.

Parameters:

ctx (ChannelContext)

Return type:

None

collect(item, key)#

Accept one part and yield any actor outputs it completes.

Parameters:

  • item (Any)

  • key (str)

Return type:

Iterable[tuple[str, Any]]

Actor output#

TrajectoryStep resolves one joined policy/environment pair, including reward, intervention, and transition semantics. Trajectory then stacks steps into [T, B, ...] Actor tensors. Boundary and final-value sequences may have T + E entries.

class rlinf.data.schema.embodied_types.TrajectoryStep#

Complete training data for one joined action chunk.

classmethod from_parts(policy, env, *, rewards, collect_prev_infos, collect_transitions, enable_rlt, include_final_value)#

Resolve one joined policy/env pair into self-contained step data.

This is the single conversion boundary for policy statistics, environment boundaries, interventions, and optional transitions. An execution mask requires logprobs shaped [B, C, …]. Models receive it in forward_inputs and must also mask their recomputed statistics.

Parameters:

  • policy (PolicyPart)

  • env (EnvPart)

  • rewards (Tensor | None)

  • collect_prev_infos (bool)

  • collect_transitions (bool)

  • enable_rlt (bool)

  • include_final_value (bool)

Return type:

TrajectoryStep

apply_interventions(intervene_actions, intervene_flags)#

Replace executed actions and synchronize their training metadata.

Parameters:

  • intervene_actions (Tensor)

  • intervene_flags (Tensor | None)

Return type:

None

set_intervention_flags(intervene_flags)#

Expand chunk-level flags to the flattened action layout.

Parameters:

intervene_flags (Tensor)

Return type:

None

set_transition_observations(policy, env, *, collect_transitions, enable_rlt)#

Extract the transition representation required by the algorithm.

Parameters:
Return type:

None

add_reward(env_id, reward)#

Add one delayed reward to an environment in this step.

Parameters:

  • env_id (int)

  • reward (Tensor)

Return type:

None

__init__(*, actions=None, intervene_flags=None, prev_logprobs=None, prev_values=None, dones=None, truncations=None, terminations=None, rewards=None, forward_inputs=<factory>, versions=None, curr_obs=<factory>, next_obs=<factory>, initial_dones=None, initial_truncations=None, initial_terminations=None, final_prev_values=None)#
Parameters:

  • actions (Tensor | None)

  • intervene_flags (Tensor | None)

  • prev_logprobs (Tensor | None)

  • prev_values (Tensor | None)

  • dones (Tensor | None)

  • truncations (Tensor | None)

  • terminations (Tensor | None)

  • rewards (Tensor | None)

  • forward_inputs (dict[str, Any])

  • versions (Tensor | None)

  • curr_obs (dict[str, Any])

  • next_obs (dict[str, Any])

  • initial_dones (Tensor | None)

  • initial_truncations (Tensor | None)

  • initial_terminations (Tensor | None)

  • final_prev_values (Tensor | None)

Return type:

None

class rlinf.data.schema.embodied_types.Trajectory#

trajectory contains multiple episodes.

classmethod from_steps(steps, *, max_episode_length=0)#

Stack complete chunk steps into one actor-facing trajectory.

Epoch-start boundaries and final bootstrap values are inserted in temporal order; ordinary fields retain one entry per action chunk.

Parameters:
Return type:

Trajectory

split(split_sizes)#

Split the batch dimension into actor-consumable shards.

An integer requests equal shards. A list provides each shard size explicitly, as required by pipeline actor routing.

Parameters:

split_sizes (int | list[int])

Return type:

list[Trajectory]

classmethod to_batch(trajectories)#

Concatenate trajectories into one [T, B, ...] Actor batch.

Parameters:

trajectories (list[Trajectory])

Return type:

dict[str, Any]

__init__(max_episode_length=0, model_weights_id='', actions=None, intervene_flags=None, rewards=None, terminations=None, truncations=None, dones=None, prev_logprobs=None, prev_values=None, versions=None, forward_inputs=<factory>, curr_obs=<factory>, next_obs=<factory>)#
Parameters:

  • max_episode_length (int)

  • model_weights_id (str)

  • actions (Tensor | None)

  • intervene_flags (Tensor | None)

  • rewards (Tensor | None)

  • terminations (Tensor | None)

  • truncations (Tensor | None)

  • dones (Tensor | None)

  • prev_logprobs (Tensor | None)

  • prev_values (Tensor | None)

  • versions (Tensor | None)

  • forward_inputs (dict[str, Any])

  • curr_obs (dict[str, Any])

  • next_obs (dict[str, Any])

Return type:

None

On this page