    def apply_gradients(self, *, grads: Any) -> "JaxRLTrainState":
        """
        Only difference from flax's TrainState is that `grads` must have
        `self.txs` as a tree prefix (i.e. where `self.txs` has a leaf, `grads`
        has a subtree with the same structure as `self.params`.)
        """
        updates_and_new_states = self._tx_tree_map(
            lambda tx, opt_state, grad: tx.update(grad, opt_state, self.params),
            self.txs,
            self.opt_states,
            grads,
        )
        updates = self._tx_tree_map(lambda _, x: x[0], self.txs, updates_and_new_states)
        new_opt_states = self._tx_tree_map(
            lambda _, x: x[1], self.txs, updates_and_new_states
        )

        # not the cleanest, I know, but this flattens the leaves of `updates`
        # into a list where leaves are defined by `self.txs`
        updates_flat = []
        self._tx_tree_map(
            lambda _, update: updates_flat.append(update), self.txs, updates
        )

        # apply all the updates additively
        updates_acc = jax.tree_map(
            lambda *xs: jnp.sum(jnp.array(xs), axis=0), *updates_flat
        )
        new_params = optax.apply_updates(self.params, updates_acc)

        return self.replace(
            step=self.step + 1, params=new_params, opt_states=new_opt_states
        )
Only difference from flax's TrainState is that `grads` must have
`self.txs` as a tree prefix (i.e. where `self.txs` has a leaf, `grads`
has a subtree with the same structure as `self.params`.)
serl_launcher/serl_launcher/common/common.py