    def target_update(self, tau: float) -> "JaxRLTrainState":
        """
        Performs an update of the target params via polyak averaging. The new
        target params are given by:

            new_target_params = tau * params + (1 - tau) * target_params
        """
        new_target_params = jax.tree_map(
            lambda p, tp: p * tau + tp * (1 - tau), self.params, self.target_params
        )
        return self.replace(target_params=new_target_params)
Performs an update of the target params via polyak averaging. The new
target params are given by:

new_target_params = tau * params + (1 - tau) * target_params
serl_launcher/serl_launcher/common/common.py