Skip to main content

SparkRT - edge-side low-latency inference runtime for SparkMind2 robot/VLA/world models

Project description

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

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

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

sparkrt-0.1.0rc1.tar.gz (50.4 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

sparkrt-0.1.0rc1-py3-none-any.whl (59.3 kB view details)

Uploaded Python 3

File details

Details for the file sparkrt-0.1.0rc1.tar.gz.

File metadata

  • Download URL: sparkrt-0.1.0rc1.tar.gz
  • Upload date:
  • Size: 50.4 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.12.13

File hashes

Hashes for sparkrt-0.1.0rc1.tar.gz
Algorithm Hash digest
SHA256 6aa1018cbb9b465659f87ca4280c6dcec32ad2bf72543fdf1c16286911c09bba
MD5 8d445a9a552538576fcc137e304acb40
BLAKE2b-256 6a14f6cd01223a0a08129978f9cbd6ed5ff210d67e68b06946cce2bc5592307c

See more details on using hashes here.

File details

Details for the file sparkrt-0.1.0rc1-py3-none-any.whl.

File metadata

  • Download URL: sparkrt-0.1.0rc1-py3-none-any.whl
  • Upload date:
  • Size: 59.3 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.12.13

File hashes

Hashes for sparkrt-0.1.0rc1-py3-none-any.whl
Algorithm Hash digest
SHA256 9c86502b06dfbbbe32e77fed83546a854d96bb2e5dd6d72b9d70b88ddeaf74ab
MD5 444213e51631fc4d81e1b6f7c2217466
BLAKE2b-256 9fb528b6a7e299dde49cc8cf3b5c0efdbff5c96cc2a911f1c9c0b468ab6cbecb

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page