Returns ranks in [0, len(x)-1]. Note: This is different from scipy.stats.rankdata, which returns ranks in [1, len(x)].

Get ranks in [0, len(x)-1]. This is different from `scipy.stats.rankdata`, which returns ranks in [1, len(x)].

Create an EnvPool environment with Gymnasium API. Unlike other jax-based env, most wrappers are handled inside the envpool.

Create a gym env based on Gymnasium. Unlike other jax-based env, most wrappers are handled inside the gymnasium.

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

Collect one-step data.

Collect one-step data.

Define which methods should be jitted. By default, the workflow's `step()` and `evaluate()` methods are jitted.

Define which methods should be shmaped. This method defines the multi-device behavior. By default, the workflow's `step()` and `evaluate()` m

Collect one-step data.

Evaulate a batch of episodic trajectories. The retruned metrics are defined by `metric_names`.

Evaluate the policy with the mean of ES.

Fast evaulate a batch of episodic trajectories. The retruned metrics are defined by `metric_names`.

Fast evaulate a batch of episodic trajectories. A even faster implementation than `eval_rollout_episode()`. It achieves early termination when it

Flatten the trajectory from [#pop, T, B, ...] to [#pop*T*B, ...].

Fold multiple steps into a single step function.

Return the global rank for each device. Returns: The sharded ranks across devices. Each device has a unique rank.

Return the node id in multi-node distributed env.

Handle episode return array with MISSING_REWARD, i.e., returned from multiple call of average_episode_discount_return.

Initialize the recorder.

Whether the JAX's distributed setting is initialized.

IMPALA loss. Args: trajectory: [T, B, ...] a sequence of transitions, not shuffled timesteps

Creates a policy network.

Mutation for MLP. Args: key: PRNGKey x: single individual, vec_relative_prob: probability of mutating a vector(1-d) param

Reset the vmapped env. Args: key: support batched keys [B,2] or single key [2]

Recompute max_priority from current buffer entries.

Save state to a file. Args: path: Checkpoint path. state: The state to be saved.

Select parents to replace worse individuals.

Force set attribute after __init__ of the dataclass.

Setup checkpoint manager.

Slide a window over the fist axis of the array. Change shape from [T, ...] to [L, W, ...], where W = (T - L) // S + 1 is the number of windows.