The basic step function for the workflow to update agent.

Tournament selection operator for single objective.

Pytree version of `jax.device_put`.

Pytree version of `jnp.ones_like`.

Update priorities at the last sampled indices. Args: buffer_state: Current buffer state (must have valid sample_indices).

Applies a linear transformation to the inputs along the last dimension. Args: inputs: The nd-array to be transformed. Re

Initialize the distribution. Args: distribution: The distribution to transform. threshold: Clipping value of the action w

Initialize the ParamVectorSpec. Args: params: Provide the structure of the parameters. It should be a single instance of model pa

Initialize a RLWorkflow instance. Args: config: the config object.

Initialize an OnPolicyWorkflow instance. Args: env: Environment object. agent: Workflow-sepecific agent object.

Initialize an OffPolicyWorkflow instance. Args: env: Environment object. agent: Workflow-sepecific agent object.

Initialize the ECWorkflow instance. Args: env: Environment object. agent: Workflow-sepecific agent object.

Initialize the ChainRecorder. Args: recorders: A sequence of recorders to use.

Initialize the log recorder. Args: log_path: The path to the log file. console: Whether to print the log to the conso

Initializes the env wrapper. Args: env: the wrapped env should be a single un-vectorized environment. episode_length:

Initialize the env wrapper. Args: env: the original env num_envs: number of parallel envs.

Initialize the env wrapper. Args: env: the original env num_envs: number of parallel envs.

Initialize the env wrapper. Args: env: the original env num_envs: number of parallel envs.

Initialize the env wrapper. Args: env: the orginal env.

Multi-Agent Brax environment. Compared to the original MABraxEnv, this version disables the autoreset and other wrappers in internal brax env

This implementation follows `The CMA Evolution Strategy: A Tutorial <https://arxiv.org/pdf/1604.00772.pdf>`_. .. note:: CMA-ES in