Speed Improvements

[alberthli]

As per discussion, there are a number of possible speed improvements we can make to hydrax to move towards deploying on hardware. Here are a few things (list is living, can be updated any time):

• use XLA_FLAGS=--xla_gpu_triton_gemm_any=true. We can set this using os in the __init__.py file of hydrax to just always have it on
• the get_action call, which simply queries the control spline at a certain time, actually has a measurable effect on the speed of the planning loop. we should consider doing all the spline querying on CPU and just moving the nominal spline knots from GPU to CPU, which should incur a low device transfer cost
  • moreover, when deploying online, the spline querying should not happen on the hydrax side anyway, so this is a low prio fix that should only affect the runtime of the sim loop
• the replace call has a small but measurable effect on the runtime. this is probably unavoidable, however.
• the majority of the speed improvements should come from messing with the internals of the controller.optimize function.