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
Edited by Kevin Haninger