Write data to the recorder.

Convert the dataclass to native python structures recursively. The data in the metric object will be converted to local data types: list, tup

Unified env creator. Args: env_cfg: The environment configuration. **kwargs: Additional keyword arguments for the environment cre

Define a pytree field in our dataclass. Args: static: Setting to False will mark the field as static for pytree, where changing data in t

Get the elements of each array in the pytree.

Collect trajectories with length of `rollout_length`. This method is a general rollout method used for collecting trajectories from a vectorized

Get raw value and statistics of a 1D array. Helper function for logging in WandB. Args: data: 1D numpy array. If data has multiple d

Perform soft update on target network. Args: target_params: Target network parameters. source_params: Source network parameters.

Pytree version of `jax.lax.stop_gradient`.

Add prefix to the keys of a dictionary.

Get a tanh transformed normal distribution.

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

Scan with mean aggregation. Usage: same like `jax.lax.scan`, but the scan results will be averaged.

Get a categorical distribution.

Clean the trajectory to make it suitable for the replay buffer.

Creates a policy network.

Unified Version of `jax.random.split` for both single key and batched keys.

Utility function to remove replay_buffer_state from state. Usually used when saving the off-policy workflow state to disk.

Vmap on different gpu.

Compute discount return from the episodic trajectory. Args: rewards: Rewards collected from an episodic trajectory. dones: Dones

Define a metric field in `MetricBase`. Args: reduce_fn: A function to reduce the metric value across different devices. For example, `jax

Set part of each array in the pytree. A Pytree version of `src[idx_or_slice]=target`. Args: src: The source pytree. target:

Create wrapped Brax environment for training. Args: env_name: Environment name. episode_length: Max episode length. paral

Reset the environment to initial state.

Generate new candidate solutions.

Scan and return last iteration results. Usage: same like `jax.lax.scan`, but return the last scan iteration results.

Convert the flatten params to the original params. Args: x: The flatten params. Returns: The original params

Weighted sum. Args: x: (n, ...) w: (n,)

Build the workflow instance from the config. This is the public API to call for instantiating a new workflow object from config. Normally, it

Create wrapped Brax environment for training. Args: env_name: Environment name. episode_length: Max episode length. paral

Random sample from log space. Suitable for hyperparameters that need explore different magnitudes in positive range. eg: [1e-3, 100].

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

Split the key according to the structure of the target pytree.

Split the key to each device.

Create Brax environment. Args: env_name: Environment name. kwargs: Arguments passing into Brax. If you want to use jax backend, s

Wrapper of the loss function that apply gradient updates. Args: loss_fn: The loss function. (params, ...) -> loss optimizer: The

Shift the array to the right with padding.

Calculates the Generalized Advantage Estimation (GAE). Args: rewards: A float32 tensor of shape [T, B] containing rewards generated by

Creates an MLP network.

Creates a Q network: (obs, action) -> value.

Take a step in the environment.

Pytree version of `jnp.zeros_like`.

Initialize the env wrapper. Args: env: the original env num_envs: number of envs to vectorize vmap_step:

Estimate the average episode return from a segmented trajectory. This method does not require the trajectory data from a complete episode.

Average L1 Norm used in TD7.

Close the workflow's components.

Compute the length of the episode. Args: dones: Dones collected from an episodic trajectory. Returns: Length of the episode.

Evaluate the agent based on its state. Args: agent_state: The state of the agent. key: The PRNG key. num_

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

Initialize the state of the . This is the public API to call for instance state initialization.

Detect if a function is wrapped by jit or pmap.

Creates a V network: (obs) -> value.

Convert the original params to flatten params. see `jax.flatten_util.ravel_pytree` Args: x: The original params.

Assert obs has correct batch dimensions.