Skip to content

Advantage Estimation

lerax.advantage.AbstractAdvantageEstimator

Bases: eqx.Module

Estimates advantages and returns from a trajectory segment.

Attributes:

Name Type Description
gamma eqx.AbstractVar[Float[ArrayLike, '']]

Discount factor.

gamma instance-attribute 

gamma: eqx.AbstractVar[Float[ArrayLike, '']]

__call__ abstractmethod 

__call__(
    rewards: Float[Array, " T"],
    values: Float[Array, " T"],
    dones: Bool[Array, " T"],
    last_value: Float[ArrayLike, ""],
) -> tuple[Float[Array, " T"], Float[Array, " T"]]

Compute advantages and returns.

Parameters:

Name Type Description Default
rewards Float[Array, ' T']

Per-step rewards.

 required
values Float[Array, ' T']

Per-step value estimates.

 required
dones Bool[Array, ' T']

Whether each step ended an episode.

 required
last_value Float[ArrayLike, '']

Bootstrap value after the trajectory segment. Strict episodic estimators may ignore it.

 required

Returns:

Type Description
tuple[Float[Array, ' T'], Float[Array, ' T']]

The advantages and returns, each with shape (T,).

lerax.advantage.discounted_returns 

discounted_returns(
    rewards: Float[Array, " T"],
    dones: Bool[Array, " T"],
    final_value: Float[ArrayLike, ""],
    gamma: Float[ArrayLike, ""],
) -> Float[Array, " T"]

Compute discounted returns from a trajectory segment.

Parameters:

Name Type Description Default
rewards Float[Array, ' T']

Per-step rewards

 required
dones Bool[Array, ' T']

Whether each step ended an episode

 required
final_value Float[ArrayLike, '']

Return used to bootstrap the trajectory boundary

 required
gamma Float[ArrayLike, '']

Discount factor

 required

Returns:

Type Description
Float[Array, ' T']

The discounted returns, with shape (T,).

lerax.advantage.GAE

Bases: AbstractAdvantageEstimator

Generalized Advantage Estimation.

Parameters:

Name Type Description Default
gamma Float[ArrayLike, '']

Discount factor

 0.99
lam Float[ArrayLike, '']

Bias-variance tradeoff

 0.95

Attributes:

Name Type Description
gamma Float[ArrayLike, '']

Discount factor
lam Float[ArrayLike, '']

Bias-variance tradeoff

gamma class-attribute instance-attribute 

gamma: Float[ArrayLike, ''] = gamma

lam class-attribute instance-attribute 

lam: Float[ArrayLike, ''] = lam

__init__ 

__init__(
    gamma: Float[ArrayLike, ""] = 0.99,
    lam: Float[ArrayLike, ""] = 0.95,
)

__call__ 

__call__(
    rewards: Float[Array, " T"],
    values: Float[Array, " T"],
    dones: Bool[Array, " T"],
    last_value: Float[ArrayLike, ""],
) -> tuple[Float[Array, " T"], Float[Array, " T"]]

lerax.advantage.BootstrappedReturn

Bases: AbstractAdvantageEstimator

Discounted returns bootstrapped at an incomplete segment boundary.

Parameters:

Name Type Description Default
gamma Float[ArrayLike, '']

Discount factor.

 0.99

Attributes:

Name Type Description
gamma Float[ArrayLike, '']

Discount factor.

gamma class-attribute instance-attribute 

gamma: Float[ArrayLike, ''] = gamma

__init__ 

__init__(gamma: Float[ArrayLike, ''] = 0.99)

__call__ 

__call__(
    rewards: Float[Array, " T"],
    values: Float[Array, " T"],
    dones: Bool[Array, " T"],
    last_value: Float[ArrayLike, ""],
) -> tuple[Float[Array, " T"], Float[Array, " T"]]

lerax.advantage.NStepReturn

Bases: AbstractAdvantageEstimator

N-step bootstrapped returns.

Parameters:

Name Type Description Default
n int

Maximum number of rewards accumulated before bootstrapping.

 required
gamma Float[ArrayLike, '']

Discount factor.

 0.99

Attributes:

Name Type Description
n int

Maximum number of rewards accumulated before bootstrapping.
gamma Float[ArrayLike, '']

Discount factor.

n class-attribute instance-attribute 

n: int = n

gamma class-attribute instance-attribute 

gamma: Float[ArrayLike, ''] = gamma

__init__ 

__init__(n: int, gamma: Float[ArrayLike, ''] = 0.99)

__call__ 

__call__(
    rewards: Float[Array, " T"],
    values: Float[Array, " T"],
    dones: Bool[Array, " T"],
    last_value: Float[ArrayLike, ""],
) -> tuple[Float[Array, " T"], Float[Array, " T"]]