Metadata-Version: 2.4
Name: sparkmind
Version: 2.0.0rc4
Summary: Unified Robot Skill Learning Framework
Requires-Python: >=3.12
Description-Content-Type: text/markdown
Requires-Dist: accelerate<2.0.0,>=1.10.0
Requires-Dist: av<16.0.0,>=15.0.0
Requires-Dist: datasets<5.0.0,>=4.0.0
Requires-Dist: diffusers<0.36.0,>=0.27.2
Requires-Dist: draccus==0.10.0
Requires-Dist: einops<0.9.0,>=0.8.0
Requires-Dist: fsspec>=2024.0.0
Requires-Dist: gymnasium<2.0.0,>=1.1.1
Requires-Dist: huggingface-hub<2.0.0,>=1.0.0
Requires-Dist: hydra-core>=1.3.0
Requires-Dist: imageio[ffmpeg]<3.0.0,>=2.34.0
Requires-Dist: numpy<2.3.0,>=2.0.0
Requires-Dist: opencv-python-headless<4.14.0,>=4.9.0
Requires-Dist: omegaconf>=2.3.0
Requires-Dist: packaging<26.0,>=24.2
Requires-Dist: pandas<4.0.0,>=2.2.0
Requires-Dist: Pillow<13.0.0,>=10.0.0
Requires-Dist: pyarrow>=15.0.0
Requires-Dist: PyYAML>=6.0.0
Requires-Dist: safetensors<1.0.0,>=0.4.3
Requires-Dist: scipy<2.0.0,>=1.14.0
Requires-Dist: shutup>=0.2.0
Requires-Dist: tensorboard>=2.0.0
Requires-Dist: termcolor<4.0.0,>=2.4.0
Requires-Dist: torch<2.11.0,>=2.7.0
Requires-Dist: torchvision<0.26.0,>=0.22.0
Requires-Dist: transformers<6.0.0,>=5.4.0
Requires-Dist: tqdm>=4.66.0
Requires-Dist: typing-extensions<5.0.0,>=4.12.0
Requires-Dist: wandb<0.25.0,>=0.24.0
Provides-Extra: aloha
Requires-Dist: gym-aloha<0.2.0,>=0.1.2; extra == "aloha"
Requires-Dist: scipy<2.0.0,>=1.14.0; extra == "aloha"
Provides-Extra: pusht
Requires-Dist: gym-pusht<0.2.0,>=0.1.5; extra == "pusht"
Requires-Dist: pymunk<7.0.0,>=6.6.0; extra == "pusht"
Provides-Extra: libero
Requires-Dist: hf-libero<0.2.0,>=0.1.3; sys_platform == "linux" and extra == "libero"
Requires-Dist: scipy<2.0.0,>=1.14.0; extra == "libero"
Provides-Extra: envs
Requires-Dist: gym-aloha<0.2.0,>=0.1.2; extra == "envs"
Requires-Dist: gym-pusht<0.2.0,>=0.1.5; extra == "envs"
Requires-Dist: pymunk<7.0.0,>=6.6.0; extra == "envs"
Requires-Dist: hf-libero<0.2.0,>=0.1.3; sys_platform == "linux" and extra == "envs"
Provides-Extra: dev
Requires-Dist: pytest>=8.1.0; extra == "dev"
Requires-Dist: pre-commit<5.0.0,>=3.7.0; extra == "dev"
Requires-Dist: debugpy<1.9.0,>=1.8.1; extra == "dev"
Requires-Dist: mypy>=1.19.1; extra == "dev"
Provides-Extra: pi
Requires-Dist: transformers<6.0.0,>=5.4.0; extra == "pi"
Requires-Dist: scipy<2.0.0,>=1.14.0; extra == "pi"
Requires-Dist: sentencepiece>=0.1.99; extra == "pi"
Requires-Dist: tiktoken>=0.6.0; extra == "pi"
Provides-Extra: aws
Requires-Dist: gym-aloha<0.2.0,>=0.1.2; extra == "aws"
Requires-Dist: hf-libero<0.2.0,>=0.1.3; sys_platform == "linux" and extra == "aws"
Requires-Dist: scipy<2.0.0,>=1.14.0; extra == "aws"
Requires-Dist: sentencepiece>=0.1.99; extra == "aws"
Requires-Dist: tiktoken>=0.6.0; extra == "aws"

# 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.

```bash
git clone <repo-url> SparkMind2
cd SparkMind2

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

Optional extras:

```bash
# 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:

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

For headless LIBERO evaluation or training, set:

```bash
export MUJOCO_GL=egl
```

## Quick Start

All examples accept regular Hydra overrides after the script arguments.

### ACT

```bash
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:

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

### Diffusion Policy

```bash
cd examples/learning_il

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

### SmolVLA

```bash
cd examples/learning_vla

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

### Pi0

```bash
cd examples/learning_vla

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

### Pi0.5

```bash
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:

```bash
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:

```bash
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:

```bash
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:

```bash
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:

```bash
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:

```text
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:

```yaml
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:

```bash
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:

```bash
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:

```bash
pip install -e ".[dev]"
```

## Project Layout

```text
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.
