Skip to content

Curriculum

lerax.curriculum.ScheduledCurriculum

Bases: AbstractStatelessCallback

Modify an environment field on a fixed schedule.

Compose instances with CallbackList to schedule multiple fields.

Attributes:

Name Type Description
where Callable

Environment-field selector such as lambda env: env.mass.

schedule_fn Callable

Maps iteration count to the scheduled value.

Example::

from lerax.curriculum import ScheduledCurriculum, linear_schedule

curriculum = ScheduledCurriculum(
    where=lambda env: env.m,
    schedule_fn=linear_schedule(start=0.5, end=2.0, total=1000),
)
algo.learn(env, policy, total_timesteps=..., key=key, callback=curriculum)

where class-attribute instance-attribute

where: Callable = eqx.field(static=True)

schedule_fn class-attribute instance-attribute

schedule_fn: Callable = eqx.field(static=True)

reset

reset(
    ctx: ResetContext, *, key: Key[Array, ""]
) -> EmptyCallbackState

step_reset

step_reset(
    ctx: ResetContext, *, key: Key[Array, ""]
) -> EmptyCallbackStepState

on_step

on_step(ctx: StepContext, *, key: Key[Array, ''])

on_iteration

on_iteration(ctx: IterationContext, *, key: Key[Array, ''])

on_training_start

on_training_start(ctx, *, key: Key[Array, ''])

on_training_end

on_training_end(ctx, *, key: Key[Array, ''])

continue_training

continue_training(
    ctx: IterationContext, *, key: Key[Array, ""]
)

apply_curriculum

apply_curriculum[S: "AbstractAlgorithmState"](
    state: S, callback_state: EmptyCallbackState
) -> tuple[S, EmptyCallbackState]

lerax.curriculum.AbstractAdaptiveCurriculum

Bases: AbstractCallback[AdaptiveCurriculumState, AdaptiveCurriculumStepState]

Abstract base for adaptive curricula tracking a performance metric.

Attributes:

Name Type Description
metric_fn Callable

(done, reward, locals_dict) -> scalar function called each step and accumulated per episode.

smoothing float

Running-metric EMA factor; higher values favor recent episodes.

metric_fn class-attribute instance-attribute

metric_fn: Callable = eqx.field(static=True)

smoothing instance-attribute

smoothing: float

reset

reset(
    ctx: ResetContext, *, key: Key[Array, ""]
) -> AdaptiveCurriculumState

step_reset

step_reset(
    ctx: ResetContext, *, key: Key[Array, ""]
) -> AdaptiveCurriculumStepState

on_step

on_step(
    ctx: StepContext, *, key: Key[Array, ""]
) -> AdaptiveCurriculumStepState

on_iteration

on_iteration(
    ctx: IterationContext, *, key: Key[Array, ""]
) -> AdaptiveCurriculumState

on_training_start

on_training_start(
    ctx: TrainingContext, *, key: Key[Array, ""]
) -> AdaptiveCurriculumState

on_training_end

on_training_end(
    ctx: TrainingContext, *, key: Key[Array, ""]
) -> AdaptiveCurriculumState

continue_training

continue_training(
    ctx: IterationContext, *, key: Key[Array, ""]
) -> Bool[Array, ""]

apply_curriculum abstractmethod

apply_curriculum[S: "AbstractAlgorithmState"](
    state: S, callback_state: AdaptiveCurriculumState
) -> tuple[S, AdaptiveCurriculumState]

Modify algorithm state from the tracked metric.

Called after on_iteration each iteration with the running metric and current level in callback_state.

lerax.curriculum.LevelCurriculum

Bases: AbstractAdaptiveCurriculum

Adapt an environment field through discrete parameter levels.

Advance when the running performance metric exceeds threshold.

Parameters:

Name Type Description Default
where Callable

Selector for the env field to modify.

required
levels Float[Array, ' num_levels']

Parameter values per level.

required
metric_fn Callable

Per-step metric extraction function.

required
threshold float

Advancement threshold.

required
smoothing float

EMA smoothing factor. Defaults to 0.05.

0.05

Example::

from lerax.curriculum import LevelCurriculum

curriculum = LevelCurriculum(
    where=lambda env: env.max_speed,
    levels=jnp.array([4.0, 6.0, 8.0]),
    metric_fn=lambda done, reward, locals: reward,
    threshold=100.0,
)
algo.learn(env, policy, total_timesteps=..., key=key, callback=curriculum)

where class-attribute instance-attribute

where: Callable = where

levels instance-attribute

levels: Float[Array, ' num_levels'] = levels

threshold instance-attribute

threshold: float = threshold

metric_fn instance-attribute

metric_fn = metric_fn

smoothing instance-attribute

smoothing = smoothing

reset

reset(
    ctx: ResetContext, *, key: Key[Array, ""]
) -> AdaptiveCurriculumState

step_reset

step_reset(
    ctx: ResetContext, *, key: Key[Array, ""]
) -> AdaptiveCurriculumStepState

on_step

on_step(
    ctx: StepContext, *, key: Key[Array, ""]
) -> AdaptiveCurriculumStepState

on_iteration

on_iteration(
    ctx: IterationContext, *, key: Key[Array, ""]
) -> AdaptiveCurriculumState

on_training_start

on_training_start(
    ctx: TrainingContext, *, key: Key[Array, ""]
) -> AdaptiveCurriculumState

on_training_end

on_training_end(
    ctx: TrainingContext, *, key: Key[Array, ""]
) -> AdaptiveCurriculumState

continue_training

continue_training(
    ctx: IterationContext, *, key: Key[Array, ""]
) -> Bool[Array, ""]

__init__

__init__(
    where: Callable,
    levels: Float[Array, " num_levels"],
    metric_fn: Callable,
    threshold: float,
    smoothing: float = 0.05,
)

apply_curriculum

apply_curriculum[S: "AbstractAlgorithmState"](
    state: S, callback_state: AdaptiveCurriculumState
) -> tuple[S, AdaptiveCurriculumState]

lerax.curriculum.linear_schedule

linear_schedule(
    start: float, end: float, total: int
) -> Callable[[Int[Array, ""]], Float[Array, ""]]

Interpolate linearly from start to end over total iterations.

Clamps to [start, end] outside the range.

Parameters:

Name Type Description Default
start float

Value at iteration 0.

required
end float

Value at iteration total.

required
total int

Number of iterations for the full transition.

required

Returns:

Type Description
Callable[[Int[Array, '']], Float[Array, '']]

A function mapping iteration count to the scheduled value.

lerax.curriculum.step_schedule

step_schedule(
    values: list[float], boundaries: list[int]
) -> Callable[[Int[Array, ""]], Float[Array, ""]]

Jump between values at specified iteration boundaries.

Parameters:

Name Type Description Default
values list[float]

Stage values; length must be len(boundaries) + 1.

required
boundaries list[int]

Iterations that transition to the next value.

required

Returns:

Type Description
Callable[[Int[Array, '']], Float[Array, '']]

A function mapping iteration count to the scheduled value.

lerax.curriculum.cosine_schedule

cosine_schedule(
    start: float, end: float, total: int
) -> Callable[[Int[Array, ""]], Float[Array, ""]]

Cosine annealing from start to end over total iterations.

Parameters:

Name Type Description Default
start float

Value at iteration 0.

required
end float

Value at iteration total.

required
total int

Number of iterations for the full transition.

required

Returns:

Type Description
Callable[[Int[Array, '']], Float[Array, '']]

A function mapping iteration count to the scheduled value.