`evorl.evaluators.mo_brax_evaluator`¶

Module Contents¶

Classes¶

BraxEvaluator

Mutli-objective evaluator for Brax environments.

Functions¶

`eval_env_step`
`eval_metrics`
`eval_rollout_episode`	Evaulate a batch of episodic trajectories.
`fast_eval_metrics`	Fast evaulate a batch of episodic trajectories.

API¶

class evorl.evaluators.mo_brax_evaluator.BraxEvaluator[source]¶

Bases: evorl.evaluators.evaluator.Evaluator

Mutli-objective evaluator for Brax environments.

Variables:: metric_names – The names of the metrics to evaluate, default is (“reward”, “episode_lengths”)

evaluate(agent_state: chex.ArrayTree, key: chex.PRNGKey, num_episodes: int) → chex.ArrayTree[source]¶

metric_names: tuple[str]¶: ‘pytree_field(…)’

evorl.evaluators.mo_brax_evaluator.eval_env_step(env_fn: evorl.envs.EnvStepFn, action_fn: evorl.agent.AgentActionFn, env_state: evorl.envs.EnvState, agent_state: evorl.agent.AgentState, key: chex.PRNGKey, metric_names: tuple[str] = ()) → tuple[evorl.rollout.SampleBatch, evorl.envs.EnvState][source]¶

evorl.evaluators.mo_brax_evaluator.eval_metrics(env_fn: collections.abc.Callable[[evorl.envs.EnvState, evorl.types.Action], evorl.envs.EnvState], action_fn: collections.abc.Callable[[evorl.agent.AgentState, evorl.rollout.SampleBatch, chex.PRNGKey], tuple[evorl.types.Action, evorl.types.PolicyExtraInfo]], env_state: evorl.envs.EnvState, agent_state: evorl.agent.AgentState, key: chex.PRNGKey, rollout_length: int, discount: float, metric_names: tuple[str] = ()) → tuple[evorl.types.PyTreeDict, evorl.envs.EnvState][source]¶

evorl.evaluators.mo_brax_evaluator.eval_rollout_episode(env_fn: collections.abc.Callable[[evorl.envs.EnvState, evorl.types.Action], evorl.envs.EnvState], action_fn: collections.abc.Callable[[evorl.agent.AgentState, evorl.rollout.SampleBatch, chex.PRNGKey], tuple[evorl.types.Action, evorl.types.PolicyExtraInfo]], env_state: evorl.envs.EnvState, agent_state: evorl.agent.AgentState, key: chex.PRNGKey, rollout_length: int, metric_names: tuple[str] = ()) → tuple[evorl.rollout.SampleBatch, evorl.envs.EnvState][source]¶

Evaulate a batch of episodic trajectories.

The retruned metrics are defined by metric_names.

evorl.evaluators.mo_brax_evaluator.fast_eval_metrics(env_fn: collections.abc.Callable[[evorl.envs.EnvState, evorl.types.Action], evorl.envs.EnvState], action_fn: collections.abc.Callable[[evorl.agent.AgentState, evorl.rollout.SampleBatch, chex.PRNGKey], tuple[evorl.types.Action, evorl.types.PolicyExtraInfo]], env_state: evorl.envs.EnvState, agent_state: evorl.agent.AgentState, key: chex.PRNGKey, rollout_length: int, metric_names: tuple[str] = ()) → tuple[evorl.types.PyTreeDict, evorl.envs.EnvState][source]¶

Fast evaulate a batch of episodic trajectories.