inference_client

inference_client

Inference client connecting trained policies to environments via transports.

EpisodeStatus

Bases: StrEnum

Status values controlling the episode loop.

EnvironmentState dataclass

EnvironmentState(observation_buffer, temporal_aggregator=None)

Per-environment observation buffer and temporal aggregator.

InferenceClient

InferenceClient(policy_runtime, observation_transport, action_transport, temporal_aggregation=False, action_execution_horizon=None, favor_more_recent=True, exponential_decay=0.01, compression_type=value, max_timesteps=800, timing_log=False, update_rate_hz=None, online_explanation_source=None)

Connects a trained policy to a simulation or real-world robot environment via custom transport protocols.

Initialize the inference client.

Parameters:

Name	Type	Description	Default
policy_runtime	PolicyRuntime	Runtime providing policy inference and metadata.	required
observation_transport	ObservationTransport	Protocol for receiving observations from the environment.	required
action_transport	ActionTransport	Protocol for sending actions to the environment.	required
temporal_aggregation	bool	Whether to use temporal ensemble. When enabled, the policy is queried every step and the overlapping action predictions from consecutive chunks are averaged with exponential weighting. When disabled, the policy predicts a chunk of actions and all of them are sent to the environment before re-querying.	False
action_execution_horizon	int | None	How many actions from each predicted chunk to execute when temporal_aggregation is False. Defaults to prediction_horizon.	None
favor_more_recent	bool	In temporal ensemble, weight predictions from later inference calls higher than earlier ones for the same timestep.	True
exponential_decay	float	Exponential decay rate for the temporal ensemble weights. Higher values discount older predictions more aggressively.	0.01
compression_type	str	Compression type for image data transfer.	value
max_timesteps	int	Maximum episode length for temporal aggregation.	800
timing_log	bool	Whether to log per-step timing breakdown.	False
update_rate_hz	float | None	Target action-send frequency in Hz.	None
online_explanation_source	OnlineExplanationSource | None	Optional source that converts the exact ready observation batch used for policy inference into explanations.	None

Source code in src/versatil/inference/inference_client.py

def __init__(
    self,
    policy_runtime: PolicyRuntime,
    observation_transport: ObservationTransport,
    action_transport: ActionTransport,
    temporal_aggregation: bool = False,
    action_execution_horizon: int | None = None,
    favor_more_recent: bool = True,
    exponential_decay: float = 0.01,
    compression_type: str = CompressionType.RAW.value,
    max_timesteps: int = 800,
    timing_log: bool = False,
    update_rate_hz: float | None = None,
    online_explanation_source: OnlineExplanationSource | None = None,
) -> None:
    """Initialize the inference client.

    Args:
        policy_runtime: Runtime providing policy inference and metadata.
        observation_transport: Protocol for receiving observations from the environment.
        action_transport: Protocol for sending actions to the environment.
        temporal_aggregation: Whether to use temporal ensemble. When enabled,
            the policy is queried every step and the overlapping action predictions
            from consecutive chunks are averaged with exponential weighting.
            When disabled, the policy predicts a chunk of actions and all of them
            are sent to the environment before re-querying.
        action_execution_horizon: How many actions from each predicted chunk to execute
            when temporal_aggregation is False. Defaults to prediction_horizon.
        favor_more_recent: In temporal ensemble, weight predictions from later inference
            calls higher than earlier ones for the same timestep.
        exponential_decay: Exponential decay rate for the temporal ensemble weights.
            Higher values discount older predictions more aggressively.
        compression_type: Compression type for image data transfer.
        max_timesteps: Maximum episode length for temporal aggregation.
        timing_log: Whether to log per-step timing breakdown.
        update_rate_hz: Target action-send frequency in Hz.
        online_explanation_source: Optional source that converts the exact
            ready observation batch used for policy inference into
            explanations.
    """
    self.policy_runtime = policy_runtime
    self.observation_transport = observation_transport
    self.action_transport = action_transport
    self.temporal_aggregation = temporal_aggregation
    self.action_execution_horizon = (
        action_execution_horizon
        if action_execution_horizon is not None
        else policy_runtime.prediction_horizon
    )
    if (
        self.action_execution_horizon > 1
        and policy_runtime.observation_horizon > 1
        and not temporal_aggregation
    ):
        logging.warning(
            f"Executing {self.action_execution_horizon} actions per "
            f"observation with "
            f"observation_horizon={policy_runtime.observation_horizon}: "
            f"the policy's history will hold frames "
            f"{self.action_execution_horizon} steps apart, while training "
            "windows are contiguous."
        )
    if self.action_execution_horizon > policy_runtime.prediction_horizon:
        raise ValueError(
            f"action_execution_horizon ({self.action_execution_horizon}) cannot exceed "
            f"prediction_horizon ({policy_runtime.prediction_horizon})."
        )
    self.compression_type = compression_type
    self.timing_log = timing_log
    self.update_rate_hz = update_rate_hz
    self.online_explanation_source = online_explanation_source
    self.timestep = 0
    observation_space = policy_runtime.observation_space
    action_space = policy_runtime.action_space
    self.camera_keys, self.state_keys, self.has_language = (
        self._bucket_observation_keys(observation_space=observation_space)
    )
    self.all_observation_keys = list(observation_space.observations_metadata.keys())
    self.action_keys_to_dimensions = {
        key: metadata.prediction_dimension
        for key, metadata in action_space.actions_metadata.items()
        if metadata.requires_prediction_head
    }
    self.observation_preprocessor = ObservationPreprocessor(
        camera_keys=self.camera_keys,
        state_keys=self.state_keys,
        has_language=self.has_language,
        camera_metadata=policy_runtime.observation_space.cameras,
        compression_type=compression_type,
        rotate_images=infer_rotate_images(config=policy_runtime.config),
        depth_clamp_ranges=policy_runtime.depth_clamp_ranges,
        state_dtypes={
            key: metadata.dtype
            for key, metadata in (observation_space.observations_metadata.items())
            if key in self.state_keys and metadata.is_numerical
        },
    )
    self.action_postprocessor = ActionPostprocessor(
        action_space=action_space,
        denoising_thresholds=policy_runtime.denoising_thresholds,
    )
    self._temporal_config = {
        "favor_more_recent": favor_more_recent,
        "exponential_decay": exponential_decay,
        "max_timesteps": max_timesteps,
    }
    self.environment_states: dict[int, EnvironmentState] = {}

run_episode

run_episode(max_steps)

Run a full inference episode.

Parameters:

Name	Type	Description	Default
max_steps	int	Maximum number of steps in the episode.	required

Source code in src/versatil/inference/inference_client.py

def run_episode(self, max_steps: int) -> None:
    """Run a full inference episode.

    Args:
        max_steps: Maximum number of steps in the episode.
    """
    self.observation_transport.register(
        client_name=self.policy_runtime.client_identifier
    )
    for _step_idx in range(max_steps):
        try:
            status = self.step()
        except Exception:
            logging.exception(f"Fatal error at step {_step_idx}")
            raise
        if status == EpisodeStatus.FINISHED.value:
            break

step

step()

Execute a single step: receive, predict, send.

Returns:

Type	Description
str	Status string indicating episode state.

Source code in src/versatil/inference/inference_client.py

def step(self) -> str:
    """Execute a single step: receive, predict, send.

    Returns:
        Status string indicating episode state.
    """
    step_start = time.time() if self.timing_log else None

    if self.timing_log:
        preprocessing_start = time.time()

    status = self._ingest_observation()
    if status != EpisodeStatus.CONTINUE.value:
        return status

    if self.timing_log:
        preprocessing_duration = time.time() - preprocessing_start
        inference_start = time.time()

    actions_by_environment = self._get_actions_for_ready_environments()

    if self.timing_log:
        inference_duration = time.time() - inference_start
        postprocessing_start = time.time()

    if actions_by_environment:
        action_metadata = self.action_postprocessor.build_action_metadata()
        # All environments have the same number of actions per chunk.
        # Send one step at a time — the server steps the environment on each send.
        num_steps = len(next(iter(actions_by_environment.values())))
        for step in range(num_steps):
            step_actions = {
                env_idx: action_list[step]
                for env_idx, action_list in actions_by_environment.items()
            }
            self.action_transport.send(
                actions=step_actions,
                action_metadata=action_metadata,
            )
            if self.update_rate_hz is not None:
                time.sleep(1.0 / self.update_rate_hz)

    if self.timing_log:
        postprocessing_duration = time.time() - postprocessing_start
        total_duration = time.time() - step_start
        logging.info(
            f"[TIMING] Step {self.timestep}: "
            f"preprocess={preprocessing_duration:.4f}s "
            f"inference={inference_duration:.4f}s "
            f"postprocess={postprocessing_duration:.4f}s "
            f"total={total_duration:.4f}s "
            f"fps={1.0 / total_duration:.1f}"
        )

    self.timestep += 1

    return EpisodeStatus.CONTINUE.value

reset

reset()

Reset all environment states for a new episode.

Source code in src/versatil/inference/inference_client.py

def reset(self) -> None:
    """Reset all environment states for a new episode."""
    for state in self.environment_states.values():
        state.observation_buffer.reset()
        if state.temporal_aggregator is not None:
            state.temporal_aggregator.reset()
    self.environment_states.clear()

shutdown

shutdown()

Close transport connections.

Source code in src/versatil/inference/inference_client.py

def shutdown(self) -> None:
    """Close transport connections."""
    try:
        self.observation_transport.close()
    except Exception:
        logging.warning("Error closing observation transport", exc_info=True)
    try:
        self.action_transport.close()
    except Exception:
        logging.warning("Error closing action transport", exc_info=True)

infer_rotate_images

infer_rotate_images(config)

Whether incoming simulator images must be rotated 180 degrees.

The LIBERO simulator emits images flipped relative to the lerobot dataset orientation the policy was trained on, so the mismatch is a property of the training data source, derived from the checkpoint's dataset schema.

Parameters:

Name	Type	Description	Default
config	DictConfig	Full training config saved with the checkpoint, either as the composed DictConfig or as an instantiated MainConfig whose schema node is a real DatasetSchema.	required

Returns:

Type	Description
bool	True for policies trained on LIBERO lerobot datasets.

Source code in src/versatil/inference/inference_client.py

def infer_rotate_images(config: DictConfig) -> bool:
    """Whether incoming simulator images must be rotated 180 degrees.

    The LIBERO simulator emits images flipped relative to the lerobot dataset
    orientation the policy was trained on, so the mismatch is a property of
    the training data source, derived from the checkpoint's dataset schema.

    Args:
        config: Full training config saved with the checkpoint, either as the
            composed DictConfig or as an instantiated MainConfig whose schema
            node is a real DatasetSchema.

    Returns:
        True for policies trained on LIBERO lerobot datasets.
    """
    if isinstance(config, DictConfig):
        schema_node = OmegaConf.select(config, "task.dataset_schema")
        if schema_node is None:
            return False
        if isinstance(schema_node, DictConfig):
            target = str(schema_node.get("_target_", ""))
            dataset_type = str(schema_node.get("dataset_type", ""))
            return (
                target.endswith("LeRobotDatasetSchemaV30")
                and dataset_type == DatasetType.LIBERO.value
            )
        schema = schema_node
    else:
        schema = config.task.dataset_schema
    return (
        isinstance(schema, LeRobotDatasetSchemaV30)
        and schema.dataset_type == DatasetType.LIBERO.value
    )