Causally Robust Reward Learning from Reason-Augmented Preference Feedback
Organizations: Thomas Lord Department of Computer Science, University of Southern California
Abstract
Preference-based reward learning is widely used for shaping agent behavior to match a user's preference, yet its sparse binary feedback makes it especially vulnerable to causal confusion. The learned reward often latches onto spurious features that merely co-occur with preferred trajectories during training, collapsing when those correlations disappear or reverse at test time. We introduce ReCouPLe, a lightweight framework that uses natural language rationales to provide the missing causal signal. Each rationale is treated as a guiding projection axis in an embedding space, training the model to score trajectories based on features aligned with that axis while de-emphasizing context that is unrelated to the stated reason. Because the same rationales (e.g., "avoids collisions", "completes the task faster") can appear across multiple tasks, ReCouPLe naturally reuses the same causal direction whenever tasks share semantics, and transfers preference knowledge to novel tasks without extra data or language-model fine-tuning. Our learned reward model can ground preferences on the articulated reason, aligning better with user intent and generalizing beyond spurious features. ReCouPLe outperforms baselines by up to 1.5x in reward accuracy under distribution shifts, and 2x in downstream policy performance in novel tasks. We have released our code at https://github.com/mj-hwang/ReCouPLe
Figures & tables
| 2-task | 4-task | |||||||||||
| In Distribution | Color Swapped | In Distribution | Color Swapped | |||||||||
| Model | Pick | Place | Pick | Place | Pick | Push | Place | Pull | Pick | Push | Place | Pull |
| Single Task | ||||||||||||
| BT | 0.980 | 1.000 | 0.540 | 0.830 | 0.980 | 1.000 | 1.000 | 1.000 | 0.540 | 0.987 | 0.830 | 0.867 |
| Multi-Task | ||||||||||||
| BT-Multi | 0.953 | 1.000 | 0.600 | 0.820 | 0.987 | 1.000 | 1.000 | 1.000 | 0.707 | 1.000 | 0.840 | 0.907 |
| Model | Training Tasks | Novel Task | ||
| Push | Push-Wall | Pick-Place-Wall | Pick-Place | |
| Multi-Task | ||||
| BT-Multi | 0.873 | 0.893 | 0.577 | 0.547 |
| Multi-Task w/ Reasons | ||||
| RFP | 0.870 | 0.900 | 0.647 | 0.553 |
| ReCouPLe-EC | 0.863 | 0.843 | 0.650 | 0.663 |
| Model | 2-task | 4-task | ||
|---|---|---|---|---|
| ID | OOD | ID | OOD | |
| ReCouPLe | 0.995 | 0.872 | 1.000 | 0.878 |
| ReCouPLe-no-consistency | 0.980 | 0.726 | 0.977 | 0.745 |
| ReCouPLe-no-consistency-no-ratio | 0.987 | 0.727 | 0.990 | 0.730 |
| 2-task | 4-task | |||||||||||
| In Distribution | Color Swapped | In Distribution | Color Swapped | |||||||||
| Model | Pick | Push | Pick | Push | Pick | Push | Place | Pull | Pick | Push | Place | Pull |
| Single Task | ||||||||||||
| BT | 0.833 | 1.000 | 0.167 | 0.610 | 0.833 | 1.000 | 0.980 | 1.000 | 0.167 | 0.610 | 0.460 | 0.053 |
| Multi-Task | ||||||||||||
| BT-Multi | 0.870 | 0.999 | 0.167 | 0.673 | 0.867 | 1.000 | 0.987 | 1.000 | 0.533 | 0.867 | 0.833 | 0.587 |
| Model | Push (Visual OOD) | Push-Wall (Visual OOD) |
|---|---|---|
| BT (Baseline) | 0.233 | 0.113 |
| BT-Multi | 0.367 | 0.227 |
| ReCouPLe-EC | 0.713 | 0.760 |
| ReCouPLe-IC | 0.793 | 0.900 |
Appendix figures & tables8 assets
Supplementary material from the paper’s appendix.
Appendix
| Task | Task Description | Larger Cube Color | Smaller Cube Color |
|---|---|---|---|
| MS-Pick-Larger | “pick up larger cube to target sphere” | Red | Blue |
| MS-Place-Larger | “place larger cube in target bin” | Blue | Red |
| MS-Push-Larger | “push larger cube toward target line” | Blue | Red |
| MS-Pull-Larger | “pull larger cube toward green line” | Red | Blue |
| Task | Task Description |
|---|---|
| Push | “make contact and push puck to goal” |
| Push-Wall | “make contact, bypass wall via waypoint, and push puck to goal” |
| Pick-Place | “pick up puck, lift it, and place it on goal” |
| Pick-Place-Wall | “pick up puck, use waypoint to bypass wall, and place it on goal” |
| Code | Reason Text |
|---|---|
| Push | “pushes more decisively after making contact” |
| Push | “pushes puck closer to goal” |
| Push-Wall | “pushes puck toward waypoint” |
| Push-Wall | “guides puck past wall” |
| Push-Wall | “pushes puck toward goal after clearing wall” |
| Push & Push-Wall | “maintains firm grip on puck” |
| Model | Pick (OOD) | Place (OOD) |
|---|---|---|
| BT (Baseline) | 0.540 | 0.830 |
| ReCouPLe-EC (Original) | 0.820 | 0.940 |
| ReCouPLe-EC (Diversified) | 0.867 | 0.967 |
| ReCouPLe-IC (Original) | 0.633 | 0.807 |
| ReCouPLe-IC (Diversified) | 0.620 | 0.820 |
| Model | Pick (OOD) | Place (OOD) |
|---|---|---|
| BT (Baseline) | 0.540 | 0.830 |
| ReCouPLe-EC (25% Reasons) | 0.783 | 0.940 |
| ReCouPLe-EC (50% Reasons) | 0.767 | 0.947 |
| ReCouPLe-IC (25% Reasons) | 0.673 | 0.807 |
| ReCouPLe-IC (50% Reasons) | 0.660 | 0.847 |
| Preference Queries | 200 | 500 | 1000 | 2000 |
|---|---|---|---|---|
| Task: Pick (OOD) | ||||
| BT-Multi | 0.573 | 0.627 | 0.600 | 0.693 |
| RFP | 0.540 | 0.613 | 0.620 | 0.760 |
| ReCouPLe-EC | 0.733 | 0.820 | 0.820 | 0.913 |
| ReCouPLe-IC | 0.547 | 0.580 | 0.633 | 0.793 |
| Task: Place (OOD) | ||||
Reward Modeling
Reward modeling learns scores or rankings that capture human goals and task success, guiding AI selection, reinforcement learning, and alignment.
See all 313 papers