This is the PyTorch implementation for training and deployment of X-Mobility.
X-Mobility: A generalizable navigation model using auto-regressive world modeling, multi-head decoders, and decoupled policy learning for robust, zero-shot Sim2Real and cross-embodiment transfer.
-
Install docker
-
Build the docker image
docker build --network=host -t <image-name> .
- Download the datasets from Hugging Face: https://huggingface.co/datasets/nvidia/X-Mobility
- x_mobility_isaac_sim_nav2_100k.zip: Teacher policy dataset to train world model and action network together.
- x_mobility_isaac_sim_random_160k.zip: Random action dataset to pre-train world model without action network.
- [Optional] Download the pre-trained checkpoints from Hugging Face: https://huggingface.co/nvidia/X-Mobility
- x_mobility-nav2-semantic_action_path.ckpt: Trained with teacher policy dataset with semantic segmenetation, action and path decoding enabled.
- Launch the docker image:
docker run --shm-size=512g -v <path-to-datasets>:/workspace/datasets -it <image-name> bash
- Set WANDB_API_KEY inside container:
export WANDB_API_KEY=<wandb-api-key>
- Training with action policy enabled:
python3 train.py -c configs/train_config.gin -d datasets/ -o <output-dir> -e <wandb-entity> -n <wandb-project> -r <wandb-run>
- Training with world model only:
python3 train.py -c configs/pretrained_gwm_train_config.gin -d datasets/ -o <output-dir> -e <wandb-entity> -n <wandb-project> -r <wandb-run>
- Evaluating checkpoint:
python3 evaluate.py -c configs/train_config.gin -d datasets/ -p <checkpoint> -e <wandb-entity> -n <wandb-project> -r <wandb-run>
- ONNX & TensorRT conversion:
python3 onnx_conversion.py -p <checkpoint> -o <onnx_file_path>
python3 trt_conversion.py -o <onnx_file_path> -t <trt_file_path>
Note: TensorRT engines are specific to both the TensorRT version and the GPU on which they are created. Therefore, it's recommended to rebuild the engine outside of docker on the target platform to run inference.
- E2E navigation with ROS2:
With the TensorRT engine, follow this instruction to setup ROS2, Isaac Sim to run a demo example of X-Mobility.
X-Mobility is released under the Apache License 2.0. See LICENSE for additional details.
Wei Liu, Huihua Zhao, Chenran Li, Joydeep Biswas, Billy Okal, Pulkit Goyal, Soha Pouya, Yan Chang
Stay tuned for the upcoming RL enhancements for embodiments using Isaac Lab.
We would like to acknowledge the following projects where parts of the codes in this repo is derived from:
If you find this repository useful, please consider citing:
@article{liu2024x,
title={X-mobility: End-to-end generalizable navigation via world modeling},
author={Liu, Wei and Zhao, Huihua and Li, Chenran and Biswas, Joydeep and Okal, Billy and Goyal, Pulkit and Chang, Yan and Pouya, Soha},
journal={arXiv preprint arXiv:2410.17491},
year={2024}
}