Analytical derivatives - #94
Conversation
- Fixed se3.py to use "coadjointbar_se3" and "pcs_SoftJointDifferentialKinematics". - Added a serial inverse-dynamics backward pass returning `dID/dq` and `dID/dqd`. - Added analytical forward-dynamics Jacobians: - `dyd/dy` - `dyd/du` - `dyd/dtau_ext` - Added base PCS force derivative helpers for elastic, damping, and actuation terms. - Added a validation script comparing analytical derivatives against JAX autodiff. forward_dynamics_derivatives dqdd/dq: rel error = 5.782666e-12 max abs error = 1.495151e-06
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3967dc6671
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| import jax | ||
| import jax.numpy as jnp | ||
|
|
||
| jax.config.update("jax_enable_x64", True) |
There was a problem hiding this comment.
Because the CI workflow runs python -m pytest from the repo root and there is no testpaths override, pytest's default test_*.py discovery will import this example module during collection. That import executes this top-level jax.config.update, globally switching JAX to x64 for the rest of the test session even though the script has no tests, so unrelated tests can run under different dtype settings depending on collection. Rename the example or move the config under the if __name__ == "__main__" block.
Useful? React with 👍 / 👎.
| d_elastic_dq = self.elastic_force_derivative_q(q) | ||
| d_damping_dq, d_damping_dqd = self.damping_force_derivatives(q, qd) | ||
| d_actuation_dq = self.actuation_force_derivative_q(q, u) |
There was a problem hiding this comment.
When these new public Jacobian methods are called on TendonActuatedPCS, these calls still resolve to the base helper implementations, but that subclass overrides elastic_force, damping_matrix, and actuation_matrix with q-dependent tendon terms. In that context dtau_dq drops the passive-tendon elastic derivative, the q derivative of D(q) @ qd, and ∂(A(q)u)/∂q, so forward_dynamics_state_jacobian is silently wrong for tendon robots with passive tendons or nonzero tendon inputs. Override these helpers in the tendon subclass or make the analytical Jacobian methods base-PCS-only.
Useful? React with 👍 / 👎.
mstoelzle
left a comment
mstoelzle
left a comment
There was a problem hiding this comment.
Thanks, @mohammedtarnini, for implementing this! This is a great step! And it is very exciting that the tests comparing against the autdiff jacobians already seem to pass! 🎉
I have some comments. Most of them are related to the goal of making the interface and style of the implementation coherent and consistent with the rest of the package.
Some overarching points for improvement:
- Please add full docstrings for all methods/functions that explains what the function does, lists all arguments and return values, including array shapes, physical units, assumptions, etc.
- Please add more code comments in general to explain the implementation and the code to make it easier to maintain and extend
- Please make the variable naming more consistent with the rest of the codebase
- Please make the argument order more consistent with the rest of the codebase
Also, please merge the current main branch into this branch and resolve the merge conflict with se3.py
Next, I would suggest adding some benchmarking code so we have an initial idea of how much faster or slower the analytical derivative is compared to an autodiff-based derivative?
| return (Omega, Z, g, T, S, Sd, f, fd, adjOmegap, dSdq_qd, dSdq_qdd, dSddq_qd) | ||
|
|
||
|
|
||
| def pcs_SoftJointDifferentialKinematics( |
There was a problem hiding this comment.
We recently clean-up the re-organized the lie algebra functions a bit. Specifically, we tried to separate more between core SO3 and SE3 methods and constant-strain-specific solutions. This function would likely be better suited for constant_strain.py?
|
|
||
|
|
||
| def pcs_SoftJointDifferentialKinematics( | ||
| eps, H, Phi, xi_star, q, qd, qdd |
There was a problem hiding this comment.
Please make the argument names and argument order more coherent with the rest of the codebase. For example, eps usually comes last. Also, we (likely) use strain_basis instead of Phi, etc.
| eps, H, Phi, xi_star, q, qd, qdd | ||
| ) -> tuple[ | ||
| Array, Array, Array, Array, Array, Array, Array, Array, Array, Array, Array, Array | ||
| ]: |
There was a problem hiding this comment.
Please add full docstrings, including list of arguments and return values. AI agents can also help you with that.
| Omegahatp2 = Omegahat @ Omegahat | ||
| Omegahatp3 = Omegahatp2 @ Omegahat | ||
|
|
||
| adjOmegap = adjOmegap.at[0:6, :].set(adjoint_se3(Omega)) |
There was a problem hiding this comment.
As mentioned during the last meeting, I would prefer a stacking/concatenation of array instead of first initializing and then overwriting blocks of the matrix as this is sometimes slower.
|
|
||
| fd = jnp.array( | ||
| [ | ||
| t3 / (2 * tp3), |
There was a problem hiding this comment.
Most/all of these vairables are already arrays. Therefore, we shouldn't use jnp.array but instead jnp.concatenate or jnp.stack for constructing the larger array (fd in this case)
| """ | ||
| Return analytical ``dqdd/dq`` and ``dqdd/dqd``. | ||
|
|
||
| This is the unconstrained PCS equivalent of the MATLAB ``ODEJacobian`` |
There was a problem hiding this comment.
I don't think the references to the MATLAB implementation should be here.
| self, | ||
| q: Array, | ||
| qd: Array, | ||
| qdd: Array, |
There was a problem hiding this comment.
Why do you require qdd as an input? This means that the user has to execute another evaluation of the forward dynamics first which causes a lot of redundant computation. Instead, this function could compute qdd in the same pass and return it as well.
| Return the local PCS interval terms used by the forward and backward passes. | ||
|
|
||
| ``eta_plus`` and ``etad_plus`` are expressed in the frame before the step, | ||
| matching the MATLAB recursion immediately before multiplying by |
There was a problem hiding this comment.
Remove the references to MATLAB
| if u is None: | ||
| u = jnp.zeros((self.num_actuators,), dtype=q.dtype) | ||
|
|
||
| mass_matrix = self.inertia_matrix(q) |
There was a problem hiding this comment.
Calling such functions like the inertia_matrix separately will be very expensive. Can you get it "for free" during the inverse_dynamics_derivatives pass?
There was a problem hiding this comment.
Later, this content should go as pytest-compatible test functions into tests/systems.
2 participants
Added backward pass and obtained analytical derivatives.
dID/dqanddID/dqd.dyd/dydyd/dudyd/dtau_extforward_dynamics_derivatives dqdd/dq:
rel error = 5.782666e-12
max abs error = 1.495151e-06
Codes can be further optimized for speed in the future, but they are giving correct values as per the test done. custom_jvp still not used and no comparison for speed against autodiff, but it should be feasible with the current state of the code.