SparkMind2

SparkMind2 is a robot imitation-learning and vision-language-action training framework built around Hydra, PyTorch, Accelerate, and LeRobot-compatible artifacts.

The repository focuses on clean training paths for open robot datasets and policies:

• IL policies: ACT and Diffusion Policy
• VLA policies: SmolVLA, Pi0, and Pi0.5
• Datasets and environments: ALOHA simulation, PushT, and LIBERO
• Checkpoints: Accelerate training checkpoints for resume, plus portable pretrained_model/ artifacts for evaluation, export, and fine-tuning

SparkMind2 keeps training behavior config-driven. The example scripts under examples/ are thin entrypoints; model, dataset, optimizer, distributed, and checkpoint behavior live in the framework and Hydra configs.

Status

SparkMind2 is under active development. The current public surface is intended for research training, checkpointing, and LeRobot-compatible policy artifact workflows. Hardware deployment and unrelated legacy RL stacks are not part of this repository's main path.

Supported Policies

Area	Policy	Example	Base config
Imitation learning	ACT	examples/learning_il/01_demo_ACT.py	BaseTaskACT
Imitation learning	Diffusion Policy	examples/learning_il/02_demo_DP.py	BaseTaskDP
VLA	SmolVLA	examples/learning_vla/01_demo_SmolVLA.py	BaseTaskSmolVLA
VLA	Pi0	examples/learning_vla/02_demo_PI0.py	BaseTaskPi0
VLA	Pi0.5	examples/learning_vla/03_demo_PI05.py	BaseTaskPi05

Installation

SparkMind2 requires Python 3.12.

git clone <repo-url> SparkMind2
cd SparkMind2

python -m venv .venv
source .venv/bin/activate
pip install -U pip
pip install -e .

Optional extras:

# Pi-family tokenizers and processor dependencies
pip install -e ".[pi]"

# Simulation environments
pip install -e ".[aloha]"
pip install -e ".[pusht]"
pip install -e ".[libero]"

# Common training image dependency set
pip install -e ".[aws]"

# Development tools
pip install -e ".[dev]"

LIBERO installs native simulation dependencies on Linux. If hf-libero needs to build native packages, install a real CMake 3.x binary first, for example:

conda install -c conda-forge "cmake>=3.18,<4"

For headless LIBERO evaluation or training, set:

export MUJOCO_GL=egl

Quick Start

All examples accept regular Hydra overrides after the script arguments.

ACT

cd examples/learning_il

python 01_demo_ACT.py \
  --task Aloha_sim/Aloha_sim_insertion_human \
  --pretrained-path lerobot/act_aloha_sim_insertion_human \
  --inference

Train:

python 01_demo_ACT.py \
  --task Aloha_sim/Aloha_sim_insertion_human \
  train.Trainer.maximum_steps=10000 \
  train.Trainer.batch_size=8

Diffusion Policy

cd examples/learning_il

python 02_demo_DP.py \
  --task PushT/PushT \
  --pretrained-path lerobot/diffusion_pusht \
  --inference

SmolVLA

cd examples/learning_vla

python 01_demo_SmolVLA.py \
  --task Libero/Libero_object \
  --pretrained-path HuggingFaceVLA/smolvla_libero \
  --inference

Pi0

cd examples/learning_vla

python 02_demo_PI0.py \
  --task Libero/Libero_object \
  --pretrained-path lerobot/pi0_libero_finetuned_v044 \
  --inference

Pi0.5

cd examples/learning_vla

python 03_demo_PI05.py \
  --task Libero/Libero_object \
  --pretrained-path lerobot/pi05_libero_finetuned_v044 \
  --inference

Training With Hydra Overrides

The main config groups are:

• sparkmind/configs/config_il.yaml
• sparkmind/configs/config_vla.yaml
• sparkmind/configs/task/
• sparkmind/configs/train/

Examples:

python examples/learning_vla/03_demo_PI05.py \
  --task Libero/Libero_object \
  train.Trainer.maximum_steps=200000 \
  train.Trainer.batch_size=8 \
  train.Trainer.num_workers=4 \
  train.Trainer.output_dir=./runs/pi05_libero_object

Fine-tune from a portable policy artifact:

python examples/learning_vla/03_demo_PI05.py \
  --task Libero/Libero_object \
  --pretrained-path lerobot/pi05_libero_finetuned_v044

Resume from a SparkMind2 training checkpoint:

python examples/learning_vla/03_demo_PI05.py \
  --task Libero/Libero_object \
  train.Trainer.resume=true \
  train.Trainer.resume_path=./runs/pi05_libero_object/checkpoints/last

Distributed Training

SparkMind2 uses Accelerate for distributed execution. Multi-GPU runs can be launched with:

accelerate launch --num_processes 8 --multi_gpu \
  examples/learning_vla/03_demo_PI05.py \
  --task Libero/Libero_object \
  train.Trainer.batch_size=8

For FSDP, select the distributed backend through Hydra:

accelerate launch --num_processes 8 --multi_gpu \
  examples/learning_vla/03_demo_PI05.py \
  --task Libero/Libero_object \
  train.Trainer.distributed_backend=fsdp

The framework keeps checkpoint state under Accelerate's save/load path. Users normally only need to choose the backend and checkpoint cadence.

Checkpoints And Artifacts

SparkMind2 writes two kinds of outputs:

checkpoints/{step}/
  training_state/accelerate_state/
  pretrained_model/
checkpoints/last

• training_state/accelerate_state/ is the canonical resume checkpoint. It contains model, optimizer, scheduler, RNG, distributed state, and SparkMind training step.
• pretrained_model/ is the portable policy artifact. It contains model weights, config, processor files, and train config for evaluation, export, and future fine-tuning.
• checkpoints/last points to the latest checkpoint.

For FSDP, training checkpoints use Accelerate/PyTorch distributed checkpointing and sharded state dicts by default. Portable pretrained_model/ artifacts remain LeRobot-compatible.

Pi Attention Backend

Pi0 and Pi0.5 expose:

train.Model.attention_backend: eager  # eager, sdpa

eager is the default because it matches the reference Gemma attention path and is the safest baseline. sdpa is available as an override for memory-sensitive experiments:

python examples/learning_vla/03_demo_PI05.py \
  --task Libero/Libero_object \
  train.Model.attention_backend=sdpa

The Pi prefix-LM/block attention mask is not a pure causal mask, so PyTorch's built-in FlashAttention backend is not used by default.

Development

Run syntax checks on touched files:

python -m py_compile \
  sparkmind/learning/VLA/models/pi0_model.py \
  sparkmind/learning/VLA/models/pi05_model.py

git diff --check

Install development dependencies with:

pip install -e ".[dev]"

Project Layout

examples/                 Thin training and inference entrypoints
sparkmind/configs/         Hydra config groups
sparkmind/data/            Dataset, processor, and optimizer utilities
sparkmind/learning/IL/     ACT and Diffusion Policy trainers and agents
sparkmind/learning/VLA/    SmolVLA, Pi0, and Pi0.5 trainers, agents, and models
docs/                      Additional command references
docker/                    Container and deployment helpers

License

Add the project license before publishing this repository publicly.