Understanding Learned Reward Functions

TL;DR

Gradient saliency, occlusion maps, and counterfactuals reveal shortcut reward algorithms in Gridworld and Atari, even when PPO behavior remains competent.

cs.LG 🟡 Intermediate 2020-12-11 34 views
Eric J. Michaud Adam Gleave Stuart Russell
reward learning RL interpretability saliency distribution shift AI safety

Key Findings

Methodology

The paper audits learned reward networks directly rather than interpreting only policies trained on them. It combines raw gradient saliency, ∂R/∂(s,a,s′), Gaussian-blur occlusion maps, and hand-crafted counterfactual inputs. Reward models are trained by regression on ground-truth labels from expert rollouts. Experiments use 11×11 grayscale Gridworld observations and Nature-DQN-preprocessed Atari 2600 inputs, then compare reward sensitivity with PPO behavior under environment changes.

Key Results

  • In CoinFlipGoal, the model is far more sensitive to s′ than s and appears to learn the wrong rule: reward the absence of a visible goal. Removing the goal produces 0.949; two goals produce 0.004; many goals produce −1.727.
  • The distorted reward still trains a nearly competent PPO policy in TwoGoals, showing that large numerical reward errors need not immediately cause behavioral failure. In Atari, Breakout attends mainly to the ball, whereas Seaquest attends mainly to the immediately updated score display.
  • After removing the score, Seaquest predictions become noisy while Breakout predictions remain stable. With 10M PPO steps, true/regressed returns were 360.3/141.9 for Breakout with score, 258.5/90.1 without score, 1746.0/712.0 for Seaquest with score, and 934.0/540.0 without score.

Significance

The study makes reward auditing a distinct problem from policy interpretation. It shows that even clean supervision from synthetic ground-truth rewards can produce models that exploit correlated but non-causal features. High policy return therefore does not establish that the learned objective is correct, especially under deployment distribution shift. The work also highlights a mathematical difference from ordinary prediction: reward values admit transformations such as potential-based shaping that can change outputs without changing optimal behavior.

Technical Contribution

This is an early systematic application of neural-network interpretability directly to learned reward functions. It combines raw gradients, Gaussian-blur occlusion, and counterfactual testing into a practical workflow: generate hypotheses from saliency, test mechanisms with interventions, and evaluate resulting policies separately. The authors also report that Integrated Gradients, Guided Grad-CAM, DeepLift, and Shapley values performed poorly in their setting, motivating reward-specific interpretability rather than simply importing policy explanations.

Novelty

Unlike policy-focused saliency work and Russell and Santos’s approach of fitting decision trees and applying Gini importance plus LIME, this paper probes the original reward network. Its deeper novelty is conceptual: sensitivity of reward output and correctness of induced policy are different targets. A reward can implement the wrong algorithm, be highly non-robust numerically, and nevertheless produce competent behavior in a particular modified environment.

Limitations

  • Saliency predicts whether perturbations change reward output, not whether the optimal policy changes. It can therefore flag behaviorally irrelevant transformations, including positive scaling or potential-based shaping.
  • The experiments use synthetic labels, small visual environments, and limited Atari games rather than noisy human preferences or complex real dynamics. Hand-crafted counterfactuals are costly, and the methods can create false confidence because they are not complete safety tests.

Future Work

Future work should develop reward-specific metrics that account for policy equivalence, reward ordering, and transformations such as potential shaping. EPIC is suggested as a starting point. Important extensions include noisy human preference data, automatically generated counterfactual environments, causal feature tests, partial observability, and formal or empirical guarantees that a learned reward preserves intended behavior under deployment shifts.

AI Executive Summary

As reinforcement learning moves toward real-world tasks, designers increasingly need to learn objectives from demonstrations or human preferences rather than write rewards by hand. This creates a subtle safety problem: a model may predict training labels accurately while learning a shortcut, such as treating a brake indicator or score display as the reason for desirable behavior. Policy performance alone cannot reveal whether the reward was learned correctly or whether an optimizer merely exploited the available correlations.

Michaud, Gleave, and Russell propose directly auditing reward networks with raw gradient saliency, Gaussian-blur occlusion maps, and counterfactual inputs. In an 11×11 Gridworld, the learned model focused mainly on the next observation and behaved as if the absence of a visible goal deserved reward. Removing the goal yielded 0.949, two goals yielded 0.004, and many goals yielded −1.727, despite the true reward being 1 only when the agent covered a goal. In Atari, Breakout models attended to the ball, while Seaquest models mainly read the score display because it updated at the reward timestep.

The results expose both the power and the limits of interpretability. Removing Seaquest’s score made predictions noisy, yet a PPO agent trained for 10M steps still achieved 540.0 true return, compared with 712.0 when the score was present. Thus saliency can reveal suspicious dependencies and forecast output sensitivity, but not reliably forecast policy transfer. The authors call for reward-specific auditing based on policy-relevant equivalence, causal interventions, and tools such as EPIC.

Deep Analysis

Background

RL systems have surpassed humans in Atari, DOTA, Go, and StarCraft, where rewards can often be specified manually. Real-world objectives involving safety, preferences, and long-term consequences are harder to formalize, motivating IRL, preference learning, and hybrid demonstration-comparison methods. EPIC compares reward functions when a ground-truth reward exists, but real deployment precisely lacks such a reference, creating a need for direct auditing.

Core Problem

The central question is whether a learned reward captures the user’s objective or merely predicts labels using spurious correlates. A model may rely on score displays, dashboard lights, object counts, or rendering artifacts. Moreover, changed reward values do not necessarily imply changed optimal policies, so standard neural-network explanations can identify sensitivity without identifying behavioral danger.

Innovation

  • ��Directly analyze the original reward network rather than a simplified surrogate.
  • ��Use gradients and occlusion to generate feature hypotheses, then counterfactuals to test mechanisms.
  • ��Separate reward-output robustness from policy transfer and measure both.
  • ��Frame reward interpretability around its special mathematical structure, including potential-based shaping and policy-equivalent transformations.

Methodology

  • ��Training: regress on ground-truth reward labels from expert rollouts; Gridworld uses R(s,s′), Atari uses R(s′).
  • ��Gradients: compute ∂R/∂(s,a,s′) to estimate local sensitivity.
  • ��Occlusion: Gaussian-blur image regions and record reward changes; Atari uses σ=3.
  • ��Counterfactuals: delete, duplicate, or multiply goals; remove Atari score displays.
  • ��Behavioral test: train PPO and evaluate returns using the environment’s true reward.
  • ��Comparison: inspect both saliency hypotheses and transfer performance rather than treating either as sufficient.

Experiments

Gridworld observations are 11×11 grayscale images. CoinFlipGoal places a goal randomly in the top-left or bottom-right; TwoGoals places both goals. Atari uses the Nature DQN preprocessing pipeline: 84×84 grayscale frames stacked four deep, with Breakout and Seaquest. The study compares true-reward and regressed-reward models, score-display interventions, and PPO training for 10M steps.

Results

The Gridworld model attends mainly to s′ and implements a rule opposite to the true coverage rule; counterfactual outputs verify the hypothesis. Yet the same model can still support competent behavior in TwoGoals. In Atari, ball saliency in Breakout contrasts with score saliency in Seaquest. Removing the score harms reward prediction but does not eliminate useful learning, demonstrating that output sensitivity is a conservative proxy for policy failure.

Applications

The workflow can audit preference-learning, IRL, and human-feedback reward models before deployment. Engineers can test whether models read dashboards, metadata, simulator artifacts, or timing leaks, then evaluate behavior after objects or interfaces are removed, duplicated, delayed, or rearranged. Relevant domains include autonomous driving, robotics, healthcare decision support, and any system whose objective is difficult to specify manually.

Limitations & Outlook

Synthetic labels and small visual environments limit direct extrapolation to noisy human preferences, partial observability, and complex real dynamics. Gradients are local, occlusion depends on spatial scale, and counterfactual design is manual. None of the methods guarantees preservation of the intended optimal policy; future systems should combine EPIC-like equivalence, causal interventions, policy-ranking tests, and formal verification.

Plain Language Accessible to non-experts

Imagine a factory supervisor whose real job is to reward workers for sending the correct package to the correct destination. During training, a green lamp happens to turn on whenever the best worker succeeds. The supervisor may then learn the shortcut “green lamp means good work,” even though the lamp is only a side effect. As long as the factory operates normally, this mistake can remain invisible.

The paper uses three kinds of detective work. A gradient method asks which tiny parts of the picture make the score twitch. An occlusion test covers one region at a time and watches what changes. A counterfactual test deliberately creates strange situations: remove the target, add two targets, or fill the scene with targets. In the Gridworld, the learned supervisor gave 0.949 when the target was removed and −1.727 when many targets were present—evidence that it had learned a backwards rule.

But a strange score does not always make the worker fail immediately. Other clues may still lead to the right package. That is why the paper separates “what the supervisor notices” from “whether the worker succeeds.” A trustworthy audit must test new layouts, missing signals, delayed displays, and other situations where accidental clues disappear.

ELI14 Explained like you're 14

Suppose you train a game bot to find a green square in a maze. The rule is simple: touch the target and get 1 point; otherwise get 0. After watching many examples, the bot seems smart. But researchers want to know whether it understands the rule or is just memorizing a weird visual trick.

They use three detective tools. One is like a magnifying glass: change each pixel slightly and see which changes the score. Another covers part of the screen with blur. The third is a prank test: remove the target, add a second target, or cover the map with targets. The bot gives 0.949 when the target disappears and −1.727 when there are many targets. Yikes! It may have learned “no target means good,” not “touch target means good.”

They also test Atari games. In Breakout, the bot watches the ball, which makes sense. In Seaquest, it mostly watches the score display because the score changes right when points are awarded. Remove the score and its predictions become messy, although a PPO player still earns 540.0 real points after 10M training steps. So the bot can use a silly clue and still play decently for a while.

That is like a student who gets answers right by copying page numbers instead of understanding the lesson. Change the worksheet and the trick may collapse. The big lesson: never judge an AI goal only by its score. Ask what it is watching, invent weird test cases, and check whether it still behaves correctly when the shortcuts vanish!

Glossary

Reward function

A mapping from states, actions, or transitions to numerical feedback. RL agents optimize accumulated reward, so errors in this function can redirect behavior.

The paper’s primary object of interpretation and safety auditing.

Raw gradient saliency

The derivative ∂R/∂x measures local sensitivity of reward to each input dimension. A large derivative indicates local influence, not necessarily causal importance.

Used to compare sensitivity to current state s and next state s′ in Gridworld.

Occlusion map

An attribution method that blurs or masks input regions and records output changes. It is suited to spatially contiguous features in images.

Applied with Gaussian blur σ=3 to Breakout and Seaquest.

Counterfactual input

A deliberately altered input used to test a hypothesis about model behavior. It can expose rules that ordinary data never reveals.

Used to add or remove goals and remove Atari score displays.

Potential-based shaping

A reward transformation R+γΦ(s′)−Φ(s) that, under standard conditions, preserves the optimal policy. It demonstrates why numerical reward differences can be behaviorally irrelevant.

Cited as a reason ordinary saliency is conservative.

EPIC

A pseudometric for comparing reward functions through policy-relevant reward equivalence. It is most informative when a reference reward or comparable transition structure is available.

Proposed as a starting point for reward-specific interpretability metrics.

Open Questions Unanswered questions from this research

  • 1 How can an auditor predict whether a reward perturbation changes the optimal policy rather than merely changing numerical output? The answer requires combining reward equivalence, dynamics, and policy ordering.
  • 2 Do saliency methods identify causal features or only correlated features? This remains unclear with noisy human preferences, partial observability, and realistic interventions.
  • 3 Can counterfactual environments be generated automatically with deployment-level safety guarantees? The paper demonstrates diagnosis, not scalable certification.

Applications

Immediate Applications

Pre-deployment reward audits

Safety teams can run gradients, occlusion maps, and counterfactual interventions on preference or IRL models. They can test dependence on scoreboards, dashboard lights, metadata, and timing leaks, then evaluate PPO behavior using the true task objective.

Visual RL data hygiene

In simulators for robotics, driving, or games, engineers can locate salient interface artifacts, remove them from training, and compare in-distribution and shifted-environment returns. This reduces the chance that a reward model exploits rendering or logging shortcuts.

Long-term Vision

Certified preference-aligned agents

A future audit stack could combine causal saliency, automatic counterfactual generation, EPIC-like reward equivalence, policy-ranking tests, and formal verification. This would make learned objectives not only trainable, but inspectable and certifiable before deployment.

Abstract

In many real-world tasks, it is not possible to procedurally specify an RL agent's reward function. In such cases, a reward function must instead be learned from interacting with and observing humans. However, current techniques for reward learning may fail to produce reward functions which accurately reflect user preferences. Absent significant advances in reward learning, it is thus important to be able to audit learned reward functions to verify whether they truly capture user preferences. In this paper, we investigate techniques for interpreting learned reward functions. In particular, we apply saliency methods to identify failure modes and predict the robustness of reward functions. We find that learned reward functions often implement surprising algorithms that rely on contingent aspects of the environment. We also discover that existing interpretability techniques often attend to irrelevant changes in reward output, suggesting that reward interpretability may need significantly different methods from policy interpretability.

cs.LG