Skip to content

Abstract Policies

lerax.policy.AbstractPolicy

Bases: Serializable

Base class for policies.

Policies map observations and internal states to actions and new internal states.

Attributes:

Name Type Description
name eqx.AbstractClassVar[str]

The name of the policy.
action_space eqx.AbstractVar[AbstractSpace[ActType, MaskType]]

The action space of the policy.
observation_space eqx.AbstractVar[AbstractSpace[ObsType, Any]]

The observation space of the policy.

name instance-attribute 

name: eqx.AbstractClassVar[str]

action_space instance-attribute 

action_space: eqx.AbstractVar[
    AbstractSpace[ActType, MaskType]
]

observation_space instance-attribute 

observation_space: eqx.AbstractVar[
    AbstractSpace[ObsType, Any]
]

serialize 

serialize(
    path: str | Path, no_suffix: bool = False
) -> None

Serialize the model to the specified path.

Parameters:

Name Type Description Default
path str | Path

The path to serialize to.

 required
no_suffix bool

If True, do not append the ".eqx" suffix

 False

deserialize classmethod 

deserialize[**Params, ClassType](
    path: str | Path,
    *args: Params.args,
    **kwargs: Params.kwargs,
) -> ClassType

Deserialize the model from the specified path. Must provide any additional arguments required by the class constructor.

Parameters:

Name Type Description Default
path str | Path

The path to deserialize from.

 required
*args Params.args

Additional arguments to pass to the class constructor

 ()
**kwargs Params.kwargs

Additional keyword arguments to pass to the class constructor

 {}

Returns:

Type Description
ClassType

The deserialized model.

__call__ abstractmethod 

__call__(
    state: StateType,
    observation: ObsType,
    *,
    key: Key | None = None,
    action_mask: MaskType | None = None,
) -> tuple[StateType, ActType]

Return the next action and new internal state given the current observation and internal state.

A key can be provided for stochastic policies. If no key is provided, the policy should behave deterministically.

Parameters:

Name Type Description Default
state StateType

The current internal state of the policy.

 required
observation ObsType

The current observation.

 required
key Key | None

An optional JAX random key for stochastic policies.

 None
action_mask MaskType | None

An optional action mask.

 None

Returns:

Type Description
tuple[StateType, ActType]

The new internal state and the action to take.

reset abstractmethod 

reset(*, key: Key) -> StateType

Return an initial internal state for the policy.

Parameters:

Name Type Description Default
key Key

A JAX random key for initializing the state.

 required

Returns:

Type Description
StateType

An initial internal state for the policy.

lerax.policy.AbstractPolicyState

Bases: eqx.Module

Base class for policy internal states.