Package: pytorch-mppi
PyTorch MPPI ImplementationThis repository implements Model Predictive Path Integral (MPPI)
with approximate dynamics in pytorch. MPPI typically requires actual
trajectory samples, butthis papershowed that it could be done with approximate dynamics (such as with a neural network)
using importance sampling.Thus it can be used in place of other trajectory optimization methods
such as the Cross Entropy Method (CEM), or random shooting.Installationpipinstallpytorch-mppifor autotuning hyperparameters, install withpipinstallpytorch-mppi[tune]for running tests, install withpipinstallpytorch-mppi[test]for development, clone the repository then install in editable modepipinstall-e.UsageSeetests/pendulum_approximate.pyfor usage with a neural network approximating
the pendulum dynamics. See thenot_batchbranch for an easier to read
algorithm. Basic use case is shown belowfrompytorch_mppiimportMPPI# create controller with chosen parametersctrl=MPPI(dynamics,running_cost,nx,noise_sigma,num_samples=N_SAMPLES,horizon=TIMESTEPS,lambda_=lambda_,device=d,u_min=torch.tensor(ACTION_LOW,dtype=torch.double,device=d),u_max=torch.tensor(ACTION_HIGH,dtype=torch.double,device=d))# assuming you have a gym-like envobs=env.reset()foriinrange(100):action=ctrl.command(obs)obs,reward,done,_=env.step(action.cpu().numpy())Requirementspytorch (>= 1.0)next state <- dynamics(state, action)function (doesn't have to be true dynamics)stateisK x nx,actionisK x nucost <- running_cost(state, action)functioncostisK x 1, state isK x nx,actionisK x nuFeaturesApproximate dynamics MPPI with importance samplingParallel/batch pytorch implementation for accelerated samplingControl bounds via sampling control noise from rectified gaussianHandle stochastic dynamic models (assuming each call is a sample) by sampling multiple state trajectories for the same
action trajectory withrollout_samplesParameter tuning and hintsterminal_state_cost- function(state (K x T x nx)) -> cost (K x 1) by default there is no terminal
cost, but if you experience your trajectory getting close to but never quite reaching the goal, then
having a terminal cost can help. The function should scale with the horizon (T) to keep up with the
scaling of the running cost.lambda_- higher values increases the cost of control noise, so you end up with more
samples around the mean; generally lower values work better (try1e-2)num_samples- number of trajectories to sample; generally the more the better.
Runtime performance scales much better withnum_samplesthanhorizon, especially
if you're using a GPU device (remember to pass that in!)noise_mu- the default is 0 for all control dimensions, which may work out
really poorly if you have control bounds and the allowed range is not 0-centered.
Remember to change this to an appropriate value for non-symmetric control dimensions.Autotunefrom version 0.5.0 onwards, you can automatically tune the hyperparameters.
A convenient tuner compatible with the popularray tunelibrary
is implemented. You can select from a variety of cutting edge black-box optimizers such asCMA-ES,HyperOpt,fmfn/BayesianOptimization, and so on.
Seetests/auto_tune_parameters.pyfor an example. A tutorial based on it follows.The tuner can be used for other controllers as well, but you will need to define the appropriateTunableParametersubclasses.First we create a toy 2D environment to do controls on and create the controller with some
default parameters.importtorchfrompytorch_mppiimportMPPIdevice="cpu"dtype=torch.double# create toy environment to do on control on (default start and goal)env=Toy2DEnvironment(visualize=True,terminal_scale=10)# create MPPI with some initial parametersmppi=MPPI(env.dynamics,env.running_cost,2,terminal_state_cost=env.terminal_cost,noise_sigma=torch.diag(torch.tensor([5.,5.],dtype=dtype,device=device)),num_samples=500,horizon=20,device=device,u_max=torch.tensor([2.,2.],dtype=dtype,device=device),lambda_=1)We then need to create an evaluation function for the tuner to tune on.
It should take no arguments and output aEvaluationResultpopulated at least by costs.
If you don't need rollouts for the cost evaluation, then you can set it to None in the return.
Tips for creating the evaluation function are described in comments below:frompytorch_mppiimportautotune# use the same nominal trajectory to start with for all the evaluations for fairnessnominal_trajectory=mppi.U.clone()# parameters for our sample evaluation function - lots of choices for the evaluation functionevaluate_running_cost=Truenum_refinement_steps=10num_trajectories=5defevaluate():costs=[]rollouts=[]# we sample multiple trajectories for the same start to goal problem, but in your case you should consider# evaluating over a diverse dataset of trajectoriesforjinrange(num_trajectories):mppi.U=nominal_trajectory.clone()# the nominal trajectory at the start will be different if the horizon's changedmppi.change_horizon(mppi.T)# usually MPPI will have its nominal trajectory warm-started from the previous iteration# for a fair test of tuning we will reset its nominal trajectory to the same random one each time# we manually warm it by refining it for some stepsforkinrange(num_refinement_steps):mppi.command(env.start,shift_nominal_trajectory=False)rollout=mppi.get_rollouts(env.start)this_cost=0rollout=rollout[0]# here we evaluate on the rollout MPPI cost of the resulting trajectories# alternative costs for tuning the parameters are possible, such as just considering terminal costifevaluate_running_cost:fortinrange(len(rollout)-1):this_cost=this_cost+env.running_cost(rollout[t],mppi.U[t])this_cost=this_cost+env.terminal_cost(rollout,mppi.U)rollouts.append(rollout)costs.append(this_cost)# can return None for rollouts if they do not need to be calculatedreturnautotune.EvaluationResult(torch.stack(costs),torch.stack(rollouts))With this we have enough to start tuning. For example, we can tune iteratively with the CMA-ES optimizer# these are subclass of TunableParameter (specifically MPPIParameter) that we want to tuneparams_to_tune=[autotune.SigmaParameter(mppi),autotune.HorizonParameter(mppi),autotune.LambdaParameter(mppi)]# create a tuner with a CMA-ES optimizertuner=autotune.Autotune(params_to_tune,evaluate_fn=evaluate,optimizer=autotune.CMAESOpt(sigma=1.0))# tune parameters for a number of iterationsiterations=30foriinrange(iterations):# results of this optimization step are returnedres=tuner.optimize_step()# we can render the rollouts in the environmentenv.draw_rollouts(res.rollouts)# get best results and apply it to the controller# (by default the controller will take on the latest tuned parameter, which may not be best)res=tuner.get_best_result()tuner.apply_parameters(res.param_values)This is a local search method that optimizes starting from the initially defined parameters.
For global searching, we use ray tune compatible searching algorithms. Note that you can modify the
search space of each parameter, but default reasonable ones are provided.# can also use a Ray Tune optimizer, see# https://docs.ray.io/en/latest/tune/api_docs/suggestion.html#search-algorithms-tune-search# rather than adapting the current parameters, these optimizers allow you to define a search space for each# and will search on that spacefrompytorch_mppiimportautotune_globalfromray.tune.search.hyperoptimportHyperOptSearchfromray.tune.search.bayesoptimportBayesOptSearch# the global version of the parameters define a reasonable search space for each parameterparams_to_tune=[autotune_global.SigmaGlobalParameter(mppi),autotune_global.HorizonGlobalParameter(mppi),autotune_global.LambdaGlobalParameter(mppi)]# be sure to close any figures before ray tune optimization or they will be duplicatedenv.visualize=Falseplt.close('all')tuner=autotune_global.AutotuneGlobal(params_to_tune,evaluate_fn=evaluate,optimizer=autotune_global.RayOptimizer(HyperOptSearch))# ray tuners cannot be tuned iteratively, but you can specify how many iterations to tune forres=tuner.optimize_all(100)res=tuner.get_best_result()tuner.apply_parameters(res.params)For example tuning hyperparameters (with CMA-ES) only on the toy problem (the nominal trajectory is reset each time so they are sampling from noise):If you want more than just the best solution found, such as if you want diversity
across hyperparameter values, or if your evaluation function has large uncertainty,
then you can directly query past results byforresintuner.optim.all_res:# the costprint(res.metrics['cost'])# extract the parametersparams=tuner.config_to_params(res.config)print(params)# apply the parameters to the controllertuner.apply_parameters(params)Alternatively you can try Quality Diversity optimization using theCMA-ME optimizer. This optimizer will
try to optimize for high quality parameters while ensuring there is diversity across
them. However, it is very slow and you might be better using aRayOptimizerand selecting
for top results while checking for diversity.
To use it, you need to installpipinstallribsYou then use it asimportpytorch_mppi.autotune_qdoptim=pytorch_mppi.autotune_qd.CMAMEOpt()tuner=autotune_global.AutotuneGlobal(params_to_tune,evaluate_fn=evaluate,optimizer=optim)iterations=10foriinrange(iterations):# results of this optimization step are returnedres=tuner.optimize_step()# we can render the rollouts in the environmentbest_params=optim.get_diverse_top_parameters(5)forresinbest_params:print(res)TestsYou'll need to installgym<=0.20to run the tests (for thePendulum-v0environment).
The easy way to install this and other testing dependencies is runningpython setup.py test.
Note thatgympast0.20deprecatedPendulum-v0forPendulum-v1with incompatible dynamics.Undertestsyou can find theMPPImethod applied to known pendulum dynamics
and approximate pendulum dynamics (with a 2 layer feedforward net
estimating the state residual). Using a continuous angle representation
(feedingcos(\theta), sin(\theta)instead of\thetadirectly) makes
a huge difference. Although both works, the continuous representation
is much more robust to controller parameters and random seed. In addition,
the problem of continuing to spin after over-swinging does not appear.Sample result on approximate dynamics with 100 steps of random policy data
to initialize the dynamics:Related projectspytorch CEM- an alternative MPC shooting method with similar API as this
project