Get centered ranks in [-0.5, 0.5].

Determine whether the input is in the space. Returns: A boolean value about whether x is in the space.

Collect given rollout_length trajectory. Tips: when use jax.jit, use: jax.jit(partial(rollout, env, agent)) Args: env: vmapped env w

Centrialized Execution: Collect given rollout_length trajectory. Args: env: vmapped env w/ autoreset Returns: env_state: las

Denormalizes values in a nested structure using the given mean/std. Only values of inexact types are denormalized. See https://numpy.org/doc/

Disable GPU memory preallocation for XLA. Call this method at the beginning of your script.

Enable deterministic mode for JAX. Call this method at the beginning of your script.

Whether multi-devices training is enabled.

Whether multi-devices training is enabled.

Mutation used in the original ERL for MLP. Args: key: PRNGKey x: single individual, vec_relative_prob: probability of mut

Collect one-step data in evaluation mode.

Evaulate a batch of episodic trajectories. It avoids unnecessary calls of `env_step()` when all environments are done. However, the agent's actio

Get the best action from the action distribution. This method is exclusively used for evaluation. Args: agent_state: the

Sample from exp(-x).

Define the exploration operation for PBT. Normally explore the local of an individual. i.e., mutation op in the context of EC. Here we us

Get statistics of a 1D array. Similar to `get_1d_array_statistics`, but instead of recording histogram, WandB will record the raw data.

Get a truncated normal distribution.

Check if the array has NaN values.

Initialize the workflow's state. Args: key: JAX PRNGKey Returns: state: the state of the workflow

Initialize the state of the replay buffer. Args: sample_spec: A single sample or sample spec that contains the pytree structure a