Bugfixes and improvements suggestions

[João Sousa Pinto]

I recently implemented a multiple-shooting variant of iLQR here: https://github.com/joaospinto/primal_dual_ilqr

Much of my implementation (apart from core algorithmic changes) is inspired on trajax. While implementing this, I found a few issues/possible improvements related to trajax itself, which I'll list below.

1. In

   trajax/trajax/optimizers.py

   Line 798 in c94a637

    Q = lax.cond(make_psd, Q, psd, Q, lambda x: x) 

   , you are projecting the Q matrices to become positive semi-definite. When the M matrices (i.e. cross-state-and-control quadratic terms) are non-zero, this is not sufficient. You should do this instead: https://github.com/joaospinto/primal_dual_ilqr/blob/main/primal_dual_ilqr/optimizers.py#L60

2. Perhaps because of the issue above, you resort to doing least-square solves in

   trajax/trajax/tvlqr.py

   Line 94 in c94a637

    K_k, *_ = np.linalg.lstsq( 

   , when you should be able to use Cholesky solves instead (didn't benchmark this part in JAX, but generally should be faster); see https://github.com/joaospinto/primal_dual_ilqr/blob/main/primal_dual_ilqr/primal_tvlqr.py#L38.

3. In

   trajax/trajax/tvlqr.py

   Line 61 in c94a637

    def lqr_step(P, p, Q, q, R, r, M, A, B, c, delta=1e-8): 

   , some of the terms you have in your code are mathematically guaranteed to be zero and can be removed; this is solved in https://github.com/joaospinto/primal_dual_ilqr/blob/main/primal_dual_ilqr/primal_tvlqr.py#L38.

4. In places like

   trajax/trajax/tvlqr.py

   Line 57 in c94a637

    return lax.fori_loop(0, T, body, (X, U)) 

   , you use fori_loop, when using a scan results in significant speed-ups. You can find almost-drop-in replacements here: https://github.com/joaospinto/primal_dual_ilqr/blob/main/primal_dual_ilqr/primal_tvlqr.py

5. It would be interesting to add support for a GPU-accelerated implementation of LQR; see https://github.com/joaospinto/primal_dual_ilqr/blob/main/primal_dual_ilqr/primal_tvlqr.py#L96.

I'd be happy to try to merge my code (https://github.com/joaospinto/primal_dual_ilqr) into this repository (either just these improvements, or actually adding the new algorithm itself), if you find that interesting. Let me know!