Model-based planning using GPU-accelerated Simulator as a World Model

This is an implementation of model-based (MB) planner that utilizes MuJoCo MJX physics simulator as a world model. This repository contains the MB-planner that is capable of generating a smooth, constrained and collision free path. In addition, the framework supports flexible task objectives by adjusting the cost function, enabling goal-driven behavior under varying conditions.

Offline Planning The entire trajectory is planned offline and executed afterwards.	Online Planning: Collision Avoidance The robot reaches target position and orientation while avoiding obstacles (red).	Online Planning: Contact Task The robot pushes the object (green sphere) to the target location (white square).
Online Planning: Real-Life The MB-planner is used to control real UR5e in dynamic environment. Online Planning: Real-Life The MB-planner is used to control real UR5e in dynamic environment.

Repository Structure

• data/: Folder to which the cost data and trajectory are saved.
• ur5e_hande_mjx/: The MuJoCo model description.
• trajectory_sampler.py : The code for trajectory sampling from multivariate normal distribution and application of constraints.
• cem_optimization.py : The code for trajectory optimization with Cross-Entropy Method.
• mpc_planner.py : The code for online planning using Model Predictive Control strategy.
• visualizer.py : The code for visualizing current MuJoCo environment and/or recorded trajectory from /data folder.

The code for the demos displayed above is located on separate branches of this repository:

Demo	Branch
Offline Planning	offline_demo
Online Planning: Collision Avoidance	online_demo_1
Online Planning: Contact Task	online_demo_2
Online Planning: Real-Life	real_life_demo

Installation

NOTE: NVIDIA graphics card is required.

$ git clone https://github.com/patsyuk03/mjx_planner.git
$ cd mjx_planner
$ pip install -r requirements.txt

Run the Example

$ python3 mpc_planner.py