evorl.algorithms.erl.cemrl_workflow

Module Contents

Classes

CEMRLTrainMetric

CEMRLWorkflowBase

Base Class for CEMRL, equipped with many useful methods.

API

class evorl.algorithms.erl.cemrl_workflow.CEMRLTrainMetric[source]

Bases: evorl.metrics.MetricBase

ec_info: evorl.types.PyTreeDict

‘metric_field(…)’

pop_episode_lengths: chex.Array

None

pop_episode_returns: chex.Array

None

rb_size: chex.Array

None

rl_metrics: evorl.metrics.MetricBase | None

None

class evorl.algorithms.erl.cemrl_workflow.CEMRLWorkflowBase(*, env: evorl.envs.Env, agent: evorl.agent.Agent, agent_state_vmap_axes: evorl.agent.AgentStateAxis, optimizer: optax.GradientTransformation, ec_optimizer: evorl.ec.optimizers.EvoOptimizer, collector: evorl.evaluators.EpisodeCollector, evaluator: evorl.evaluators.Evaluator, replay_buffer: evorl.replay_buffers.AbstractReplayBuffer, config: omegaconf.DictConfig)[source]

Bases: evorl.workflows.Workflow

Base Class for CEMRL, equipped with many useful methods.

classmethod build_from_config(config: omegaconf.DictConfig, enable_multi_devices: bool = False, enable_jit: bool = True) → typing_extensions.Self[source]
classmethod enable_jit() → None[source]
abstract evaluate(state: evorl.types.State) → tuple[evorl.metrics.MetricBase, evorl.types.State][source]
setup(key: chex.PRNGKey) → evorl.types.State[source]