Approximate KL divergence by K3 estimator (no bias, low variance). See http://joschu.net/blog/kl-approx.html Args: logratio: ratio o

Create Brax environment. Args: env_name: Environment name. kwargs: Arguments passing into Brax. Returns: Brax env.

Random sample.

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

Normalizes along dimension `axis` using an L2 norm. This specialized function exists for numerical stability reasons. Args: x: An in

Normalize obs to PyTreeDict if it's a raw dict. Some envs (e.g., mujoco_playground) return dict obs. PyTreeDict.__setattr__ auto-converts nes

Define how to replace the pop agent_state from the population params. Args: agent_state: State of the agent. params:

Update the obs_preprocessor_state based on sampled trajectories. By default, don't update obs_preprocessor_state. Args:

All-gather the data across all devices.

Get actions from the policy model + add exploraton noise. This method is exclusively used for rollout. Args: agent_state

Define which methods should be jitted. By default, the workflow's `step()` method is jitted.

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

Customize the evaluation logic for the workflow. Args: state: State of the workflow.

Filter out zero-size arrays from the pytree.

Filter out zero-size arrays from the pytree.

Get the normalization layer class based on the type.

Return the output directory of hydra.

Return a specific wrapper of an env.

Helper function that inverts a permutation array.

Creates a Q network for discrete action space: (obs) -> q_values.

Creates multiple MLP networks in parallel.

Set the absl log level.

Senitize the dirname.

Pytree version of `jnp.astype`.