Deep copy the pytree. Useful for mutable pytree structure like dict. The return also includes a deepcopy of these mutable structures.

Initialize the ECWorkflow instance. Args: config: the config object

Customize the process of building the workflow instance from the config. Args: config: Config of the workflow. Returns:

Customize the process of building the workflow instance from the config. Args: config: Config of the workflow Returns:

Step for Next-Step mode.

Convert the rollout metrics to fitnesses. By default, use the mean of episode_returns over multiple episodes as fitnesses. Args:

Step for Same-Step mode.

Post-setup state before training. By default, no post-setup is applied

Post-setup ReplayBuffer state before training.

Customize the logic of rescaling the config settings when multi-devices training is enabled. When enable_multi_devices=True, rescale config s

Customize the logic of rescaling the config settings when multi-devices training is enabled. When enable_multi_devices=True, rescale config s

Add num_updates support. Therefore this method cannot be jitted.

Calculate the fitness and update the replay buffer and ec_optimizer.

Setup Agent and Optimizer states. Args: key: JAX PRNGKey. Returns: Tuple of (agent_state, opt_state)

Setup Agent and Optimizer states. Args: key: JAX PRNGKey. Returns: Tuple of (agent_state, opt_state).

Setup Agent and ECOptimizer states. Args: key: JAX PRNGKey Returns: Tuple of (agent_state, ec_state)

Setup ReplayBuffer state.

Define Workflow metrics.

Define Workflow metrics.

Define Workflow metrics.

Verifies shapes of the batch leaves against the reference sample. Checks that batch dimensions are the same in all leaves in the batch. Check

Enhanced version of `jax.random.split` that allows batched keys. Args: key: Key or batched keys with shape (B, 2) num: Number of

Actor loss in TD3. Args: sample_barch: [B, ...] Return: LossDict[ actor_loss critic_loss

K actors + 1 shared critic.

K (actor, critic) pairs.

K actors + 1 shared critic.

Finalize the recorder.