Gym interface

Brax has a pretty decent interface for MuJoCo-based Gym environments, which would make benchmarking against other ML methods easy.

• Check that jit / jacfwd works on brax envs
• Write a BraxSys shim
  • step
  • cost = -reward
  • Change defaults in brax_sys after loading to see the results
  • need to handle done? not sure
  • get_vec_rollout override so we don't call pipeline_step twice
  • add renderer for videos
  • get constraints on input / state
• Write tests
• Check CEM on InvertedPendulum
• Check single shooting
  • Do tests of grad, jacfwd, jacrev on step / cost. Right now, it's dying w/o explanation I assume OOM error
• Add simulate to roll out
• Add x0 to CEM