evorl.types

Common type annotations and data structures.

Module Contents

Classes

PyTreeArrayMixin	batch operate pytree with jax.Array.
PyTreeData	A pytree dataclass for Data.
PyTreeDict	An easydict with pytree support.
PyTreeNode	A pytree dataclass for Node.
State	A general State class.

Functions

dataclass
pytree_field	Define a pytree field in our dataclass.

Data

Action
ActionPostprocessorParams
AgentID
Axis
Done
EnvInternalState
ExtraInfo
LossDict
MISSING_REWARD
Metrics
ObsPreprocessorParams
Observation
Params
PolicyExtraInfo
ReplayBufferState
Reward
RewardDict

API

evorl.types.Action

None

evorl.types.ActionPostprocessorParams

None

evorl.types.AgentID

None

evorl.types.Axis

None

evorl.types.Done

None

evorl.types.EnvInternalState

None

evorl.types.ExtraInfo

None

evorl.types.LossDict

None

evorl.types.MISSING_REWARD

None

evorl.types.Metrics

None

evorl.types.ObsPreprocessorParams

None

evorl.types.Observation

None

evorl.types.Params

None

evorl.types.PolicyExtraInfo

None

class evorl.types.PyTreeArrayMixin

batch operate pytree with jax.Array.

It assumes all arrays have the same head shape.

property T

concatenate(*others: chex.ArrayTree, axis: int = 0) → chex.ArrayTree

index_set(idx: jax.Array | collections.abc.Sequence[jax.Array], o: chex.ArrayTree) → chex.ArrayTree

index_sum(idx: jax.Array | collections.abc.Sequence[jax.Array], o: chex.ArrayTree) → chex.ArrayTree

reshape(shape: collections.abc.Sequence[int]) → chex.ArrayTree

slice(beg: int, end: int, strides=None) → chex.ArrayTree

take(i, axis=0) → chex.ArrayTree

class evorl.types.PyTreeData

A pytree dataclass for Data.

Like PyTreeNode, but all fileds must be set at init, and not allow set_frozen_attr() method.

class evorl.types.PyTreeDict(*args, **kwargs)

Bases: dict

An easydict with pytree support.

copy()

pop(k, d=None)

replace(**d)

tree_flatten()

classmethod tree_unflatten(aux_data, children)

update(e=None, **f)

class evorl.types.PyTreeNode

A pytree dataclass for Node.

set_frozen_attr(name, value)

Force set attribute after init of the dataclass.

evorl.types.ReplayBufferState

None

evorl.types.Reward

None

evorl.types.RewardDict

None

class evorl.types.State(*args, **kwargs)

Bases: evorl.types.PyTreeDict

A general State class.

An alias of PyTreeDict. This class is specfically used for Workflow state.

evorl.types.dataclass(clz: evorl.types._T, *, pure_data=False, **kwargs) → evorl.types._T

evorl.types.pytree_field(*, static=False, **kwargs)

Define a pytree field in our dataclass.

Parameters:

static – Setting to False will mark the field as static for pytree, where changing data in these fields will cause a re-jit of func.

Returns:

A dataclass field.