Customize rewards¶

cfg.reward_fn is the function the env calls every step to score the new state for both sides. Two paths to swap it:

1. Pass it as a constructor kwarg — ScenarioConfig(..., reward_fn=MyReward()). This is the canonical override path; __post_init__ only consults cfg.game.default_reward_fn() when reward_fn is None.
2. dataclasses.replace — cfg = replace(cfg, reward_fn=MyReward()). Useful when you already have a cfg in hand and want to swap one field without re-listing the rest.

Where defaults come from¶

Every Game subclass declares its default reward (and termination) via default_reward_fn / default_termination_fn. ScenarioConfig.__post_init__ calls those when cfg.reward_fn is None. The mapping is:

The protocol¶

A RewardFn returns a jax.Array whose shape depends on scope:

• PER_SIDE — scalar (()).
• PER_VEHICLE — (N_side,).

It receives both sides' actions in one call (so multi-side rewards like zero-sum penalties are natural), the previous and next env states, the Side it's scoring, the cfg via params, and the step counter t.

Bundled implementations¶

ZeroReward is the no-op reward (NoGame's default). It always returns 0; useful as an explicit "I will supply my own reward" placeholder:

ZeroReward(scope: RewardScope = RewardScope.PER_SIDE)

DistanceToReferenceOrbit is a reference single-agent reward — negative sum of guard distances to the reference origin, zero for the bandit. Useful when you want a guard-only signal without LBG event semantics:

DistanceToReferenceOrbit(
    scope: RewardScope = RewardScope.PER_SIDE,
)

LbgZeroSumReward is the LBG default — each side gets potential-based shaping over the distance it is trying to close, plus mirrored terminal events on catch / breach, and the guard side pays a separation charge. Reads catch_radius_m and breach_radius_m, and the dwell each event requires, from its own fields (the LadyBanditGuard.default_reward_fn builder threads those from the game knobs):

LbgZeroSumReward(
    shaping_scale_m: float = 300.0,
    shaping_gain: float = 50.0,
    shaping_discount: float = 1.0,
    home_weight: float = 0.0,
    guard_separation_m: float = 20.0,
    lady_keepout_m: float = 20.0,
    separation_cost: float = 10.0,
    r_catch: float = 1000.0,
    r_breach: float = 1000.0,
    catch_radius_m: float = 50.0,
    breach_radius_m: float = 5.0,
    catch_speed_mps: float = float("inf"),
    breach_speed_mps: float = float("inf"),
    catch_dwell_steps: int = 0,
    breach_dwell_steps: int = 0,
    escape_radius_m: float = 0.0,
    repel_on_empty_tank: bool = False,
    dv_cost: float = 0.0,
    scope: RewardScope = RewardScope.PER_SIDE,
)

potential(state, side: Side) -> jax.Array

The potentials and the separation charge are module-level functions, so a custom reward or a value estimate can reuse the same geometry:

lbg_distances(
    state,
) -> tuple[jax.Array, jax.Array, jax.Array]

lbg_potential_from_distances(
    side: Side,
    d_gb: Array,
    d_bl: Array,
    d_gl: Array,
    shaping_scale_m: float,
    home_weight: float,
) -> jax.Array

lbg_potential(
    state,
    side: Side,
    shaping_scale_m: float,
    home_weight: float,
) -> jax.Array

guard_separation_cost(
    state,
    guard_separation_m: float,
    lady_keepout_m: float,
    separation_cost: float,
) -> jax.Array

Pursuit-Evasion ships its own zero-sum reward (PursuitEvasionReward) that reads cfg.game.capture_distance_m. Sun-Blocking ships SunBlockingReward, and Observation-Blocking ships ObservationBlockingReward.

Worked example: a sparse "breach happened" reward¶

Replace the shaped default with a sparse penalty that fires only on the breach step:

    import dataclasses
    from dataclasses import dataclass

    import jax.numpy as jnp

    from orbitalgym import OrbitalGymEnv, Side, SingleAgentView, make_lady_bandit_guard
    from orbitalgym.policies import ZeroControl
    from orbitalgym.rewards.base import RewardScope

    @dataclass(frozen=True)
    class SparseBreachReward:
        """+0 every step; -1 to the guard on the step where the bandit
        breaches the lady (and +1 to the bandit, mirrored)."""

        breach_radius_m: float
        scope: RewardScope = RewardScope.PER_SIDE

        def __call__(self, prev_state, action, next_state, side, params, t):
            del prev_state, action, params, t
            bandits = next_state.bandits
            positions = bandits.rtn[:, :3] if hasattr(bandits, "rtn") else bandits.rt[:, :2]
            min_dist = jnp.min(jnp.linalg.norm(positions, axis=-1))
            breached = min_dist < self.breach_radius_m
            # Guard loses 1 on breach; bandit mirrors.
            penalty = jnp.where(breached, -1.0, 0.0)
            return jnp.where(side == Side.GUARD, penalty, -penalty)

Wire it onto the cfg:

    cfg = make_lady_bandit_guard(seed=0, max_horizon_s=200.0)
    cfg = dataclasses.replace(cfg, reward_fn=SparseBreachReward(breach_radius_m=5.0))
    env = OrbitalGymEnv(cfg)

Per-game applicability¶

See also¶

• In depth → Symmetric core & data shapes — what's in prev_state / next_state.
• API → Components → Rewards — auto-generated reference.
• Game guides — each game's reward in context.