locoformer

LocoFormer (wip)

LocoFormer - Generalist Locomotion via Long-Context Adaptation

The gist is they trained a simple Transformer-XL in simulation on robots with many different bodies (cross-embodiment) and extreme domain randomization. When transferring to the real-world, they noticed the robot now gains the ability to adapt to insults. The XL memories span across multiple trials, which allowed the robot to learn in-context adaptation.

Install

$ pip install locoformer

Usage

import torch
from locoformer.locoformer import Locoformer

# mock robot embodied with some state dimensions and action dimensions

locoformer = Locoformer(
    embedder = dict(
        dim = 512,
        dim_state = [32, 16], # support multiple bodies / robots
    ),
    unembedder = dict(
        num_continuous = 12 + 6,
        selectors = [
            list(range(12)),
            list(range(12, 12 + 6))
        ]
    ),
    transformer = dict(
        dim = 512,
        depth = 6,
        heads = 8,
        window_size = 32
    )
)

# mock state from one of the robots (0th one)

state = torch.randn(1, 1, 32)

# forward to get action logits

action_logits, _ = locoformer(
    state,
    state_embed_kwargs = dict(state_type = 'raw'),
    state_id_kwarg = dict(state_id = 0),
    action_select_kwargs = dict(selector_index = 0)
)

# sample action using the internal distribution

action = locoformer.unembedder.sample(action_logits, selector_index = 0) # (1, 1, 12)

Sim-to-Real Bipedal Locomotion (Microduck)

You can train an adaptive bipedal walking gait for the Microduck biped robot from scratch using Locoformer and PPO

An emergent bipedal walking gait starts emerging before 2M steps

$ uv run train_microduck.py

To render periodic evaluation videos or resume from the saved best checkpoint:

# render periodic evaluation videos to videos/
$ uv run train_microduck.py --eval_video True --eval_every 20

# resume training from best checkpoint
$ uv run train_microduck.py --resume True

Special thanks to Pollen Robotics 🤗 for their open sourcing of this sim2real env!

Citations

@article{liu2025locoformer,
    title   = {LocoFormer: Generalist Locomotion via Long-Context Adaptation},
    author  = {Liu, Min and Pathak, Deepak and Agarwal, Ananye},
    journal = {Conference on Robot Learning ({CoRL})},
    year    = {2025}
}

@inproceedings{anonymous2025flow,
    title   = {Flow Policy Gradients for Legged Robots},
    author  = {Anonymous},
    booktitle = {Submitted to The Fourteenth International Conference on Learning Representations},
    year    = {2025},
    url     = {https://openreview.net/forum?id=BA6n0nmagi},
    note    = {under review}
}

@misc{ashlag2025stateentropyregularizationrobust,
    title   = {State Entropy Regularization for Robust Reinforcement Learning},
    author  = {Yonatan Ashlag and Uri Koren and Mirco Mutti and Esther Derman and Pierre-Luc Bacon and Shie Mannor},
    year    = {2025},
    eprint  = {2506.07085},
    archivePrefix = {arXiv},
    primaryClass = {cs.LG},
    url     = {https://arxiv.org/abs/2506.07085},
}

@misc{osband2026delightfulpolicygradient,
    title   = {Delightful Policy Gradient},
    author  = {Ian Osband},
    year    = {2026},
    eprint  = {2603.14608},
    archivePrefix = {arXiv},
    primaryClass = {cs.LG},
    url     = {https://arxiv.org/abs/2603.14608},
}

@misc{schwarzer2021dataefficientreinforcementlearningselfpredictive,
    title   = {Data-Efficient Reinforcement Learning with Self-Predictive Representations},
    author  = {Max Schwarzer and Ankesh Anand and Rishab Goel and R Devon Hjelm and Aaron Courville and Philip Bachman},
    year    = {2021},
    eprint  = {2007.05929},
    archivePrefix = {arXiv},
    primaryClass = {cs.LG},
    url     = {https://arxiv.org/abs/2007.05929},
}

@misc{teoh2026nextlatentpredictiontransformerslearn,
    title   = {Next-Latent Prediction Transformers Learn Compact World Models},
    author  = {Jayden Teoh and Manan Tomar and Kwangjun Ahn and Edward S. Hu and Tim Pearce and Pratyusha Sharma and Akshay Krishnamurthy and Riashat Islam and Alex Lamb and John Langford},
    year    = {2026},
    eprint  = {2511.05963},
    archivePrefix = {arXiv},
    primaryClass = {cs.LG},
    url     = {https://arxiv.org/abs/2511.05963},
}

Metadata

Release files for locoformer 0.3.2

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for locoformer 0.3.2
File	Size	Uploaded
locoformer-0.3.2.tar.gz	28.0 kB	Sep 11, 2026	Details

Built distribution (wheel)

Table of built distributions (wheels) for locoformer 0.3.2
File	Interpreter	ABI	Platform	Reset
locoformer-0.3.2-py3-none-any.whl 27.1 kB Sep 11, 2026	Python 3	none	any	Details

Total release size: 55.0 kB

Release files / locoformer-0.3.2.tar.gz

Download URL	locoformer-0.3.2.tar.gz
Size	28.0 kB
Tags	Source
SHA-256 checksum How to use checksums	`3ad2ff8fa6c97dfd21a76ba96244b9b5f7e941c18843b69f9f1544b34af76dd3`
BLAKE2b-256 checksum How to use checksums	`a345bdab28e42bf101e31173aaaebad523ab534409c128c3039da8bfaee416f0`
Upload date	Sep 11, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`uv/0.8.17`

Release files / locoformer-0.3.2-py3-none-any.whl

Download URL	locoformer-0.3.2-py3-none-any.whl
Size	27.1 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`e5b712b3b49ef70c7da149598b951c80d403c72f0dddc7e622fa3a42b29d64e3`
BLAKE2b-256 checksum How to use checksums	`2432cf81f4a3cf9e187d142a1bce537c3c57357928505a714ac9e9636c262055`
Upload date	Sep 11, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`uv/0.8.17`