SparkRT

SparkRT is a low-latency inference SDK for SparkMind2 robot policies. Load a trained checkpoint and run it from your control loop in a few lines - no knowledge of SparkMind2 internals, tensor layouts, or preprocessor key names required:

import sparkrt
from sparkrt import Observation

policy = sparkrt.load_policy("/path/to/checkpoints/020000/pretrained_model")

obs = Observation(
    images={"image": rgb_frame, "image2": wrist_frame},  # numpy HWC RGB uint8
    state=robot_state,                                    # 1-D float array
    instruction="pick up the object and place it in the basket",
)
action = policy.step(obs)   # -> numpy action for your robot

Under the hood SparkRT wraps trained SparkMind2 agents and optimizes the action-inference hot path without changing SparkMind2 model definitions, checkpoint formats, or pre/post-processing logic:

• SparkMind2 owns training, checkpoints, model modules, normalization, tokenizers, environment logic, and evaluation metrics.
• SparkRT owns the ergonomic Policy SDK, inference sessions, action queues, runtime configuration, and pluggable execution backends.

Detailed benchmark results and experiment history are kept in docs/EXPERIMENTS.md. Backend integration guidance is in docs/INTEGRATION.md.

Status

SparkRT currently supports:

• ACT policies (single-shot action chunking, optional temporal ensemble).
• Pi0.5 policies (vision-language-action, flow-matching denoising).
• Eager PyTorch execution.
• CUDA Graph execution for fixed-shape hot regions.
• TorchCompile execution for fused PyTorch/Inductor kernels.
• Preset-based runtime configuration.
• A SparkMind2 LIBERO eval wrapper for development certification.

The public SDK is Python-first. Native C++/CUDA or TensorRT backends can be added later behind the same ExecutionBackend contract.

Installation

Install inside the SparkMind2 environment:

source /path/to/SparkMind2/.venv/bin/activate
pip install -e /path/to/SparkRT

SparkMind2 is a local dependency and should be installed or available on PYTHONPATH separately. torch is intentionally not pinned by SparkRT so the runtime can use the CUDA/Torch build provided by the target machine.

YAML presets require the optional preset extra:

pip install -e '/path/to/SparkRT[presets]'

Loading a Policy

Load a SparkMind2 pretrained_model checkpoint directly:

import sparkrt
from sparkrt import RuntimeConfig

policy = sparkrt.load_policy(
    "/path/to/checkpoints/020000/pretrained_model",
    config=RuntimeConfig.from_preset("latency"),  # optional; no config defaults to eager
)

Or wrap an agent that SparkMind2 has already constructed:

import sparkrt

policy = sparkrt.from_sparkmind_agent(agent, config=RuntimeConfig.from_preset("quality"))

Both return a Policy. Inspect what the policy expects before building inputs:

policy.camera_names          # e.g. ("image", "image2")  or  ("top",)
policy.state_dim             # length of the state vector (0 if unused)
policy.requires_instruction  # True for Pi0.5, False for ACT
policy.action_dim            # environment action dimensionality

Building an Observation

policy.step(obs) takes an Observation describing one timestep. You supply raw camera frames and the robot state with friendly camera names; SparkRT handles the channel-order, dtype, normalization, batching, and tokenization internally.

Field	Type	Notes
images	dict[str, ndarray]	Keyed by policy.camera_names. Each frame: NumPy (or torch) HWC RGB, uint8 [0,255] or float [0,1]. CHW is also accepted. A single-camera policy may take a bare array.
state	1-D array	Length policy.state_dim. Any range - SparkRT applies the checkpoint's normalization.
instruction	str or None	Required when policy.requires_instruction is True (Pi0.5). Ignored by ACT.

Image frames are RGB with no batch dimension; SparkRT converts HWC -> CHW, scales uint8 -> [0,1], and lets the checkpoint preprocessor handle batching.

ACT example (images + state, no language)

ACT is a deterministic single-camera policy with no instruction:

import numpy as np
import sparkrt
from sparkrt import Observation

policy = sparkrt.load_policy("/path/to/act/checkpoints/200000/pretrained_model")
# policy.camera_names == ("top",); policy.state_dim == 14; requires_instruction == False

policy.reset()  # at the start of each episode
for _ in range(horizon):
    obs = Observation(
        images={"top": camera_rgb},          # (480, 640, 3) uint8 HWC RGB
        state=joint_positions,               # shape (14,) float
    )
    action = policy.step(obs)                # numpy action, shape (action_dim,)
    robot.apply(action)

Pi0.5 example (multi-camera + state + instruction)

Pi0.5 is a vision-language-action policy with two cameras and a task prompt:

import numpy as np
import sparkrt
from sparkrt import Observation, RuntimeConfig

policy = sparkrt.load_policy(
    "/path/to/pi05/checkpoints/020000/pretrained_model",
    config=RuntimeConfig.from_preset("latency"),
)
# policy.camera_names == ("image", "image2"); policy.state_dim == 8; requires_instruction == True

policy.reset()
for _ in range(horizon):
    obs = Observation(
        images={
            "image": base_rgb,    # (256, 256, 3) uint8 HWC RGB
            "image2": wrist_rgb,  # (256, 256, 3) uint8 HWC RGB
        },
        state=robot_state,        # shape (8,) float
        instruction="pick up the object and place it in the basket",
    )
    action = policy.step(obs)
    robot.apply(action)

For advanced use (parity checks, raw chunk prediction, custom loops) the underlying session is available as policy.session.

Runtime Presets

Built-in presets live under sparkrt/config/presets. They are intended as named deployment profiles rather than one-off benchmark commands. Calling sparkrt.load_policy(..., config=None) uses the eager backend; RuntimeConfig() has the same dependency-light eager default. Presets are explicit opt-in profiles and require sparkrt[presets] / omegaconf.

Preset	Purpose
default / safe	Conservative CUDA Graph path with the checkpoint's original Pi0.5 schedule. Use when bit-exact behavior versus eager SparkMind2 is required.
quality	Conservative TorchCompile path with the full Pi0.5 denoise schedule. Use when you want compile acceleration while avoiding reduced-step solver risk.
latency	Aggressive TorchCompile path with fewer Pi0.5 denoise steps. Use when lowest single-stream latency is more important than strict task-level stability.
memory	Development profile for the INT8-artifact memory path. Artifact loading is currently handled by the quantization scripts, not the public load_policy API.

Override individual fields when needed:

from sparkrt.config import RuntimeConfig

cfg = RuntimeConfig.from_preset("quality", kind="torchcompile")
cfg = RuntimeConfig.from_preset("latency", num_steps=8)

Programmatic configuration is also available for deployments that should not load YAML:

from sparkrt.config import RuntimeConfig, BackendConfig, Pi05RuntimeConfig

cfg = RuntimeConfig(
    backend=BackendConfig(kind="torchcompile"),
    pi05=Pi05RuntimeConfig(num_steps=10, schedule="uniform"),
    device="cuda:0",
)

Backend Integration

Recommended service shape:

1. Load one SparkRT Policy per worker process and robot/environment stream.
2. Warm the policy during worker startup, especially for TorchCompile.
3. Call policy.reset() when the task changes or a new episode starts.
4. Call policy.step(obs) from the control loop.
5. Do not share a single stateful policy across concurrent tasks.

For SparkMind2 eval flows that already run the policy pre/post-processors, use sparkrt.eval.PreprocessedRuntimePolicy rather than passing a Policy / InferenceSession directly. This avoids applying normalization and action post-processing twice.

See docs/INTEGRATION.md for the full integration plan, lifecycle recommendations, concurrency model, and failure handling.

Architecture

Observation (friendly camera names)
  -> Policy (SDK surface: step / reset / discovery)
  -> InferenceSession (stateful action queue / reset semantics)
  -> SparkMind processor (normalize / tokenize)
  -> model adapter (model orchestration, including Pi0.5 denoising)
  -> backend-compiled Regions (eager / CUDA Graph / TorchCompile)
  -> SparkMind processor (action postprocess)
  -> action

Layer	Package	Responsibility
Processors	sparkrt.processors	Convert observations to model-ready batches and convert model actions back to environment actions.
Adapters	sparkrt.adapters	Describe model computation as named fixed-shape Regions and own model-specific orchestration such as the Pi0.5 denoise loop.
Backends	sparkrt.backends	Execute regions with eager PyTorch, CUDA Graph replay, or TorchCompile.
Session	sparkrt.session	Own runtime state, action-chunk queues, reset semantics, and the model-agnostic action loop.
Config	sparkrt.config	Provide frozen dataclass configuration plus optional YAML presets.

The core contracts in sparkrt.core are intentionally lightweight and import without SparkMind2 or torch.

Development

Run dependency-light tests:

python -m pytest tests/ -q

Run the main Pi0.5 runtime benchmark:

CUDA_VISIBLE_DEVICES=0 MUJOCO_GL=egl PYOPENGL_PLATFORM=egl \
PYTHONPATH=/path/to/SparkRT:/path/to/SparkMind2 \
    python scripts/bench/bench_pi05_runtime.py \
    --checkpoint /path/to/pi05/pretrained_model \
    --backend torchcompile --num-steps 10 --schedule uniform

Run the development LIBERO success-rate runner:

CUDA_VISIBLE_DEVICES=0 MUJOCO_GL=egl PYOPENGL_PLATFORM=egl \
PYTHONPATH=/path/to/SparkMind2:/path/to/SparkRT \
    python scripts/eval/run_libero_eval_sparkrt.py \
    --checkpoint /path/to/pretrained_model \
    --preset quality \
    --n-episodes 5 \
    --seed 1000 \
    --output-dir runs/my_eval

TorchCompile has a large cold-start cost. Use a persistent Inductor cache for long-running workers:

export TORCHINDUCTOR_CACHE_DIR=$HOME/.cache/sparkrt/inductor

Documentation

• docs/INTEGRATION.md: service/backend integration plan.
• docs/EXPERIMENTS.md: benchmark and accuracy results.
• scripts/README.md: developer script reference.

Repository Layout

sparkrt/
  api.py          public load_policy / from_sparkmind_agent entry points
  policy.py       Policy SDK facade (step / reset / discovery)
  observation.py  Observation input type + camera-name conversion
  config/         RuntimeConfig dataclasses and YAML presets
  core/           backend, adapter, region, and shape contracts
  processors/     SparkMind processor reuse layer
  adapters/       ACT and Pi0.5 model adapters
  backends/       eager, cudagraph, and torchcompile executors
  session/        stateful inference session
  io/             SparkMind2 checkpoint loading helpers
  eval/           SparkMind2 eval-compatible policy wrapper

scripts/
  bench/          latency, profiling, and parity scripts
  eval/           LIBERO dev eval runner
  quant/          INT8 artifact export and benchmark scripts

docs/
  EXPERIMENTS.md
  INTEGRATION.md