AID-H-012.002: Learned Reward-Signal Integrity & Calibration

Version 1.20261015 · Actionable sub-technique

Protect training-specific integrity of learned reward functions built from preference pairs or expert demonstrations. Enforce signal-mode configuration, leakage-resistant train and evaluation partitions, reward-model calibration, abstention behavior, source-concentration checks, and adversarial validation.

Pillars: model

Lifecycle: building

Scope boundary

Owns integrity and calibration of a learned reward signal derived from preference or demonstration data. It consumes label provenance but does not own the general labeling pipeline or the final policy's reward-gaming gate.

AID-H-002.005: AID-H-012.002 trains, calibrates, and validates the learned reward signal; AID-H-002.005 proves provenance and integrity of the preference pairs, labels, or demonstrations supplied to it. Trusted labels are necessary inputs but do not prove reward-model calibration.

AID-H-012.005: AID-H-012.002 governs the learned reward model as a predictive signal; AID-H-012.005 tests whether the optimizing policy exploits that reward model or verifier. A calibrated evaluator can still contain shortcuts that a policy discovers under optimization pressure.

Implementation guidance

AID-H-012.002-G001: Train and validate learned reward from provenance-bound preference or demonstration signals using isolated splits, concentration checks, calibration metrics, and adversarial holdouts.

AID-H-012.002-G001

Train and validate learned reward from provenance-bound preference or demonstration signals using isolated splits, concentration checks, calibration metrics, and adversarial holdouts.

Apply this control when an organization trains a reward function from human preference pairs or expert demonstrations. The signed release policy selects exactly one signal mode and defines the task/rater/collection split keys, required rater cohorts, abstention handling, environment semantics, and evaluation thresholds. A workload with no learned human-signal reward is NOT_APPLICABLE. Preference and demonstration data are untrusted governed inputs, not safety ground truth.

Implement masked, validated preference learning

# File: rl_rewards/preference_learning.py
from __future__ import annotations

import hashlib
import json
from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor


def tensor_sha256(value: Tensor) -> str:
    tensor = value.detach().cpu().contiguous()
    header = json.dumps(
        {"dtype": str(tensor.dtype), "shape": list(tensor.shape)},
        sort_keys=True,
        separators=(",", ":"),
    ).encode()
    return hashlib.sha256(
        header + tensor.view(torch.uint8).numpy().tobytes(order="C")
    ).hexdigest()


@dataclass(frozen=True)
class PreferenceBatch:
    pair_ids: tuple[str, ...]
    preferred_states: Tensor
    preferred_actions: Tensor
    preferred_mask: Tensor
    rejected_states: Tensor
    rejected_actions: Tensor
    rejected_mask: Tensor

    def validate(self) -> None:
        tensors = [
            self.preferred_states,
            self.preferred_actions,
            self.rejected_states,
            self.rejected_actions,
        ]
        if any(value.ndim != 3 or not torch.isfinite(value).all() for value in tensors):
            raise ValueError("state/action trajectories must be finite [batch,time,features] tensors")
        batch, preferred_time, _ = self.preferred_states.shape
        if len(self.pair_ids) != batch or len(set(self.pair_ids)) != batch:
            raise ValueError("pair IDs must be complete and unique")
        if self.preferred_actions.shape[:2] != (batch, preferred_time):
            raise ValueError("preferred state/action trajectory shapes differ")
        if self.rejected_states.shape[0] != batch or self.rejected_actions.shape[:2] != self.rejected_states.shape[:2]:
            raise ValueError("rejected state/action trajectory shapes differ")
        if self.preferred_mask.dtype != torch.bool or self.preferred_mask.shape != (batch, preferred_time):
            raise ValueError("preferred mask must cover every timestep")
        if self.rejected_mask.dtype != torch.bool or self.rejected_mask.shape != self.rejected_states.shape[:2]:
            raise ValueError("rejected mask must cover every timestep")
        if bool((self.preferred_mask.sum(dim=1) == 0).any()) or bool((self.rejected_mask.sum(dim=1) == 0).any()):
            raise ValueError("each trajectory must contain an observed timestep")


class RewardModel(nn.Module):
    def __init__(self, state_dim: int, action_dim: int, hidden_dim: int) -> None:
        super().__init__()
        if min(state_dim, action_dim, hidden_dim) <= 0:
            raise ValueError("model dimensions must be positive")
        self.net = nn.Sequential(
            nn.Linear(state_dim + action_dim, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, 1),
        )

    def forward(self, states: Tensor, actions: Tensor) -> Tensor:
        if states.shape[:-1] != actions.shape[:-1]:
            raise ValueError("state/action leading dimensions differ")
        output = self.net(torch.cat([states, actions], dim=-1)).squeeze(-1)
        if not torch.isfinite(output).all():
            raise RuntimeError("reward model returned non-finite values")
        return output


def masked_return(values: Tensor, mask: Tensor) -> Tensor:
    if values.shape != mask.shape:
        raise ValueError("reward and mask shapes differ")
    return (values * mask.to(values.dtype)).sum(dim=1)


def preference_loss(model: RewardModel, batch: PreferenceBatch) -> tuple[Tensor, dict]:
    batch.validate()
    preferred = masked_return(
        model(batch.preferred_states, batch.preferred_actions),
        batch.preferred_mask,
    )
    rejected = masked_return(
        model(batch.rejected_states, batch.rejected_actions),
        batch.rejected_mask,
    )
    margin = preferred - rejected
    loss = -F.logsigmoid(margin).mean()
    if not torch.isfinite(loss):
        raise RuntimeError("preference loss is non-finite")
    receipt = {
        "schema_version": "aidefend.preference-batch.v1",
        "pair_count": len(batch.pair_ids),
        "pair_ids_sha256": hashlib.sha256(
            json.dumps(batch.pair_ids, separators=(",", ":")).encode()
        ).hexdigest(),
        "preferred_states_sha256": tensor_sha256(batch.preferred_states),
        "preferred_actions_sha256": tensor_sha256(batch.preferred_actions),
        "preferred_mask_sha256": tensor_sha256(batch.preferred_mask),
        "rejected_states_sha256": tensor_sha256(batch.rejected_states),
        "rejected_actions_sha256": tensor_sha256(batch.rejected_actions),
        "rejected_mask_sha256": tensor_sha256(batch.rejected_mask),
        "loss": float(loss.detach()),
        "mean_margin": float(margin.detach().mean()),
    }
    return loss, receipt

Govern dataset construction and the demonstration alternative

Before training, verify a signed manifest that binds every sample ID to task, rater, rater cohort, rubric version, collection batch, source trajectory, label/abstention, and split. Reject duplicate sample IDs, source-trajectory overlap across splits, rater or collection-batch leakage forbidden by policy, missing required cohorts, and post-signature mutation. Keep adversarial policy-violation examples in an access-controlled holdout never exposed to optimization.

If the signed mode is demonstration, use a pinned maintained implementation such as imitation AIRL with the exact deployed Gymnasium environment wrapper, serialized expert-trajectory manifest, generator checkpoint, reward network, replay-buffer capacity, terminal/truncation semantics, normalization state, seeds, and update schedule. Require a loader readback proving that every admitted trajectory and only those trajectories entered the run; do not substitute a short custom discriminator loop.

Verification, evidence, and outcomes

Retain signature-verification receipts, dataset/rubric/rater-cohort/split digests, all batch receipts, dependency and environment image locks, seeds, reward-model and candidate-policy artifacts, calibration and subgroup metrics, and complete normal plus evaluator-gaming, sycophancy, collusion, hidden-state, label-reversal, low-agreement, and distribution-shift results. An independent runner reconstructs the split, checks zero forbidden overlap, reloads the exact reward artifact, replays held-out pairs or demonstrations, and evaluates the trained policy using independent evidence not used for optimization.

Action: Select and sign one human-signal mode, enforce complete sample lineage and split isolation before training, and block reward-model or policy promotion until an independent holdout replay produces PASS.

Open-source tools

imitation
Gymnasium
PyTorch
Stable-Baselines3
Label Studio
Hugging Face TRL (RewardTrainer execution; independent split and calibration gates required)

Threat mappings

MITRE ATLAS

AML.T0020 Training Data Poisoning (governed preference and demonstration signals resist poisoned RL feedback)
AML.T0059 Erode Dataset Integrity (split and lineage controls protect feedback-dataset integrity)

MAESTRO

Data Poisoning (Training Phase) (L1)
Data Poisoning (L2)

OWASP LLM Top 10 2026

LLM05:2026 Data and Model Poisoning

OWASP ML Top 10 2023

ML08:2023 Model Skewing

OWASP Top 10 for Agentic Applications 2026

NIST Adversarial Machine Learning 2025

NISTAML.013 Data Poisoning
NISTAML.037 Training Data Attacks (includes RLHF poisoning)
NISTAML.027 Misaligned Outputs

Cisco Integrated AI Security and Safety Framework

AISubtech-6.1.2 Reinforcement Biasing
AISubtech-6.1.3 Reinforcement Signal Corruption
AITech-2.3 Reward Hacking
AITech-6.1 Training Data Poisoning

Google Secure AI Framework 2.0 - Risks

DP: Data Poisoning (preference and demonstration signals are governed training data)

Databricks AI Security Framework 3.0

Datasets 3.1: Data poisoning (preference and demonstration signal poisoning)
Agents - Core 13.7: Misaligned & Deceptive Behaviors