    def create(
        cls, *, apply_fn, params, txs, target_params=None, rng=jax.random.PRNGKey(0)
    ):
        """
        Initializes a new train state.

        Args:
            apply_fn: The function used to apply the model, typically `model_def.apply`.
            params: The model parameters, typically from `model_def.init`.
            txs: The optimizer or pytree of optimizers.
            target_params: The target model parameters.
            rng: The rng key used to initialize the rng chain for `apply_loss_fns`.
        """
        return cls(
            step=0,
            apply_fn=apply_fn,
            params=params,
            target_params=target_params,
            txs=txs,
            opt_states=cls._tx_tree_map(lambda tx: tx.init(params), txs),
            rng=rng,
        )
Initializes a new train state.

Args:
apply_fn: The function used to apply the model, typically `model_def.apply`.
params: The model parameters, typically from `model_def.init`.
txs: The optimizer or pytree of optimizers.
target_params: The target model parameters.
rng: The rng key used to initialize the rng chain for `apply_loss_fns`.
serl_launcher/serl_launcher/common/common.py