Trajectory Optimization API

[alberthli]

This PR adds a very generic and flexible API for trajectory optimization. As a specific instantiation, it also implements the extremely simple predictive sampling algorithm.

• defines a generic API for trajectory optimization
• defines a generic API for cost functions (and their gradients) along with the ability to pass a generic set of caches/parameters to relevant functions
  • implements a simple quadratic cost function with fixed goal state
• implements predictive sampling

[vincekurtz]

Overall API looks good, just missing a few tests.

I would add tests for:

• Computing the cost for quadratic with static goal (vs for loop)
• Computing the gradient and Hessian for the quadratic with static goal (vs autodiff)
• A smoke test for creating a VanillaPredictiveSampler and calling optimize (maybe not necessary depending on below)
• Some sanity check that predictive sampling is doing the right thing. That could be a pendulum swingup example (should work even if the implementation isn't blazing fast yet) or a little unit test that shows the cost generally going down, or something else.

[alberthli]

Agreed on the first two requested tests - those should be trivial to spin up.

Sanity check of choice: randomly sample a batch of initial policies and also shoot them forward. pass those guesses to vanilla predictive sampling, which returns new trajectories that should be no worse than the guesses. verify this property in a test.

TODO list:

• refactor q and v to x everywhere
• test cost function and its derivatives (thanks for suggesting this, literally all my implementations were wrong)
• smoke test for VanillaPredictiveSampler + optimize + jit
• sanity check for VanillaPredictiveSampler
• drop example file from PR after all other tests are implemented

[vincekurtz]

LGTM