RewardShapingPipeline
RewardShapingPipeline(reward_shapers: RewardShaperConfiguration, context: RewardShapingContext) -> NoneSequential trainer-owned reward-shaping pipeline.
Initialize the reward-shaping pipeline.
Parameters:
reward_shapers(RewardShaperConfiguration) – None, a single shaper type, or a list of shaper types.context(RewardShapingContext) – Shared shaping runtime context.
Properties
has_intrinsic_reward_stream
has_intrinsic_reward_stream: boolReturn whether any configured shaper contributes to the intrinsic reward stream.
Returns:
bool–Truewhen at least one shaper targets the intrinsic stream.
is_enabled
is_enabled: boolReturn whether any reward shapers are configured.
Returns:
bool–Truewhen at least one shaper is active.
Methods
complete_episode
complete_episode(env_id: int, step: int) -> NoneLog and clear reward-shaping episode metrics for one environment.
Parameters:
env_id(int) – The environment ID.step(int) – The global step at episode end.
flush_metrics
flush_metrics(step: int) -> NoneEmit any throttled actor or learner reward-shaping metrics.
Parameters:
step(int) – The step at which to log.
learn
learn(batch: RewardShapingBatch) -> NoneRun learner-side reward-shaper updates for one batch.
Parameters:
batch(RewardShapingBatch) – The learner batch to consume.
load_state_dict
load_state_dict(state: dict[str, Any]) -> NoneRestore pipeline state from a checkpoint.
Parameters:
state(dict[str, Any]) – The serialised pipeline state.
on_episode_start
on_episode_start(env_id: int, initial_observation: np.ndarray) -> NoneNotify all configured shapers about a reset.
Parameters:
env_id(int) – The environment ID.initial_observation(ndarray) – The reset observation.
restore_from_checkpoint
restore_from_checkpoint(checkpoint: Checkpoint) -> NoneRestore reward-shaper state from a checkpoint.
Parameters:
checkpoint(Checkpoint) – The checkpoint to restore from.
Raises:
OSError– If the checkpoint contents do not match the configured reward shapers.
shape
shape(step: RewardShapingStep) -> RewardShapingResultApply all configured shapers in sequence and log reward metrics.
Parameters:
step(RewardShapingStep) – The step information to shape.
Returns:
RewardShapingResult– The shaped reward result.
shape_batch
shape_batch(steps: Sequence[RewardShapingStep]) -> list[RewardShapingResult]Apply all configured shapers to a batch of steps.
Shapers are run sequentially, but each individual shaper can override RewardShaper.shape_batch to do efficient batched inference.
Parameters:
steps(Sequence[RewardShapingStep]) – The steps to shape.
Returns:
list[RewardShapingResult]– The shaped reward results in the same order.
state_dict
state_dict() -> dict[str, Any]Capture pipeline state for checkpointing.
Returns:
dict[str, Any]– A serializable state dictionary.
to_checkpoint_file
to_checkpoint_file(compress: bool = True) -> CheckpointFile | NoneBuild the checkpoint file for this pipeline.
Parameters:
compress(bool) – Whether to zstd-compress the checkpoint file.
Returns:
CheckpointFile | None– The checkpoint file, orNonewhen reward shaping is disabled.