Skip to content

ahadjawaid/nac

Repository files navigation

NAC

NAC: Neural Action Codec for Vision-Language-Action Models

Ahad Jawaid and Yu Xiang

Website Apache License Twitter/X

NAC overview

NAC (Neural Action Codec) compresses continuous robot actions into a discrete, multi-scale token vocabulary using a quantized autoencoder with adversarial training. At inference, a vision-language-action policy predicts action tokens that the NAC decoder turns back into continuous control commands.

News

  • 2026-06-28 — Initial public release.

Setup

Clone with submodules so third_party/LIBERO is available:

git clone --recurse-submodules [email protected]:ahadjawaid/nac.git
cd nac

Install with uv (Python 3.10, pinned via .python-version):

uv sync
uv pip install -e .

If uv doesn't work on your cluster, micromamba env create -f conda_env.yaml is an alternative.

Usage

Training runs through scripts/run_workspace.py with a namespaced config: --config-name=model/<env>/train_<name>, where <env> is libero or robomimic.

Datasets

LIBERO — download the prebuilt libero10 zarr from Hugging Face and extract it to data/libero/, or build it locally:

uv run third_party/LIBERO/benchmark_scripts/download_libero_datasets.py --datasets libero_100
uv run scripts/convert_libero_dataset.py --root_dir data/libero --hdf5_dir_name hdf5_datasets
uv run scripts/compose_libero_multitask_dataset.py --multitask_name libero10 --root_dir data/libero

RoboMimic — download the prebuilt mt4 multitask zarr from Hugging Face and extract it to data/robomimic/ so data/robomimic/mt4_N800.zarr exists, or build it locally:

uv run python -m robomimic.scripts.download_datasets \
    --download_dir data/robomimic/hdf5_datasets \
    --tasks lift can square tool_hang --dataset_types ph --hdf5_types raw
uv run scripts/convert_robomimic_dataset.py --root_dir data/robomimic --hdf5_dir_name hdf5_datasets -ds ph
uv run scripts/compose_robomimic_multitask_dataset.py -mt mt4 --root_dir data/robomimic -ds ph

ToolHang requires robosuite>=1.5 for local conversion (uv pip install --force-reinstall git+https://github.com/Chaoqi-LIU/robosuite.git@trajtok).

This writes data/<env>/<task>_N<num_demo>.zarr. Set training.num_demo on the command line to match the N in the filename (LIBERO defaults to 500, RoboMimic to 800).

Train the NAC tokenizer

uv run accelerate launch scripts/run_workspace.py \
    --config-name=model/libero/train_nactok \
    task/tokenizer=libero/libero10

Train the policy

MUJOCO_GL=egl uv run accelerate launch scripts/run_workspace.py \
    --config-name=model/libero/train_nacpolicy \
    task/policy=libero/libero10 \
    policy.action_tokenizer.checkpoint=[path/to/nactok.ckpt]

For RoboMimic, use --config-name=model/robomimic/... with task/tokenizer=robomimic/mt4 / task/policy=robomimic/mt4. Baselines follow the same pattern: tokenizers train_{bintok,fasttok,oattok,vqvlatok} with their matching policies train_{bin,fast,oat,vqvla}policy, plus train_diffpolicy (no tokenizer needed).

Evaluate

MUJOCO_GL=egl uv run scripts/eval_policy_sim.py \
    --checkpoint [path/to/policy.ckpt] \
    --output_dir output/eval/libero10 \
    --num_exp 5

--checkpoint accepts a single .ckpt or a directory of checkpoints; results are written to eval_log.json.

Training on LeRobot data (optional)

scripts/convert_lerobot_dataset.py converts LeRobot datasets to the repo's zarr format. lerobot is not in the core dependencies (it conflicts with transformers>=5.8), so run the conversion in a throwaway venv:

uv venv .venv-lerobot
uv pip install -p .venv-lerobot lerobot==0.4.4 "zarr<3" "numcodecs<0.16" click
.venv-lerobot/bin/python scripts/convert_lerobot_dataset.py \
    --out data/lerobot/demos.zarr --repo_id lerobot/pusht

Then train against it by overriding task.tokenizer.dataset.zarr_path / task.policy.dataset.zarr_path.

Acknowledgement

Built on top of OAT and Sim-Env, with VQ-VLA as a baseline in this repo.

About

No description, website, or topics provided.

Resources

License

Stars

14 stars

Watchers

1 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors

Languages