search

google-deepmind / mujoco_mpc

Real-time behaviour synthesis with MuJoCo, using Predictive Control
https://github.com/deepmind/mujoco_mpc
Apache License 2.0
952 stars 142 forks source link

MuJoCo MPC running on GPU (MuJoCo XLA) #221

Closed hessad closed 8 months ago

hessad commented 9 months ago

Has anyone tried to use MJPC (e.g. predictive sampling) with MuJoCo XLA to take advantage of GPU parallelization? Is it being added in future MJPC versions or could it already be tested with the current MJPC version and some modification?

thowell commented 8 months ago

I haven't seen an implementation of predictive sampling written in JAX, but this is definitely something of interest.

I recommend checking out trajax which includes a number of planners written in JAX. This tool should be compatible with MuJoCo XLA.