Skip to content

Overview

Overview

All kernels in rl-triton share a single architectural idea: express the RL recurrence as a linear recurrence of the form

\[A_t = a_t + b_t \cdot A_{t+1}\]

and solve it in \(O(\log N)\) parallel steps using an associative scan on the GPU, rather than the \(O(N)\) sequential loop a naïve implementation requires. The scan runs entirely inside Streaming Multiprocessor SRAM, avoiding repeated round-trips to High Bandwidth Memory (HBM).

For sequences that fit in a single thread block (seq_len ≤ 131072), a fully-fused kernel computes every intermediate quantity - TD errors, IS ratios, decay products - without materialising any intermediate tensors. Longer sequences fall back to a chunked scan that stitches results across blocks.

Choosing the Right Kernel

Situation Kernel
On-policy PPO / A2C advantage estimation compute_gae
Off-policy actor-critic (IMPALA, APPO) compute_vtrace
Off-policy Q-learning, discrete actions compute_retrace
Critic targets with bias-variance control compute_lambda_returns
Simple reward-to-go baseline compute_discounted_returns
TD(λ) parameter updates with traces compute_eligibility_traces
Episode-scoped cumulative statistics compute_episodic_prefix_sum

API Reference

rl_triton.ops.gae.compute_gae(rewards, values, terminateds, truncateds=None, gamma=0.99, lambda_=0.95, bootstrap_values=None, last_value=None)

Compute Generalized Advantage Estimation via a backward associative scan.

Recurrence:

  • A[t] = δ[t] + β[t] * A[t+1], A[T] = 0
  • δ[t] = r[t] + γ·(1-terminated[t]) * v_next[t] - V(s_t)
  • β[t] = γ·λ·(1-done[t]) (scan decay coefficient)
  • done[t] = terminated[t] | truncated[t]

terminated[t] gates the one-step bootstrap inside δ[t]: set to 1 for true episode ends (s_{t+1} is a reset state with no meaningful value). truncated[t] stops trace propagation and injects the true continuation value bootstrap_values[env, t] = V(s_{t+1}^true) into δ[t].

bootstrap_values supplies the true continuation value V(s_{t+1}) wherever the stored values[t+1] is invalid. Two situations require this, under one rule: (a) truncated steps, where values[t+1] belongs to the next episode; and (b) the final column t=T-1, where values[t+1] lies past the buffer. In both cases set bootstrap_values to the true continuation value; leave it zero everywhere else.

The final column feeds delta[T-1] as the next-state value ONLY. The scan's additive boundary carry A[T] is always 0: an advantage carry represents trace mass from steps past the buffer, and there is none there. (β[T-1], the multiplicative decay coefficient, is unaffected by this -- only the additive A[T] term changes.) A single entry in bootstrap_values[:, -1] is sufficient -- set it to V(s_T) if the episode continues past the window, or 0 if it terminated at T-1.

Most users have no interior truncations and should use the last_value argument (shape [num_envs]) instead, which populates the boundary column automatically.

Parameters:

Name Type Description Default
rewards Tensor

Per-step rewards, [num_envs, seq_len], float32, CUDA.

 required
values Tensor

V(s_t), [num_envs, seq_len], float32, CUDA.

 required
terminateds Tensor

True termination flags (1.0=terminated), [num_envs, seq_len], float32, CUDA.

 required
truncateds Tensor | None

Time-limit truncation flags (1.0=truncated), [num_envs, seq_len], float32, CUDA. If None, terminateds is used for both gating roles (conservative: treats all boundaries as terminations).

 None
gamma float

Discount factor (default 0.99).

 0.99
lambda_ float

GAE trace parameter in [0, 1] (default 0.95).

 0.95
bootstrap_values Tensor | None

True continuation values V(s_{t+1}^true), [num_envs, seq_len], float32, CUDA. Set bootstrap_values[env, t] = V(s_{t+1}^true) at every truncated step and at t=T-1 if the window ends mid-episode. Zero elsewhere. If None, defaults to all zeros. Mutually exclusive with last_value.

 None
last_value Tensor | None

Convenience arg for the common case of no interior truncations: V(s_T) per environment, shape [num_envs], float32, CUDA. Populates bootstrap_values[:, -1] automatically. Mutually exclusive with bootstrap_values.

 None

Returns:

Name Type Description
advantages Tensor

A[t], shape [num_envs, seq_len], float32.
Source code in rl_triton/ops/gae.py 
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
def compute_gae(
    rewards: torch.Tensor,
    values: torch.Tensor,
    terminateds: torch.Tensor,
    truncateds: torch.Tensor | None = None,
    gamma: float = 0.99,
    lambda_: float = 0.95,
    bootstrap_values: torch.Tensor | None = None,
    last_value: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Compute Generalized Advantage Estimation via a backward associative scan.

    Recurrence:

    - A[t] = δ[t] + β[t] * A[t+1],  A[T] = 0
    - δ[t] = r[t] + γ·(1-terminated[t]) * v_next[t] - V(s_t)
    - β[t] = γ·λ·(1-done[t])       (scan decay coefficient)
    - done[t] = terminated[t] | truncated[t]

    `terminated[t]` gates the one-step bootstrap inside δ[t]: set to 1 for true
    episode ends (s_{t+1} is a reset state with no meaningful value).
    `truncated[t]` stops trace propagation and injects the true continuation
    value bootstrap_values[env, t] = V(s_{t+1}^true) into δ[t].

    bootstrap_values supplies the true continuation value V(s_{t+1}) wherever
    the stored values[t+1] is invalid.  Two situations require this, under one
    rule: (a) truncated steps, where values[t+1] belongs to the next episode;
    and (b) the final column t=T-1, where values[t+1] lies past the buffer.
    In both cases set bootstrap_values to the true continuation value; leave
    it zero everywhere else.

    The final column feeds delta[T-1] as the next-state value ONLY.  The scan's
    additive boundary carry A[T] is always 0: an advantage carry represents
    trace mass from steps past the buffer, and there is none there.  (β[T-1],
    the multiplicative decay coefficient, is unaffected by this -- only the
    additive A[T] term changes.)  A single entry in bootstrap_values[:, -1] is
    sufficient -- set it to V(s_T) if the episode continues past the window,
    or 0 if it terminated at T-1.

    Most users have no interior truncations and should use the `last_value`
    argument (shape [num_envs]) instead, which populates the boundary column
    automatically.

    Args:
        rewards:          Per-step rewards, [num_envs, seq_len], float32, CUDA.
        values:           V(s_t), [num_envs, seq_len], float32, CUDA.
        terminateds:      True termination flags (1.0=terminated),
                          [num_envs, seq_len], float32, CUDA.
        truncateds:       Time-limit truncation flags (1.0=truncated),
                          [num_envs, seq_len], float32, CUDA.
                          If None, terminateds is used for both gating roles
                          (conservative: treats all boundaries as terminations).
        gamma:            Discount factor (default 0.99).
        lambda_:          GAE trace parameter in [0, 1] (default 0.95).
        bootstrap_values: True continuation values V(s_{t+1}^true),
                          [num_envs, seq_len], float32, CUDA.
                          Set bootstrap_values[env, t] = V(s_{t+1}^true) at every
                          truncated step and at t=T-1 if the window ends mid-episode.
                          Zero elsewhere.  If None, defaults to all zeros.
                          Mutually exclusive with last_value.
        last_value:       Convenience arg for the common case of no interior
                          truncations: V(s_T) per environment, shape [num_envs],
                          float32, CUDA.  Populates bootstrap_values[:, -1]
                          automatically.  Mutually exclusive with bootstrap_values.

    Returns:
        advantages: A[t], shape [num_envs, seq_len], float32.
    """
    num_envs, seq_len = rewards.shape
    has_truncations   = truncateds is not None

    # Cheap structural checks -- always-on.
    for name, t in [("rewards", rewards), ("values", values), ("terminateds", terminateds)]:
        assert t.is_cuda,                f"{name} must be on CUDA"
        assert t.dtype == torch.float32, f"{name}: expected float32, got {t.dtype}"
        assert t.shape == rewards.shape, f"{name} shape {t.shape} != rewards shape {rewards.shape}"
    if has_truncations:
        assert truncateds.is_cuda,                "truncateds must be on CUDA"
        assert truncateds.dtype == torch.float32, "truncateds: expected float32"
        assert truncateds.shape == rewards.shape, \
            f"truncateds shape {truncateds.shape} != rewards shape {rewards.shape}"
    if last_value is not None:
        assert bootstrap_values is None, \
            "pass either last_value (shape [num_envs], convenience for the " \
            "window boundary) or bootstrap_values (shape [num_envs, seq_len], " \
            "full per-step control), not both."
        assert last_value.shape == (num_envs,), \
            f"last_value must have shape [{num_envs}], got {last_value.shape}"
        assert not has_truncations, \
            "last_value cannot be combined with truncateds; use bootstrap_values instead."
    if bootstrap_values is not None:
        assert bootstrap_values.is_cuda,                "bootstrap_values must be on CUDA"
        assert bootstrap_values.dtype == torch.float32, "bootstrap_values: expected float32"
        assert bootstrap_values.shape == rewards.shape, \
            f"bootstrap_values shape {bootstrap_values.shape} != rewards shape {rewards.shape}"

    # Expensive tensor scans -- correctness-warning path only (not in benchmark hot loop).
    if _CORRECTNESS_WARNINGS():
        if has_truncations:
            assert not (terminateds.bool() & truncateds.bool()).any(), \
                "terminated and truncated are mutually exclusive: a step cannot be both"
        if has_truncations and bootstrap_values is not None:
            interior = torch.ones_like(truncateds, dtype=torch.bool)
            interior[:, -1] = False
            stray = (bootstrap_values != 0) & (truncateds == 0) & interior
            assert not stray.any(), (
                "bootstrap_values must be zero at non-truncated interior steps. "
                "Nonzero entries there double-count via the additive v_next trick "
                "and corrupt delta. Populate bootstrap_values only at truncated "
                "steps and at the final column (window boundary)."
            )

    rewards     = rewards.contiguous()
    values      = values.contiguous()
    terminateds = terminateds.contiguous()
    if has_truncations:
        truncateds = truncateds.contiguous()

    if last_value is not None:
        scalar_bootstrap = last_value.contiguous()
        bootstrap_values = None
    elif bootstrap_values is not None:
        bootstrap_values = bootstrap_values.contiguous()
        scalar_bootstrap = bootstrap_values[:, -1].contiguous()
    else:
        scalar_bootstrap = None

    out = torch.empty_like(rewards)

    if seq_len <= _FLAT_MAX_SEQ_LEN:
        BLOCK_SIZE = triton.next_power_of_2(seq_len)
        num_warps  = _WARPS.get(BLOCK_SIZE, 16)
        num_stages = 2 if BLOCK_SIZE >= 2048 else 1

        if has_truncations:
            if bootstrap_values is None:
                bootstrap_values = torch.zeros_like(rewards)
            gae_fused_kernel[(num_envs,)](
                rewards, values, terminateds, truncateds,
                out, bootstrap_values,
                seq_len, rewards.stride(0),
                gamma=gamma, lambda_=lambda_,
                BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps, num_stages=num_stages,
                HAS_TRUNCATIONS=True, HAS_BOOTSTRAP=True,
            )
        else:
            has_bootstrap = scalar_bootstrap is not None
            gae_fused_kernel[(num_envs,)](
                rewards, values, terminateds, None,
                out, scalar_bootstrap,
                seq_len, rewards.stride(0),
                gamma=gamma, lambda_=lambda_,
                BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps, num_stages=num_stages,
                HAS_TRUNCATIONS=False, HAS_BOOTSTRAP=has_bootstrap,
            )
        return out

    # Chunked path for seq_len > 131072.
    if not has_truncations:
        truncateds = torch.zeros_like(terminateds)
    if bootstrap_values is None:
        bootstrap_values = torch.zeros_like(rewards)
        if scalar_bootstrap is not None:
            bootstrap_values[:, -1] = scalar_bootstrap

    not_terminated = 1.0 - terminateds
    not_done       = 1.0 - (terminateds + truncateds).clamp(max=1.0)
    next_values    = torch.empty_like(values)
    next_values[:, :-1] = values[:, 1:]
    next_values[:, -1]  = 0.0
    next_values = next_values * (1.0 - truncateds)
    next_values[:, -1]  = 0.0
    next_values = next_values + bootstrap_values

    deltas = rewards + gamma * not_terminated * next_values - values
    decays = gamma * lambda_ * not_done   # beta[t], the scan decay coefficient
    # Additive boundary carry A[T] = 0: bootstrap_values[:, -1] already entered
    # deltas[:, -1] via next_values above (weight 1). Seeding the scan's
    # boundary carry with it too would double-count it -- see kernels/gae.py's
    # module docstring. beta (decays) itself is unaffected by this fix.
    return _run_scan(deltas, decays)

rl_triton.ops.vtrace.compute_vtrace(log_pi_target, log_pi_behavior, values, rewards, terminateds, truncateds=None, gamma=0.99, rho_bar=1.0, c_bar=1.0, bootstrap_values=None, last_value=None)

Compute V-Trace targets and advantages via a backward associative scan.

V-Trace (Espeholt et al. 2018, IMPALA) corrects for off-policy data using clipped importance sampling ratios ρ and c.

Recurrence:

  • Δ[t] = δ[t] + β[t] * Δ[t+1], Δ[T] = 0
  • δ[t] = ρ[t] * (r[t] + γ * V(s_{t+1}) * (1 - terminated[t]) - V(s_t))
  • β[t] = γ * c[t] * (1 - done[t]), done[t] = terminated[t] | truncated[t] (scan decay coefficient)
  • ρ[t] = min(ρ_bar, π_target[t] / π_behavior[t])
  • c[t] = min(c_bar, π_target[t] / π_behavior[t])

Outputs:

  • vs[t] = Δ[t] + V(s_t) (critic targets)
  • A[t] = ρ[t] * (r[t] + γ * vs[t+1] * (1 - terminated[t]) - V(s_t)) (actor advantages)

bootstrap_values supplies the true continuation value V(s_{t+1}) wherever the stored values[t+1] is invalid. Two situations require this, under one rule: (a) truncated steps, where values[t+1] belongs to the next episode; and (b) the final column t=T-1, where values[t+1] lies past the buffer. In both cases set bootstrap_values to the true continuation value; leave it zero everywhere else.

The final column feeds delta[T-1] as the next-state value ONLY. The scan's additive boundary carry Δ[T] is always 0: an advantage carry represents trace mass from steps past the buffer, and there is none there. (β[T-1], the multiplicative decay coefficient above, is unaffected by this -- only the additive Δ[T] term changes.) A single entry in bootstrap_values[:, -1] is sufficient -- set it to V(s_T) if the episode continues past the window, or 0 if it terminated at T-1. (The same bootstrap_values[:, -1] is also used directly, not as a carry, when computing next_vtrace_targets[:, -1] for the advantage formula -- that use is single-counted and correct as-is.)

Most users have no interior truncations and should use the last_value argument (shape [num_envs]) instead, which populates the boundary column automatically.

terminated[t] gates the one-step bootstrap in δ[t]: set to 1 at true episode ends. truncated[t] stops trace propagation and injects bootstrap_values[env, t] = V(s_{t+1}^true), correcting for stale reset observations in the values buffer under Gymnasium next-step autoreset.

Parameters:

Name Type Description Default
log_pi_target Tensor

Log probabilities under target policy, [num_envs, seq_len], float32, CUDA.

 required
log_pi_behavior Tensor

Log probabilities under behavior policy, [num_envs, seq_len], float32, CUDA.

 required
values Tensor

V(s_t), [num_envs, seq_len], float32, CUDA.

 required
rewards Tensor

Per-step rewards, [num_envs, seq_len], float32, CUDA.

 required
terminateds Tensor

True termination flags (1.0=terminated), [num_envs, seq_len], float32, CUDA.

 required
truncateds Tensor | None

Time-limit truncation flags (1.0=truncated), [num_envs, seq_len], float32, CUDA. If None, terminateds is used for both gating roles.

 None
gamma float

Discount factor (default 0.99).

 0.99
rho_bar float

IS ratio clip for δ (default 1.0).

 1.0
c_bar float

IS ratio clip for decay (default 1.0).

 1.0
bootstrap_values Tensor | None

True continuation values V(s_{t+1}^true), [num_envs, seq_len], float32, CUDA. Set bootstrap_values[env, t] = V(s_{t+1}^true) at every truncated step and at t=T-1 if the window ends mid-episode. Zero elsewhere. If None, defaults to all zeros. Mutually exclusive with last_value.

 None
last_value Tensor | None

Convenience arg for the common case of no interior truncations: V(s_T) per environment, shape [num_envs], float32, CUDA. Populates bootstrap_values[:, -1] automatically. Mutually exclusive with bootstrap_values.

 None

Returns:

Name Type Description
vtrace_targets Tensor

vs[t], shape [num_envs, seq_len], float32.
vtrace_advantages Tensor

A[t], shape [num_envs, seq_len], float32.
Source code in rl_triton/ops/vtrace.py 
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
def compute_vtrace(
    log_pi_target: torch.Tensor,
    log_pi_behavior: torch.Tensor,
    values: torch.Tensor,
    rewards: torch.Tensor,
    terminateds: torch.Tensor,
    truncateds: torch.Tensor | None = None,
    gamma: float = 0.99,
    rho_bar: float = 1.0,
    c_bar: float = 1.0,
    bootstrap_values: torch.Tensor | None = None,
    last_value: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Compute V-Trace targets and advantages via a backward associative scan.

    V-Trace (Espeholt et al. 2018, IMPALA) corrects for off-policy data using
    clipped importance sampling ratios ρ and c.

    Recurrence:

    - Δ[t] = δ[t] + β[t] * Δ[t+1],   Δ[T] = 0
    - δ[t] = ρ[t] * (r[t] + γ * V(s_{t+1}) * (1 - terminated[t]) - V(s_t))
    - β[t] = γ * c[t] * (1 - done[t]),  done[t] = terminated[t] | truncated[t]
             (scan decay coefficient)
    - ρ[t] = min(ρ_bar, π_target[t] / π_behavior[t])
    - c[t] = min(c_bar, π_target[t] / π_behavior[t])

    Outputs:

    - vs[t] = Δ[t] + V(s_t)   (critic targets)
    - A[t]  = ρ[t] * (r[t] + γ * vs[t+1] * (1 - terminated[t]) - V(s_t))   (actor advantages)

    bootstrap_values supplies the true continuation value V(s_{t+1}) wherever
    the stored values[t+1] is invalid.  Two situations require this, under one
    rule: (a) truncated steps, where values[t+1] belongs to the next episode;
    and (b) the final column t=T-1, where values[t+1] lies past the buffer.
    In both cases set bootstrap_values to the true continuation value; leave
    it zero everywhere else.

    The final column feeds delta[T-1] as the next-state value ONLY.  The
    scan's additive boundary carry Δ[T] is always 0: an advantage carry
    represents trace mass from steps past the buffer, and there is none
    there.  (β[T-1], the multiplicative decay coefficient above, is
    unaffected by this -- only the additive Δ[T] term changes.)  A single
    entry in bootstrap_values[:, -1] is sufficient -- set it to V(s_T) if the
    episode continues past the window, or 0 if it terminated at T-1.  (The
    same bootstrap_values[:, -1] is also used directly, not as a carry, when
    computing next_vtrace_targets[:, -1] for the advantage formula -- that use
    is single-counted and correct as-is.)

    Most users have no interior truncations and should use the `last_value`
    argument (shape [num_envs]) instead, which populates the boundary column
    automatically.

    `terminated[t]` gates the one-step bootstrap in δ[t]: set to 1 at true episode
    ends.  `truncated[t]` stops trace propagation and injects bootstrap_values[env, t]
    = V(s_{t+1}^true), correcting for stale reset observations in the values buffer
    under Gymnasium next-step autoreset.

    Args:
        log_pi_target:    Log probabilities under target policy, [num_envs, seq_len], float32, CUDA.
        log_pi_behavior:  Log probabilities under behavior policy, [num_envs, seq_len], float32, CUDA.
        values:           V(s_t), [num_envs, seq_len], float32, CUDA.
        rewards:          Per-step rewards, [num_envs, seq_len], float32, CUDA.
        terminateds:      True termination flags (1.0=terminated),
                          [num_envs, seq_len], float32, CUDA.
        truncateds:       Time-limit truncation flags (1.0=truncated),
                          [num_envs, seq_len], float32, CUDA.
                          If None, terminateds is used for both gating roles.
        gamma:            Discount factor (default 0.99).
        rho_bar:          IS ratio clip for δ (default 1.0).
        c_bar:            IS ratio clip for decay (default 1.0).
        bootstrap_values: True continuation values V(s_{t+1}^true),
                          [num_envs, seq_len], float32, CUDA.
                          Set bootstrap_values[env, t] = V(s_{t+1}^true) at every
                          truncated step and at t=T-1 if the window ends mid-episode.
                          Zero elsewhere.  If None, defaults to all zeros.
                          Mutually exclusive with last_value.
        last_value:       Convenience arg for the common case of no interior
                          truncations: V(s_T) per environment, shape [num_envs],
                          float32, CUDA.  Populates bootstrap_values[:, -1]
                          automatically.  Mutually exclusive with bootstrap_values.

    Returns:
        vtrace_targets:    vs[t], shape [num_envs, seq_len], float32.
        vtrace_advantages: A[t],  shape [num_envs, seq_len], float32.
    """
    num_envs, seq_len = rewards.shape
    has_truncations   = truncateds is not None

    if _CORRECTNESS_WARNINGS():
        for name, t in [
            ("log_pi_target",   log_pi_target),
            ("log_pi_behavior", log_pi_behavior),
            ("values",          values),
            ("rewards",         rewards),
            ("terminateds",     terminateds),
        ]:
            assert t.is_cuda,                f"{name} must be on CUDA"
            assert t.dtype == torch.float32, f"{name}: expected float32, got {t.dtype}"
            assert t.shape == rewards.shape, f"{name} shape {t.shape} != rewards shape {rewards.shape}"
        if has_truncations:
            assert truncateds.is_cuda,                "truncateds must be on CUDA"
            assert truncateds.dtype == torch.float32, "truncateds: expected float32"
            assert truncateds.shape == rewards.shape, \
                f"truncateds shape {truncateds.shape} != rewards shape {rewards.shape}"
            assert not (terminateds.bool() & truncateds.bool()).any(), \
                "terminated and truncated are mutually exclusive: a step cannot be both"
        if last_value is not None:
            assert bootstrap_values is None, \
                "pass either last_value (shape [num_envs], convenience for the " \
                "window boundary) or bootstrap_values (shape [num_envs, seq_len], " \
                "full per-step control), not both."
            assert last_value.shape == (num_envs,), \
                f"last_value must have shape [{num_envs}], got {last_value.shape}"
        if bootstrap_values is not None:
            assert bootstrap_values.is_cuda,                "bootstrap_values must be on CUDA"
            assert bootstrap_values.dtype == torch.float32, "bootstrap_values: expected float32"
            assert bootstrap_values.shape == rewards.shape, \
                f"bootstrap_values shape {bootstrap_values.shape} != rewards shape {rewards.shape}"
        if has_truncations and bootstrap_values is not None:
            interior = torch.ones_like(truncateds, dtype=torch.bool)
            interior[:, -1] = False
            stray = (bootstrap_values != 0) & (truncateds == 0) & interior
            assert not stray.any(), (
                "bootstrap_values must be zero at non-truncated interior steps. "
                "Nonzero entries there double-count via the additive v_next trick "
                "and corrupt delta. Populate bootstrap_values only at truncated "
                "steps and at the final column (window boundary)."
            )

    log_pi_target   = log_pi_target.contiguous()
    log_pi_behavior = log_pi_behavior.contiguous()
    values          = values.contiguous()
    rewards         = rewards.contiguous()
    terminateds     = terminateds.contiguous()
    if has_truncations:
        truncateds = truncateds.contiguous()

    # Fused kernel for seq_len <= 131072.
    # Pass truncateds=None for the no-truncation path so compute_vtrace_fused
    # dispatches HAS_TRUNCATIONS=False -- no zero-tensor allocations for truncateds
    # or 2D bootstrap_values in that path.
    if seq_len <= _FLAT_MAX_SEQ_LEN:
        return compute_vtrace_fused(
            log_pi_target, log_pi_behavior,
            values, rewards, terminateds,
            truncateds=truncateds if has_truncations else None,
            gamma=gamma, rho_bar=rho_bar, c_bar=c_bar,
            bootstrap_values=bootstrap_values,
            last_value=last_value,
        )

    # Chunked path for seq_len > 131072.  Materialize None inputs here only.
    if not has_truncations:
        truncateds = torch.zeros_like(terminateds)
    if last_value is not None:
        bootstrap_values = torch.zeros_like(rewards)
        bootstrap_values[:, -1] = last_value
    elif bootstrap_values is None:
        bootstrap_values = torch.zeros_like(rewards)

    not_terminated = 1.0 - terminateds
    not_done       = 1.0 - (terminateds + truncateds).clamp(max=1.0)
    next_values    = torch.empty_like(values)
    next_values[:, :-1] = values[:, 1:]
    next_values[:, -1]  = 0.0
    next_values = next_values * (1.0 - truncateds)
    next_values[:, -1]  = 0.0
    next_values = next_values + bootstrap_values

    is_ratios      = torch.exp(log_pi_target - log_pi_behavior)
    rho            = torch.clamp(is_ratios, max=rho_bar)
    c              = torch.clamp(is_ratios, max=c_bar)
    u = rho * (rewards + gamma * next_values * not_terminated - values)
    v = gamma * c * not_done   # beta[t], the scan decay coefficient

    # Additive boundary carry Delta[T] = 0: bootstrap_values[:, -1] already
    # entered u[:, -1] via next_values above (weight 1). Seeding the scan's
    # boundary carry with it too would double-count it -- see kernels/gae.py's
    # module docstring for the same bug class in GAE. beta (v) itself is
    # unaffected by this fix.
    value_deltas   = _run_scan(u, v)
    vtrace_targets = value_deltas + values

    next_vtrace_targets = torch.empty_like(vtrace_targets)
    next_vtrace_targets[:, :-1] = vtrace_targets[:, 1:]
    next_vtrace_targets[:, -1]  = 0.0
    next_vtrace_targets = next_vtrace_targets * (1.0 - truncateds)
    next_vtrace_targets[:, -1]  = 0.0
    next_vtrace_targets = next_vtrace_targets + bootstrap_values

    vtrace_advantages = rho * (rewards + gamma * next_vtrace_targets * not_terminated - values)
    return vtrace_targets, vtrace_advantages

rl_triton.ops.retrace.compute_retrace(action_probs_target, action_probs_behavior, q_values, next_q_values_all, actions, rewards, terminateds, truncateds, gamma, lambda_=1.0, c_bar=1.0, rho_bar=1.0)

Compute Retrace(λ) Q-value targets and advantages via a backward associative scan.

Retrace(λ) (Munos et al. 2016) corrects off-policy Q-value estimates using truncated IS traces applied only to the decay factor, not the TD error. Discrete actions only -- use compute_vtrace for continuous action spaces.

Recurrence:

  • Δ[t] = δ[t] + decay[t] * Δ[t+1], Δ[T] = 0
  • δ[t] = r[t] + γ · E_π[Q(s_{t+1},a)] · (1-terminated[t]) - Q(s_t,a_t)
  • decay[t] = γ · c[t+1] · (1-done[t]), done[t] = terminated[t] | truncated[t]
  • c[t] = λ · min(c_bar, π(a_t|s_t) / μ(a_t|s_t))

Q-value targets: Q_ret[t] = Q(s_t, a_t) + Δ[t] Advantages: A[t] = ρ[t] · (r[t] + γ · Q_ret[t+1] · (1-terminated[t]) - Q(s_t,a_t)) ρ[t] = min(rho_bar, π(a_t|s_t) / μ(a_t|s_t))

The Q-bootstrap γ·E_π[Q(s_{t+1},·)] is folded into δ[t] via next_q_values_all, so no separate bootstrap_values argument is needed. terminated gates the one-step bootstrap in δ[t]; terminated | truncated gates trace decay in β[t].

Parameters:

Name Type Description Default
action_probs_target Tensor

Target policy probabilities over all actions, [num_envs, seq_len, num_actions], float32, CUDA.

 required
action_probs_behavior Tensor

Behavior policy probability of the taken action, [num_envs, seq_len], float32, CUDA.

 required
q_values Tensor

Q(s_t, a_t) for the taken action, [num_envs, seq_len], float32, CUDA.

 required
next_q_values_all Tensor

Q(s_{t+1}, a) for all actions, [num_envs, seq_len, num_actions], float32, CUDA.

 required
actions Tensor

Indices of the taken action, [num_envs, seq_len], int64, CUDA.

 required
rewards Tensor

Per-step rewards, [num_envs, seq_len], float32, CUDA.

 required
terminateds Tensor

True termination flags (1.0 = terminated). Zeros the bootstrap γ·E_π[Q(s_{t+1},·)] in δ[t]. [num_envs, seq_len], float32, CUDA.

 required
truncateds Tensor

Time-limit truncation flags (1.0 = truncated). Keeps the bootstrap in δ[t] but severs the trace. [num_envs, seq_len], float32, CUDA.

 required
gamma float

Discount factor.

 required
lambda_ float

Trace decay parameter (default 1.0).

 1.0
c_bar float

IS ratio clip for trace weights (default 1.0).

 1.0
rho_bar float

IS ratio clip for advantage scaling (default 1.0).

 1.0

Returns:

Name Type Description
retrace_targets Tensor

Q_ret[t], shape [num_envs, seq_len], float32.
advantages Tensor

A[t], shape [num_envs, seq_len], float32.
Source code in rl_triton/ops/retrace.py 
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
def compute_retrace(
    action_probs_target: torch.Tensor,
    action_probs_behavior: torch.Tensor,
    q_values: torch.Tensor,
    next_q_values_all: torch.Tensor,
    actions: torch.Tensor,
    rewards: torch.Tensor,
    terminateds: torch.Tensor,
    truncateds: torch.Tensor,
    gamma: float,
    lambda_: float = 1.0,
    c_bar: float = 1.0,
    rho_bar: float = 1.0,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Compute Retrace(λ) Q-value targets and advantages via a backward associative scan.

    Retrace(λ) (Munos et al. 2016) corrects off-policy Q-value estimates using
    truncated IS traces applied only to the decay factor, not the TD error.
    Discrete actions only -- use compute_vtrace for continuous action spaces.

    Recurrence:

    - Δ[t] = δ[t] + decay[t] * Δ[t+1],   Δ[T] = 0
    - δ[t]     = r[t] + γ · E_π[Q(s_{t+1},a)] · (1-terminated[t]) - Q(s_t,a_t)
    - decay[t] = γ · c[t+1] · (1-done[t]),  done[t] = terminated[t] | truncated[t]
    - c[t]     = λ · min(c_bar, π(a_t|s_t) / μ(a_t|s_t))

    Q-value targets: Q_ret[t] = Q(s_t, a_t) + Δ[t]
    Advantages:      A[t]     = ρ[t] · (r[t] + γ · Q_ret[t+1] · (1-terminated[t]) - Q(s_t,a_t))
                     ρ[t]     = min(rho_bar, π(a_t|s_t) / μ(a_t|s_t))

    The Q-bootstrap γ·E_π[Q(s_{t+1},·)] is folded into δ[t] via
    `next_q_values_all`, so no separate bootstrap_values argument is needed.
    `terminated` gates the one-step bootstrap in δ[t]; `terminated | truncated`
    gates trace decay in β[t].

    Args:
        action_probs_target:   Target policy probabilities over all actions,
                               [num_envs, seq_len, num_actions], float32, CUDA.
        action_probs_behavior: Behavior policy probability of the taken action,
                               [num_envs, seq_len], float32, CUDA.
        q_values:              Q(s_t, a_t) for the taken action,
                               [num_envs, seq_len], float32, CUDA.
        next_q_values_all:     Q(s_{t+1}, a) for all actions,
                               [num_envs, seq_len, num_actions], float32, CUDA.
        actions:               Indices of the taken action,
                               [num_envs, seq_len], int64, CUDA.
        rewards:               Per-step rewards, [num_envs, seq_len], float32, CUDA.
        terminateds:           True termination flags (1.0 = terminated).
                               Zeros the bootstrap γ·E_π[Q(s_{t+1},·)] in δ[t].
                               [num_envs, seq_len], float32, CUDA.
        truncateds:            Time-limit truncation flags (1.0 = truncated).
                               Keeps the bootstrap in δ[t] but severs the trace.
                               [num_envs, seq_len], float32, CUDA.
        gamma:                 Discount factor.
        lambda_:               Trace decay parameter (default 1.0).
        c_bar:                 IS ratio clip for trace weights (default 1.0).
        rho_bar:               IS ratio clip for advantage scaling (default 1.0).

    Returns:
        retrace_targets: Q_ret[t], shape [num_envs, seq_len], float32.
        advantages:      A[t],     shape [num_envs, seq_len], float32.
    """
    # Cheap structural checks -- always-on.
    for name, t in [
        ("action_probs_behavior", action_probs_behavior),
        ("q_values",              q_values),
        ("rewards",               rewards),
        ("terminateds",           terminateds),
        ("truncateds",            truncateds),
    ]:
        assert t.is_cuda,                 f"{name} must be on CUDA"
        assert t.dtype == torch.float32,  f"{name}: expected float32, got {t.dtype}"
        assert t.shape == rewards.shape,  f"{name} shape {t.shape} != rewards shape {rewards.shape}"

    assert action_probs_target.is_cuda,                "action_probs_target must be on CUDA"
    assert action_probs_target.dtype == torch.float32, "action_probs_target: expected float32"
    assert action_probs_target.shape[:2] == rewards.shape, (
        f"action_probs_target shape {action_probs_target.shape} incompatible with rewards {rewards.shape}"
    )
    assert next_q_values_all.is_cuda,                "next_q_values_all must be on CUDA"
    assert next_q_values_all.dtype == torch.float32, "next_q_values_all: expected float32"
    assert next_q_values_all.shape == action_probs_target.shape, (
        f"next_q_values_all {next_q_values_all.shape} != action_probs_target {action_probs_target.shape}"
    )
    assert actions.is_cuda,              "actions must be on CUDA"
    assert actions.dtype == torch.int64, f"actions: expected int64, got {actions.dtype}"
    assert actions.shape == rewards.shape, \
        f"actions shape {actions.shape} != rewards shape {rewards.shape}"

    # Expensive tensor scan -- correctness-warning path only.
    if _CORRECTNESS_WARNINGS():
        assert not (terminateds.bool() & truncateds.bool()).any(), \
            "terminated and truncated are mutually exclusive: a step cannot be both"

    action_probs_target   = action_probs_target.contiguous()
    action_probs_behavior = action_probs_behavior.contiguous()
    q_values              = q_values.contiguous()
    next_q_values_all     = next_q_values_all.contiguous()
    actions               = actions.contiguous()
    rewards               = rewards.contiguous()
    terminateds           = terminateds.contiguous()
    truncateds            = truncateds.contiguous()

    num_envs, seq_len = rewards.shape

    # Fused Triton kernel for seq_len <= _TRITON_SEQ_LEN_CEILING (confirmed win);
    # generic materialize-u/v + _run_scan path above that (fused kernel is a
    # confirmed loss there -- see _TRITON_SEQ_LEN_CEILING comment; this reroute
    # is a smaller loss too, not a win -- see the same comment).
    # done[t] = terminated[t] | truncated[t] is computed in-kernel from the two
    # raw flags for the fused path -- no separate PyTorch combine here.
    if seq_len <= _TRITON_SEQ_LEN_CEILING:
        return compute_retrace_fused(
            action_probs_target, action_probs_behavior,
            q_values, next_q_values_all, actions,
            rewards, truncateds, terminateds,
            gamma=gamma, lambda_=lambda_, c_bar=c_bar, rho_bar=rho_bar,
        )

    # seq_len > _TRITON_SEQ_LEN_CEILING: materialize u/v with plain PyTorch ops
    # (bandwidth-bound, no in-kernel 3D reads) and hand off to the generic
    # scan kernel via _run_scan, which internally picks the flat associative
    # scan for seq_len <= 131072 or the chunked kernel above that.
    dones = (terminateds + truncateds).clamp(max=1.0)
    expected_next_q = (action_probs_target * next_q_values_all).sum(dim=-1)
    u = rewards + gamma * expected_next_q * (1.0 - terminateds) - q_values

    pi_a   = action_probs_target.gather(-1, actions.unsqueeze(-1)).squeeze(-1)
    c      = lambda_ * torch.clamp(pi_a / action_probs_behavior, max=c_bar)
    c_next = torch.empty_like(c)
    c_next[:, :-1] = c[:, 1:]
    c_next[:, -1]  = 0.0
    v = gamma * c_next * (1.0 - dones)

    retrace_targets = _run_scan(u, v) + q_values

    rho = torch.clamp(pi_a / action_probs_behavior, max=rho_bar)
    next_q_ret = torch.empty_like(retrace_targets)
    next_q_ret[:, :-1] = retrace_targets[:, 1:]
    next_q_ret[:, -1]  = 0.0
    advantages = rho * (rewards + gamma * next_q_ret * (1.0 - terminateds) - q_values)

    return retrace_targets, advantages

rl_triton.ops.returns.compute_lambda_returns(rewards, next_values, terminateds, truncateds=None, gamma=0.99, lambda_=0.95, bootstrap_values=None)

Compute TD(λ) targets (λ-returns) via a backward associative scan.

Recurrence:

  • G[t] = r[t] + γ(1-done[t]) * [(1-λ)V(s_{t+1}) + λG[t+1]] + γtruncated[t]*bootstrap[t]
  • done[t] = terminated[t] | truncated[t]

Special cases: λ=0 reduces to one-step TD; λ=1 reduces to discounted returns.

terminated[t] stops both the value bootstrap and trace propagation. truncated[t] stops trace propagation and injects bootstrap_values[env, t] = V(s_{t+1}^true) as the continuation value. The caller must zero next_values[env, t] at truncated steps to avoid double-counting.

Parameters:

Name Type Description Default
rewards Tensor

Per-step rewards, [num_envs, seq_len], float32, CUDA.

 required
next_values Tensor

V(s_{t+1}), [num_envs, seq_len], float32, CUDA. Must be zeroed by caller at truncated steps.

 required
terminateds Tensor

True termination flags (1.0=terminated), [num_envs, seq_len], float32, CUDA.

 required
truncateds Tensor | None

Time-limit truncation flags (1.0=truncated), [num_envs, seq_len], float32, CUDA. If None, terminateds is used for both gating roles.

 None
gamma float

Discount factor (default 0.99).

 0.99
lambda_ float

Trace parameter in [0, 1] (default 0.95).

 0.95
bootstrap_values Tensor | None

True continuation values V(s_{t+1}^true), [num_envs, seq_len], float32, CUDA. Nonzero at truncated steps and at t=T-1 when the window ends mid-episode; zero elsewhere. If None, defaults to all zeros.

 None

Returns:

Name Type Description
lambda_returns Tensor

G[t], shape [num_envs, seq_len], float32.
Source code in rl_triton/ops/returns.py 
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
def compute_lambda_returns(
    rewards: torch.Tensor,
    next_values: torch.Tensor,
    terminateds: torch.Tensor,
    truncateds: torch.Tensor | None = None,
    gamma: float = 0.99,
    lambda_: float = 0.95,
    bootstrap_values: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Compute TD(λ) targets (λ-returns) via a backward associative scan.

    Recurrence:

    - G[t] = r[t] + γ*(1-done[t]) * [(1-λ)*V(s_{t+1}) + λ*G[t+1]]
             + γ*truncated[t]*bootstrap[t]
    - done[t] = terminated[t] | truncated[t]

    Special cases: λ=0 reduces to one-step TD; λ=1 reduces to discounted returns.

    `terminated[t]` stops both the value bootstrap and trace propagation.
    `truncated[t]` stops trace propagation and injects bootstrap_values[env, t]
    = V(s_{t+1}^true) as the continuation value.  The caller must zero
    next_values[env, t] at truncated steps to avoid double-counting.

    Args:
        rewards:          Per-step rewards, [num_envs, seq_len], float32, CUDA.
        next_values:      V(s_{t+1}), [num_envs, seq_len], float32, CUDA.
                          Must be zeroed by caller at truncated steps.
        terminateds:      True termination flags (1.0=terminated),
                          [num_envs, seq_len], float32, CUDA.
        truncateds:       Time-limit truncation flags (1.0=truncated),
                          [num_envs, seq_len], float32, CUDA.
                          If None, terminateds is used for both gating roles.
        gamma:            Discount factor (default 0.99).
        lambda_:          Trace parameter in [0, 1] (default 0.95).
        bootstrap_values: True continuation values V(s_{t+1}^true),
                          [num_envs, seq_len], float32, CUDA.
                          Nonzero at truncated steps and at t=T-1 when the window
                          ends mid-episode; zero elsewhere.
                          If None, defaults to all zeros.

    Returns:
        lambda_returns: G[t], shape [num_envs, seq_len], float32.
    """
    num_envs, seq_len = rewards.shape
    has_truncations   = truncateds is not None

    # Cheap structural checks -- always-on.
    for name, t in [("rewards", rewards), ("next_values", next_values), ("terminateds", terminateds)]:
        assert t.is_cuda,                f"{name} must be on CUDA"
        assert t.dtype == torch.float32, f"{name}: expected float32, got {t.dtype}"
        assert t.shape == rewards.shape, f"{name} shape {t.shape} != rewards shape {rewards.shape}"
    if has_truncations:
        assert truncateds.is_cuda,                "truncateds must be on CUDA"
        assert truncateds.dtype == torch.float32, "truncateds: expected float32"
        assert truncateds.shape == rewards.shape, \
            f"truncateds shape {truncateds.shape} != rewards shape {rewards.shape}"
    if bootstrap_values is not None:
        assert bootstrap_values.is_cuda,                "bootstrap_values must be on CUDA"
        assert bootstrap_values.dtype == torch.float32, "bootstrap_values: expected float32"
        assert bootstrap_values.shape == rewards.shape, \
            f"bootstrap_values shape {bootstrap_values.shape} != rewards shape {rewards.shape}"

    # Expensive tensor scans -- correctness-warning path only.
    if _CORRECTNESS_WARNINGS():
        if has_truncations:
            assert not (terminateds.bool() & truncateds.bool()).any(), \
                "terminated and truncated are mutually exclusive: a step cannot be both"

    rewards     = rewards.contiguous()
    next_values = next_values.contiguous()
    terminateds = terminateds.contiguous()
    if has_truncations:
        truncateds = truncateds.contiguous()
    if bootstrap_values is not None:
        bootstrap_values = bootstrap_values.contiguous()

    out = torch.empty_like(rewards)

    if seq_len <= _FLAT_MAX_SEQ_LEN:
        BLOCK_SIZE = triton.next_power_of_2(seq_len)
        num_warps  = _WARPS_LAMBDA.get(BLOCK_SIZE, 16)
        num_stages = 2 if BLOCK_SIZE >= 2048 else 1
        if has_truncations:
            if bootstrap_values is None:
                bootstrap_values = torch.zeros_like(rewards)
            lambda_returns_fused_kernel[(num_envs,)](
                rewards, next_values, terminateds, truncateds,
                out, bootstrap_values,
                seq_len, rewards.stride(0),
                gamma=gamma, lambda_=lambda_,
                BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps, num_stages=num_stages,
                HAS_TRUNCATIONS=True, HAS_BOOTSTRAP=True,
            )
        else:
            scalar_bootstrap = bootstrap_values[:, -1].contiguous() \
                               if bootstrap_values is not None else None
            has_bootstrap = scalar_bootstrap is not None
            lambda_returns_fused_kernel[(num_envs,)](
                rewards, next_values, terminateds, None,
                out, scalar_bootstrap,
                seq_len, rewards.stride(0),
                gamma=gamma, lambda_=lambda_,
                BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps, num_stages=num_stages,
                HAS_TRUNCATIONS=False, HAS_BOOTSTRAP=has_bootstrap,
            )
        return out

    # Chunked fallback for seq_len > 131072.
    if not has_truncations:
        truncateds = torch.zeros_like(terminateds)
    if bootstrap_values is None:
        bootstrap_values = torch.zeros_like(rewards)
    not_done = 1.0 - (terminateds + truncateds).clamp(max=1.0)
    carry    = bootstrap_values[:, -1]
    u = rewards + gamma * (1.0 - lambda_) * not_done * next_values \
        + gamma * truncateds * bootstrap_values
    v = gamma * lambda_ * not_done
    return _run_scan(u, v, carry)

rl_triton.ops.returns.compute_discounted_returns(rewards, terminateds, truncateds=None, gamma=0.99, bootstrap_values=None)

Compute discounted returns (reward-to-go) via a backward associative scan.

Recurrence:

  • G[t] = r[t] + γ(1-done[t])G[t+1] + γtruncated[t]bootstrap[t]
  • done[t] = terminated[t] | truncated[t]

terminated[t] stops the return propagation (episode ends, G[t+1]=0). truncated[t] also stops propagation but injects bootstrap_values[env, t] = V(s_{t+1}^true) as the continuation value instead of zero.

Parameters:

Name Type Description Default
rewards Tensor

Per-step rewards, [num_envs, seq_len], float32, CUDA.

 required
terminateds Tensor

True termination flags (1.0=terminated), [num_envs, seq_len], float32, CUDA.

 required
truncateds Tensor | None

Time-limit truncation flags (1.0=truncated), [num_envs, seq_len], float32, CUDA. If None, terminateds is used for both gating roles.

 None
gamma float

Discount factor (default 0.99).

 0.99
bootstrap_values Tensor | None

True continuation values V(s_{t+1}^true), [num_envs, seq_len], float32, CUDA. Nonzero at truncated steps and at t=T-1 when the window ends mid-episode; zero elsewhere. If None, defaults to all zeros.

 None

Returns:

Name Type Description
returns Tensor

G[t], shape [num_envs, seq_len], float32.
Source code in rl_triton/ops/returns.py 
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
def compute_discounted_returns(
    rewards: torch.Tensor,
    terminateds: torch.Tensor,
    truncateds: torch.Tensor | None = None,
    gamma: float = 0.99,
    bootstrap_values: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Compute discounted returns (reward-to-go) via a backward associative scan.

    Recurrence:

    - G[t] = r[t] + γ*(1-done[t])*G[t+1] + γ*truncated[t]*bootstrap[t]
    - done[t] = terminated[t] | truncated[t]

    `terminated[t]` stops the return propagation (episode ends, G[t+1]=0).
    `truncated[t]` also stops propagation but injects bootstrap_values[env, t]
    = V(s_{t+1}^true) as the continuation value instead of zero.

    Args:
        rewards:          Per-step rewards, [num_envs, seq_len], float32, CUDA.
        terminateds:      True termination flags (1.0=terminated),
                          [num_envs, seq_len], float32, CUDA.
        truncateds:       Time-limit truncation flags (1.0=truncated),
                          [num_envs, seq_len], float32, CUDA.
                          If None, terminateds is used for both gating roles.
        gamma:            Discount factor (default 0.99).
        bootstrap_values: True continuation values V(s_{t+1}^true),
                          [num_envs, seq_len], float32, CUDA.
                          Nonzero at truncated steps and at t=T-1 when the window
                          ends mid-episode; zero elsewhere.
                          If None, defaults to all zeros.

    Returns:
        returns: G[t], shape [num_envs, seq_len], float32.
    """
    num_envs, seq_len = rewards.shape
    has_truncations   = truncateds is not None

    # Cheap structural checks -- always-on.
    for name, t in [("rewards", rewards), ("terminateds", terminateds)]:
        assert t.is_cuda,                f"{name} must be on CUDA"
        assert t.dtype == torch.float32, f"{name}: expected float32, got {t.dtype}"
        assert t.shape == rewards.shape, f"{name} shape {t.shape} != rewards shape {rewards.shape}"
    if has_truncations:
        assert truncateds.is_cuda,                "truncateds must be on CUDA"
        assert truncateds.dtype == torch.float32, "truncateds: expected float32"
        assert truncateds.shape == rewards.shape, \
            f"truncateds shape {truncateds.shape} != rewards shape {rewards.shape}"
    if bootstrap_values is not None:
        assert bootstrap_values.is_cuda,                "bootstrap_values must be on CUDA"
        assert bootstrap_values.dtype == torch.float32, "bootstrap_values: expected float32"
        assert bootstrap_values.shape == rewards.shape, \
            f"bootstrap_values shape {bootstrap_values.shape} != rewards shape {rewards.shape}"

    # Expensive tensor scans -- correctness-warning path only.
    if _CORRECTNESS_WARNINGS():
        if has_truncations:
            assert not (terminateds.bool() & truncateds.bool()).any(), \
                "terminated and truncated are mutually exclusive: a step cannot be both"

    rewards     = rewards.contiguous()
    terminateds = terminateds.contiguous()
    if has_truncations:
        truncateds = truncateds.contiguous()
    if bootstrap_values is not None:
        bootstrap_values = bootstrap_values.contiguous()

    out = torch.empty_like(rewards)

    if seq_len <= _FLAT_MAX_SEQ_LEN:
        BLOCK_SIZE = triton.next_power_of_2(seq_len)
        num_warps  = _WARPS_LIGHT.get(BLOCK_SIZE, 16)
        num_stages = 2 if BLOCK_SIZE >= 1024 else 1
        if has_truncations:
            if bootstrap_values is None:
                bootstrap_values = torch.zeros_like(rewards)
            discounted_returns_fused_kernel[(num_envs,)](
                rewards, terminateds, truncateds,
                out, bootstrap_values,
                seq_len, rewards.stride(0),
                gamma=gamma,
                BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps, num_stages=num_stages,
                HAS_TRUNCATIONS=True, HAS_BOOTSTRAP=True,
            )
        else:
            scalar_bootstrap = bootstrap_values[:, -1].contiguous() \
                               if bootstrap_values is not None else None
            has_bootstrap = scalar_bootstrap is not None
            discounted_returns_fused_kernel[(num_envs,)](
                rewards, terminateds, None,
                out, scalar_bootstrap,
                seq_len, rewards.stride(0),
                gamma=gamma,
                BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps, num_stages=num_stages,
                HAS_TRUNCATIONS=False, HAS_BOOTSTRAP=has_bootstrap,
            )
        return out

    # Chunked fallback for seq_len > 131072.
    if not has_truncations:
        truncateds = torch.zeros_like(terminateds)
    if bootstrap_values is None:
        bootstrap_values = torch.zeros_like(rewards)
    done  = (terminateds + truncateds).clamp(max=1.0)
    carry = bootstrap_values[:, -1]
    u = rewards + gamma * truncateds * bootstrap_values
    v = gamma * (1.0 - done)
    return _run_scan(u, v, carry)

rl_triton.ops.returns.compute_eligibility_traces(gradients, dones, gamma, lambda_, seed_values=None)

Compute accumulating eligibility traces via a forward associative scan.

Recurrence:

  • z[t] = g[t] + γ·λ·(1 - done[t]) * z[t-1], z[-1] = seed

g[t] is the per-step input: the value-function gradient ∇_w V̂(s_t) for general function approximation, or the feature vector x(s_t) in the linear case. Unlike all other kernels, this scan runs forward in time. Limited to seq_len <= 131072.

Parameters:

Name Type Description Default
gradients Tensor

g[t] -- value-function gradients or feature vectors, [num_envs, seq_len], float32, CUDA.

 required
dones Tensor

Episode termination flags (1.0=done), [num_envs, seq_len], float32, CUDA.

 required
gamma float

Discount factor.

 required
lambda_ float

Trace decay parameter.

 required
seed_values Tensor | None

Initial trace z[-1] per environment, shape [num_envs]. Defaults to zeros.

 None

Returns:

Name Type Description
traces Tensor

z[t], shape [num_envs, seq_len], float32.
Source code in rl_triton/ops/returns.py 
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
def compute_eligibility_traces(
    gradients: torch.Tensor,
    dones: torch.Tensor,
    gamma: float,
    lambda_: float,
    seed_values: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Compute accumulating eligibility traces via a forward associative scan.

    Recurrence:

    - z[t] = g[t] + γ·λ·(1 - done[t]) * z[t-1],  z[-1] = seed

    g[t] is the per-step input: the value-function gradient ∇_w V̂(s_t) for
    general function approximation, or the feature vector x(s_t) in the linear
    case.  Unlike all other kernels, this scan runs forward in time.
    Limited to seq_len <= 131072.

    Args:
        gradients:   g[t] -- value-function gradients or feature vectors,
                     [num_envs, seq_len], float32, CUDA.
        dones:       Episode termination flags (1.0=done), [num_envs, seq_len], float32, CUDA.
        gamma:       Discount factor.
        lambda_:     Trace decay parameter.
        seed_values: Initial trace z[-1] per environment, shape [num_envs].
                     Defaults to zeros.

    Returns:
        traces: z[t], shape [num_envs, seq_len], float32.
    """
    num_envs, seq_len = gradients.shape

    # Cheap structural checks -- always-on.
    assert gradients.is_cuda and dones.is_cuda, "gradients and dones must be on CUDA"
    assert gradients.dtype == torch.float32, f"gradients: expected float32, got {gradients.dtype}"
    assert dones.dtype == torch.float32,     f"dones: expected float32, got {dones.dtype}"
    assert gradients.shape == dones.shape,   "gradients and dones must have the same shape"
    if seed_values is not None:
        assert seed_values.shape == (num_envs,), \
            f"seed_values must have shape [{num_envs}], got {seed_values.shape}"
        assert seed_values.is_cuda, "seed_values must be on CUDA"

    assert seq_len <= _FLAT_MAX_SEQ_LEN, (
        f"seq_len={seq_len} exceeds the flat kernel limit {_FLAT_MAX_SEQ_LEN}. "
        "A chunked forward scan kernel has not been implemented yet."
    )

    gradients = gradients.contiguous()
    dones     = dones.contiguous()

    has_seed = seed_values is not None
    if has_seed:
        seed_values = seed_values.contiguous()

    out = torch.empty_like(gradients)

    BLOCK_SIZE = triton.next_power_of_2(seq_len)
    num_warps  = _WARPS_LIGHT.get(BLOCK_SIZE, 16)
    num_stages = 2 if BLOCK_SIZE >= 1024 else 1

    eligibility_traces_fused_kernel[(num_envs,)](
        gradients, dones,
        out, seed_values,
        seq_len, gradients.stride(0),
        gamma=gamma, lambda_=lambda_,
        BLOCK_SIZE=BLOCK_SIZE,
        num_warps=num_warps,
        num_stages=num_stages,
        HAS_SEED=has_seed,
    )
    return out

rl_triton.ops.prefix_sum.compute_episodic_prefix_sum(inputs, dones, seed_values=None)

Cumulative sum that resets to zero at episode boundaries.

Recurrence:

  • C[t] = x[t] + (1 - done[t]) * C[t-1], C[-1] = seed
  • done[t]=1: accumulator resets, C[t] = x[t]
  • done[t]=0: accumulator continues, C[t] = x[t] + C[t-1]

Limited to seq_len <= 131072.

Parameters:

Name Type Description Default
inputs Tensor

Values to accumulate x[t], [num_envs, seq_len], float32, CUDA.

 required
dones Tensor

Episode termination flags (1.0=done), [num_envs, seq_len], float32, CUDA.

 required
seed_values Tensor | None

Initial carry C[-1] per environment, shape [num_envs]. Defaults to zeros.

 None

Returns:

Name Type Description
prefix_sums Tensor

C[t], shape [num_envs, seq_len], float32.
Source code in rl_triton/ops/prefix_sum.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
def compute_episodic_prefix_sum(
    inputs: torch.Tensor,
    dones: torch.Tensor,
    seed_values: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Cumulative sum that resets to zero at episode boundaries.

    Recurrence:

    - C[t] = x[t] + (1 - done[t]) * C[t-1],  C[-1] = seed
    - done[t]=1: accumulator resets, C[t] = x[t]
    - done[t]=0: accumulator continues, C[t] = x[t] + C[t-1]

    Limited to seq_len <= 131072.

    Args:
        inputs:      Values to accumulate x[t], [num_envs, seq_len], float32, CUDA.
        dones:       Episode termination flags (1.0=done), [num_envs, seq_len], float32, CUDA.
        seed_values: Initial carry C[-1] per environment, shape [num_envs].
                     Defaults to zeros.

    Returns:
        prefix_sums: C[t], shape [num_envs, seq_len], float32.
    """
    num_envs, seq_len = inputs.shape

    if _CORRECTNESS_WARNINGS():
        assert inputs.is_cuda and dones.is_cuda, "inputs and dones must be on CUDA"
        assert inputs.dtype == torch.float32, f"inputs: expected float32, got {inputs.dtype}"
        assert dones.dtype == torch.float32,  f"dones: expected float32, got {dones.dtype}"
        assert inputs.shape == dones.shape,   "inputs and dones must have the same shape"
        if seed_values is not None:
            assert seed_values.shape == (num_envs,), \
                f"seed_values must have shape [{num_envs}], got {seed_values.shape}"
            assert seed_values.is_cuda, "seed_values must be on CUDA"

    assert seq_len <= _FLAT_MAX_SEQ_LEN, (
        f"seq_len={seq_len} exceeds the flat kernel limit {_FLAT_MAX_SEQ_LEN}. "
        "A chunked forward scan kernel has not been implemented yet."
    )

    inputs = inputs.contiguous()
    dones  = dones.contiguous()

    has_seed = seed_values is not None
    if has_seed:
        seed_values = seed_values.contiguous()

    out = torch.empty_like(inputs)

    BLOCK_SIZE = triton.next_power_of_2(seq_len)
    num_warps  = _WARPS.get(BLOCK_SIZE, 16)
    num_stages = 2 if BLOCK_SIZE >= 1024 else 1

    prefix_sum_fused_kernel[(num_envs,)](
        inputs, dones,
        out, seed_values,
        seq_len, inputs.stride(0),
        BLOCK_SIZE=BLOCK_SIZE,
        num_warps=num_warps,
        num_stages=num_stages,
        HAS_SEED=has_seed,
    )
    return out