evorl.workflows

Package Contents

Classes

ECWorkflow	Base Workflow for EC (Evolutionary Computation) algorithms.
ECWorkflowTemplate	Workflow template for EC algorithms.
MultiObjectiveECWorkflowTemplate	Workflow template for multi-objective EC algorithms.
OffPolicyWorkflow	Workflow template for Off-Policy RL algorithms.
OnPolicyWorkflow	Workflow template for On-Policy RL algorithms.
RLWorkflow	Base Workflow for RL algorithms.
Workflow	The base class for all Workflows.

API

class evorl.workflows.ECWorkflow(config: omegaconf.DictConfig)

Bases: evorl.workflows.workflow.Workflow

Base Workflow for EC (Evolutionary Computation) algorithms.

classmethod build_from_config(config: omegaconf.DictConfig, enable_multi_devices: bool = False, enable_jit: bool = True)

Build the ec workflow instance from the config.

Parameters:

• config – Config of the workflow

• enable_multi_devices – Whether multi-devices training is enabled

• enable_jit – Whether jit is enabled

classmethod enable_jit() → None

Define which methods should be jitted.

By default, the workflow’s step() method is jitted.

property enable_multi_devices: bool

Whether multi-devices training is enabled.

classmethod enable_shmap(axis_name) → None

Define which methods should be shmaped.

This method defines the multi-device behavior. By default, the workflow’s step() method is shmaped.

class evorl.workflows.ECWorkflowTemplate(*, env: evorl.envs.Env, agent: evorl.agent.Agent, ec_optimizer: evorl.ec.optimizers.EvoOptimizer, ec_evaluator: evorl.evaluators.Evaluator | evorl.evaluators.EpisodeCollector, agent_state_vmap_axes: evorl.agent.AgentStateAxis = 0, config: omegaconf.DictConfig)

Bases: evorl.workflows.ec_workflow.ECWorkflow

Workflow template for EC algorithms.

Variables:

• env – Environment object.

• agent – Workflow-sepecific agent object.

• ec_optimizer – EC Optimizer of the agent.

• ec_evaluator – Evaluator object used in self.evaluation().

• agent_state_vmap_axes – Vmap axis for the agent state.

• config – Config of the workflow.

classmethod enable_jit() → None

classmethod enable_shmap(axis_name) → None

setup(key: chex.PRNGKey) → evorl.types.State

step(state: evorl.types.State) → tuple[evorl.metrics.MetricBase, evorl.types.State]

class evorl.workflows.MultiObjectiveECWorkflowTemplate(*, env: evorl.envs.Env, agent: evorl.agent.Agent, ec_optimizer: evorl.ec.optimizers.EvoOptimizer, ec_evaluator: evorl.evaluators.Evaluator | evorl.evaluators.EpisodeCollector, agent_state_vmap_axes: evorl.agent.AgentStateAxis = 0, config: omegaconf.DictConfig)

Bases: evorl.workflows.ec_workflow.ECWorkflowTemplate

Workflow template for multi-objective EC algorithms.

step(state: evorl.types.State) → tuple[evorl.metrics.MetricBase, evorl.types.State]

class evorl.workflows.OffPolicyWorkflow(env: evorl.envs.Env, agent: evorl.agent.Agent, optimizer: optax.GradientTransformation, evaluator: evorl.evaluators.Evaluator, replay_buffer: evorl.replay_buffers.AbstractReplayBuffer, config: omegaconf.DictConfig)

Bases: evorl.workflows.rl_workflow.RLWorkflow

Workflow template for Off-Policy RL algorithms.

This class constructs the template for Off-Policy RL algorithms, providing the general setup() and evaluate() methods.

evaluate(state: evorl.types.State) → tuple[evorl.metrics.MetricBase, evorl.types.State]

setup(key: chex.PRNGKey) → evorl.types.State

class evorl.workflows.OnPolicyWorkflow(env: evorl.envs.Env, agent: evorl.agent.Agent, optimizer: optax.GradientTransformation, evaluator: evorl.evaluators.Evaluator, config: omegaconf.DictConfig)

Bases: evorl.workflows.rl_workflow.RLWorkflow

Workflow template for On-Policy RL algorithms.

This class constructs the template for On-Policy RL algorithms, providing the general setup() and evaluate() methods.

evaluate(state: evorl.types.State) → tuple[evorl.metrics.MetricBase, evorl.types.State]

setup(key: chex.PRNGKey) → evorl.types.State

class evorl.workflows.RLWorkflow(config: omegaconf.DictConfig)

Bases: evorl.workflows.workflow.Workflow

Base Workflow for RL algorithms.

classmethod build_from_config(config: omegaconf.DictConfig, enable_multi_devices: bool = False, enable_jit: bool = True) → typing_extensions.Self

Build the rl workflow instance from the config.

Parameters:

• config – Config of the workflow.

• enable_multi_devices – Whether multi-devices training is enabled.

• enable_jit – Whether jit is enabled.

classmethod enable_jit() → None

Define which methods should be jitted.

By default, the workflow’s step() and evaluate() methods are jitted.

property enable_multi_devices: bool

Whether multi-devices training is enabled.

classmethod enable_shmap(axis_name: str) → None

Define which methods should be shmaped.

This method defines the multi-device behavior. By default, the workflow’s step() and evaluate() methods are shmaped.

Parameters:

axis_name – The axis_name for shmap.

abstract evaluate(state: evorl.types.State) → tuple[evorl.metrics.MetricBase, evorl.types.State]

Customize the evaluation logic for the workflow.

Parameters:

state – State of the workflow.

abstract step(state: evorl.types.State) → tuple[evorl.metrics.MetricBase, evorl.types.State]

Customize the training logic of one iteration.

Parameters:

state – State of the workflow.

Returns:

Tuple of (metrics, state).

class evorl.workflows.Workflow(config: omegaconf.DictConfig)

Bases: evorl.workflows.workflow.AbstractWorkflow

The base class for all Workflows.

All workflow classes are inherit from this class, and customize by implementing

add_recorders(recorders: evorl.recorders.Recorder) → None

abstract classmethod build_from_config(config: omegaconf.DictConfig, *args, **kwargs) → typing_extensions.Self

Build the workflow instance from the config.

This is the public API to call for instantiating a new workflow object from config. Normally, it will call init() and do some pre- and post-processing.

Parameters:

config – config object

Returns:

A workflow instance

close() → None

Close the workflow’s components.

init(key: chex.PRNGKey) → evorl.types.State

Initialize the state of the .

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

abstract learn(state: evorl.types.State) → evorl.types.State

Run the complete learning process.

The learning process includes:

• call multiple times of step()

• record the metrics

• save checkpoints

abstract setup(key: chex.PRNGKey) → evorl.types.State