{"id":"543a8b36-04ff-429f-a5fd-d4a4ed859c44","arxiv_id":"2412.05766","paper_version":1,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":6.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":3,"one_line_summary":"Policy-Shaped Prediction weights a world model's reconstruction loss by policy-gradient salience aggregated through SAM segmentation, plus an adversarial action head, improving MBRL robustness to learnable distractors.","lead":"This paper introduces a training method that makes model-based reinforcement learning agents ignore predictable but useless background distractions by weighting the world model's reconstruction loss toward pixels the policy actually uses. It shows large gains on a new benchmark with such distractors, while keeping performance on standard benchmarks.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"The 2x robustness gain is not shown to come from policy-shaped prediction: Table 2 shows an adversarial-head-only variant reaches 324.4 on Reafferent Cheetah versus DreamerV3's 158.4, reproducing most of full PSP's 383.1, and a value-gradient-only variant scores 445.7, yet neither is tested on…","rationale":"The reader's weakest assumption was that the policy-gradient salience signal is stable and identifies task-relevant pixels. My concern goes one step further: the evidence in Table 2 suggests that this signal may not even be the operative component behind the headline improvement. The adversarial action head alone nearly reproduces the 2x gain, and the value-gradient-only variant beats full PSP on the novel benchmark. If these simpler variants also succeed on Hopper and Distracting Control, then the central contribution ('policy-shaped prediction') is not load-bearing for the reported results; the method would remain empirically interesting but its novelty and attribution would change substantially. This is a concrete, checkable empirical gap rather than a theoretical objection, and it is consistent with the paper's own data. The reader's verdict of CONDITIONAL is therefore the right level: the method is promising, but the causal claim needs to be established by testing the missing ablations and by providing the promised code. My concern does not move the verdict; it strengthens the conditions already attached to the conditional acceptance.","tokens_in":14405,"tokens_out":14060,"duration_ms":143936,"concrete_test":"Run the full PSP pipeline together with (a) adversarial-head-only ('None+SEG+Adv'), (b) policy-gradient-only (no SEG, no Adv), and (c) a new 'Value+SEG+Adv' variant (value-gradient weighting in Eq. 1, with segmentation and adversarial head) on all reported tasks: Reafferent Hopper/Cheetah, Distracting Hopper/Cheetah, and Unmodified DMC, using the same hyperparameters and 3 seeds. If any reduced variant matches full PSP within standard error on the Reafferent and Distracting benchmarks, the 2x claim is not attributable to policy-shaped prediction; if full PSP clearly wins across tasks, the synergy is supported.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The paper's central claim is that the PSP combination yields a 2x improvement on challenging distractions, and the title and abstract attribute this to policy-shaped prediction. The ablations in Table 2 undermine that attribution on the main new benchmark (Reafferent Cheetah Run). Removing policy-gradient weighting and segmentation while keeping only the adversarial action head ('None ✗ ✓') scores 324.4 ± 2.3, already 2.05x DreamerV3's 158.4 ± 45.7. Replacing the policy gradient with a value gradient and dropping the other components ('Value ✗ ✗') scores 445.7 ± 126.9, which is higher than full PSP's 383.1 ± 23.8. Full PSP does improve unaltered Cheetah (712.3 vs 674.2/381.7), suggesting some synergy, but the central robustness result appears driven by the adversarial head and/or value-gradient weighting rather than the policy-shaped mechanism emphasized as the main novelty. Crucially, the paper never reports these simpler variants on Hopper Stand or Distracting Control, so it is unknown whether the 'unmatched' results require the full method or merely components that are not the paper's primary contribution.","agreement_with_reader":"partial"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper proposes Policy-Shaped Prediction (PSP), a modification of DreamerV3 that reweights the image reconstruction loss with segmentation-aggregated gradients of the policy with respect to input pixels, and adds an adversarial action-prediction head intended to prevent the image encoder from wasting capacity on self-generated action information. It introduces Reafferent DeepMind Control, a new benchmark whose distracting backgrounds are deterministic functions of the previous action and time, and reports that PSP outperforms DreamerV3, DreamerPro, Task Informed Abstractions, Denoised MDPs, and DrQv2 on this benchmark while maintaining comparable performance on unmodified and standard Distracting Control tasks. The paper includes ablations on Cheetah Run, experiments with SAM and SAM2 segmentation backbones, and an appendix documenting computational overhead and resource usage.","tokens_in":14626,"tokens_out":5637,"duration_ms":54059,"significance":"If the attribution claim held, PSP would be a useful contribution to distraction-robust model-based reinforcement learning, and the Reafferent DMC benchmark is a well-motivated testbed for learnable distractors. The paper makes a good-faith effort to report compute costs, includes a code repository, and tests multiple segmentation backbones, which are strengths. However, the headline 2x robustness gain is not cleanly attributable to the policy-shaped prediction mechanism: the ablations in Table 2 show that an adversarial-head-only variant reproduces most of the gain on the main benchmark, and a value-gradient-only variant scores even higher. The Hopper Stand evidence is also statistically thin, with only one of three seeds succeeding. These issues need to be addressed before the central claims can be accepted.","major_comments":[{"comment":"The claim that the PSP combination yields the 2x robustness improvement is not supported by the ablations on the main benchmark. The adversarial-head-only variant ('None ✗ ✓') scores 324.4 ± 2.3 on Reafferent Cheetah Run, already 2.05x DreamerV3's 158.4 ± 45.7, and the value-gradient-only variant ('Value ✗ ✗') scores 445.7 ± 126.9, which is higher than full PSP's 383.1 ± 23.8. The narrative that segmentation-based aggregation is 'critical' is also not evident from these rows, since the policy-gradient-only variant without segmentation or adversarial head ('Policy ✗ ✗') scores 379.0, close to full PSP's 383.1. The abstract and contribution list should be rewritten to attribute the gain to the adversarial head and/or value-gradient weighting, or the full method must be shown to beat these components on the relevant environments.","section":"Section 3.4, Table 2"},{"comment":"The ablation study is limited to Cheetah Run, so the paper never establishes that the full PSP configuration is necessary for the Hopper Stand or Distracting Control results. Given the Cheetah ablations, the Hopper Stand outcome could be driven by a component other than policy-shaped prediction. Table 1's 'unmatched performance' claim is therefore not yet supported for those environments; at minimum, the missing ablation rows should be run or the claim should be restricted to the environments where the full method is actually compared against its components.","section":"Section 3.4"},{"comment":"The Reafferent Hopper Stand result is statistically weak: PSP's mean is 128.5 with a standard deviation of 215.7, and the text and Table 3 indicate that only one of three seeds succeeds. With n=3, this is not a reliable demonstration of 'scores beyond the reach of any of the baselines.' The paper should report individual seeds, provide additional runs, or explicitly qualify the claim as based on a single successful seed.","section":"Table 1, Reafferent Hopper Stand"},{"comment":"The method depends on several constants that receive no sensitivity analysis: the interpolation factor alpha = 0.9 in Eq. (3), the 99th-percentile clipping of the raw salience map, and the adversarial gradient scale epsilon = 1e3. The paper explicitly invokes alpha to avoid a chicken-and-egg failure where neither the policy nor the world model is useful, so it is load-bearing to know whether PSP is robust to reasonable variations of these values. A small sensitivity study or at least a discussion of the chosen values would strengthen the paper.","section":"Section 2.2-2.3"},{"comment":"The comparison between PSP and the baselines involves a substantial compute disparity: PSP trials use 4 A40 GPUs for training plus 4 A40 GPUs for segmentation, while baselines run on a single A40 GPU or a desktop 2070 SUPER. Since the baselines use 'default hyperparameters' and no tuning budget is reported, the resource asymmetry could confound the headline improvement. The paper should discuss whether the baselines were given comparable compute or tuning, or should temper the comparison accordingly.","section":"Appendix C"}],"minor_comments":[{"comment":"The symbol 'a' is overloaded: it denotes the action vector in 'a = E(pi(s))' and the scalar sum 'a = sum_j a_j'. Please use distinct notation for the vector and the scalar.","section":"Section 2.1, Eq. (1)"},{"comment":"The update 'phi <- Adam(gradient(L) - epsilon * derivative(L_adv_head)/derivative(rho), phi)' is confusing because the subtracted term is written as a derivative with respect to rho while the parameter being updated is phi. Please clarify that the adversarial gradient is applied only through the encoder parameters and specify exactly which parameters receive the subtracted gradient.","section":"Algorithm 1, line 12"},{"comment":"The SAM condition on Reafferent Hopper Stand is reported as 128.5 ± 215.7 in Table 1 and 130.3 ± 214.1 in Table 3. Please reconcile the discrepancy or explain why the numbers differ.","section":"Tables 1 and 3"},{"comment":"The sentence 'segmentation-based aggregation is critical to improving our model's performance amid distractors' is not clearly supported by Table 2, where the policy-gradient-only variant without segmentation achieves 379.0 on Reafferent Cheetah Run. Please soften this claim or provide the missing comparison that demonstrates the critical role of segmentation.","section":"Section 3.4"},{"comment":"The unqualified statement that PSP 'outperforms a variety of other approaches' is too strong given the high variance on Hopper Stand and the ablation results in Table 2. Please qualify the claim to reflect the environments and seeds for which the advantage is consistent.","section":"Abstract and Section 1"},{"comment":"There are minor typographical issues, including 'paramaterized' in Algorithm 1 and inconsistent spelling of 'VaGraM'/'VaGram' in Section 4. Please proofread.","section":"Throughout"}],"recommendation":"major_revision","confidential_remarks":"To the editor: This is a solid empirical paper with a useful new benchmark and a well-motivated method, but the attribution problem in Table 2 is serious: the headline robustness gain appears to come from the adversarial head or value-gradient weighting rather than from the policy-shaped prediction mechanism that the title and abstract emphasize. I would not reject, because the issue is addressable with additional ablations and a more careful framing, and the Reafferent DMC benchmark is a valuable contribution. Please also consider whether the compute asymmetry between PSP and the baselines should be explicitly flagged to the authors."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Thanks for the report. I read the paper carefully. My take: it is a solid submission with a genuinely useful new benchmark, but the central claim—that policy-shaped prediction is what drives the robustness—does not hold up under the ablations. The full method works, but the attribution is shaky.\n\nWhat's actually new and good: the Reafferent DeepMind Control benchmark is a nice stress test. Distractors are learnable, deterministic functions of action and time, so they are perfectly predictable but useless to the policy. That directly targets the failure mode where reconstruction-based world models waste capacity. The method itself is a reasonable combination: policy-gradient saliency weighted reconstruction, SAM-based object aggregation to stabilize the noisy gradients, and an adversarial action-prediction head. The segmentation aggregation is a sensible fix, and I give credit for testing SAM2 as an alternative. On Reafferent Cheetah Run, full PSP clearly beats DreamerV3 (383 vs 158) and most baselines. On Hopper Stand, PSP gets a nonzero mean where all model-based baselines are near zero, though only one of three seeds succeeds and the variance is huge.\n\nThe soft spots are real. The abstract claims PSP \"outperforms a variety of other approaches\" and a 2x improvement, but Table 1 shows TIA and DreamerPro beat PSP on individual Distracting Control tasks. More importantly, the ablations in Table 2 undermine the attribution of the robustness result. The adversarial-head-only variant (None, no segmentation, no policy gradient) scores 324.4 on Reafferent Cheetah, already 2x DreamerV3, and the value-gradient-only variant scores 445.7, higher than full PSP's 383.1. The paper acknowledges the trade-off but never tests these simpler variants on Hopper or Distracting Control, so we don't know whether the full method is needed there. The missing code link (appendix says \"this GitHub Repository\" with no URL) is a reproducibility issue.\n\nThe paper's own limitations section is honest, and the authors clearly know the method has an object-centric bias and compute costs. But the overclaim in the abstract and the untested ablations need fixing.\n\nWho this is for: anyone working on MBRL robustness, representation learning, or visual distraction. It deserves a serious referee because the benchmark and the method are worth engaging with, but the authors should qualify the claims, run the key ablations on all environments, and provide a working code link.\n\nRecommendation: send to peer review with a request for major revision.","headline":"Worth reading for the Reafferent benchmark and the segmentation-weighted saliency idea, but the paper overclaims attribution: the adversarial head alone reproduces most of the robustness gain on their main result, and a value-gradient variant beats the full method.","tokens_in":15210,"tokens_out":2335,"would_cite":false,"duration_ms":23787,"reading_group":"yes","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"Policy-shaped world models shrug off learnable distractions","keywords":["model-based reinforcement learning","world models","distraction robustness","policy-gradient saliency","image segmentation","adversarial action prediction","DreamerV3","Reafferent Deepmind Control"],"falsifier":"Train PSP on Reafferent Cheetah Run with per-segment loss weights replaced by random permutations of the true weights, keeping segmentation and the adversarial head intact: if the score stays near 383, the policy-gradient signal is not the cause of the gain, while a drop toward DreamerV3's 158 would confirm the salience weighting is doing the work.","tokens_in":14139,"feed_emoji":"🎯","tokens_out":6424,"duration_ms":55015,"temperature":0.7,"pith_summary":"Policy-Shaped Prediction (PSP) argues that a reconstruction-based world model wastes capacity when it devotes pixels to content that is predictable but useless for choosing actions, and that this waste can be fixed by telling the model which pixels its policy actually looks at. The paper shows that weighting the DreamerV3 image-reconstruction loss by object-aggregated policy gradients, plus an adversarial head that stops the encoder from re-encoding the previous action, lets an agent keep a faithful model of the task while ignoring learnable background distractions. On the new Reafferent Deepmind Control Suite, PSP scores 383.1 versus 158.4 for DreamerV3 on Cheetah Run and obtains a nonzero Hopper Stand mean score where every model-based baseline is near zero, though that Hopper mean comes from one of three seeds. If correct, the method gives a recipe for pointing a world model's finite capacity at what matters, without changing the agent at inference time.","feed_headline":"Policy-shaped world models ignore learnable distractions","feed_subtitle":"On predictable background noise, PSP scores 2x DreamerV3 on Cheetah Run with no loss on clean tasks.","key_machinery":"Three components carry PSP. First, policy-gradient loss weighting: the scalar action sum $a = \\sum_j a_j$ is differentiated with respect to input pixels to form $\\partial a/\\partial x_i$, which scales each pixel's reconstruction error (Eq. 1). Second, segmentation aggregation: a pretrained segmenter (SAM, or SAM2 variants) masks the image into objects, and each pixel weight becomes the mean absolute gradient inside its object, with a 99th-percentile clip and a 0.9/0.1 interpolation toward a uniform weight so the model still learns background early (Eqs. 2-3). Third, adversarial action prediction head: an MLP predicts the previous action from the encoder's embedding, and its scaled gradient ($\\epsilon = 10^3$) is subtracted from the world-model update, so the encoder has no incentive to store information about its own past action (Eqs. 4-5). These modifications apply only during training, leaving the inference-time DreamerV3 policy unchanged.","core_discovery":"The central discovery is that the gradient of the policy with respect to input pixels, once averaged inside segmentation masks, is a stable enough salience signal to reshape a world model's learning objective. PSP multiplies the per-pixel reconstruction loss of DreamerV3 by this policy-shaped weight, so the model must reproduce pixels the policy depends on even if the agent is small and uncertain, while the adversarial action prediction head removes the redundant reconstruction of self-generated visual changes. The paper reports that this combination achieves roughly twice DreamerV3's score on Reafferent Cheetah Run (383.1 vs 158.4), is the only model-based method tested with a nonzero mean on Reafferent Hopper Stand (128.5 vs at most 4.6 for the model-based baselines), stays within the normal range on unmodified Deepmind Control, and outperforms DreamerV3 on Distracting Control Hopper Stand (417.7 vs 173.7).","pith_inferences":["If the salience signal is the load-bearing part, a natural test is to shuffle the per-segment weights: random weights should severely degrade Reafferent Cheetah Run, while keeping the rest of PSP intact would show whether the policy-gradient signal is doing the work.","The same recipe could transfer to other reconstruction-based world models and to model-free agents that learn a state encoder from pixels, since nothing in the weighting or adversarial head depends on DreamerV3's RSSM specifics.","The adversarial head's fixed scale ($\\epsilon = 10^3$) is a hyperparameter the paper does not sweep; on tasks with higher-dimensional action spaces the gradient subtraction may need to be larger or learned, and sweeping it would show whether the efference-copy mechanism is robust.","Because segmentation is computed once per collected image and stored in the replay buffer, the method's cost can be shifted off the training loop entirely; using video segmentation models for temporal aggregation is a likely next step that could handle dynamic object boundaries."],"forward_implications":["PSP roughly doubles DreamerV3's score on Reafferent Cheetah Run (383.1 vs 158.4) and is the only tested model-based method with a nonzero mean on Reafferent Hopper Stand.","On unmodified Deepmind Control, PSP (Cheetah 712.3, Hopper 865.6) stays within the range of DreamerV3 and the leading baselines, so distraction suppression does not come at the cost of clean-environment performance.","On Distracting Control Suite, PSP outperforms DreamerV3 on both Cheetah Run (408.6 vs 243.8) and Hopper Stand (417.7 vs 173.7).","The method's modifications are training-only, so a PSP-trained agent runs the standard DreamerV3 inference pipeline with no extra segmentation or salience computation.","Segmentation quality matters most on hard tasks: switching from SAM to SAM2-tiny preserves Cheetah performance but loses ground on Hopper, and SAM2-large recovers most of the lost Distracting-Control score."],"supporting_citations":[{"why":"DreamerV3 is the base world-model architecture that PSP modifies and the primary baseline it must beat.","marker":"[Hafner et al., 2023]"},{"why":"VaGraM supplies the value-gradient-weighted model loss idea that PSP extends from low-dimensional states to image pixels.","marker":"[Voelcker et al., 2022]"},{"why":"Segment Anything is the pretrained segmentation model used to aggregate noisy pixel gradients into object-level weights.","marker":"[Kirillov et al., 2023]"},{"why":"Documents the noisiness of raw gradient saliency that motivates the object-aggregation step, and SmoothGrad is cited as an infeasible alternative.","marker":"[Smilkov et al., 2017]"},{"why":"Provides the efference-copy concept that motivates the adversarial action prediction head.","marker":"[Crapse and Sommer, 2008]"},{"why":"Domain-adversarial training is the technique the adversarial head uses to strip action information from the encoder.","marker":"[Ganin et al., 2016]"},{"why":"Distracting Control Suite is the standard distractor benchmark and the basis of the new Reafferent environment.","marker":"[Stone et al., 2021]"},{"why":"DreamerPro is a reconstruction-free MBRL baseline that PSP must outperform on distractors.","marker":"[Deng et al., 2022]"},{"why":"Task Informed Abstractions is a structural-regularization baseline designed for distractor robustness.","marker":"[Fu et al., 2021]"},{"why":"Denoised MDPs is a controllability-aware baseline; the paper shows it reconstructs backgrounds but fails on Reafferent tasks.","marker":"[Wang et al., 2022]"}],"fun_headline_variants":["World models learn to ignore pixel noise with policy-shaped losses","Policy-shaped losses make MBRL see past predictable distractions","Model-based RL focuses on what matters: policy-guided reconstruction","2x DreamerV3 on distracted Cheetah with policy-aware world models","Shaping world model loss by policy gradients beats distractors"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The whole method rests on the assumption that the policy's pixel gradients, averaged within segmented objects, mark the parts of the scene the world model should spend its capacity on, and that this marking stays trustworthy while the policy and world model are being learned together.","fun_headline_variants_meta":{"raw":{"variants":["World models learn to ignore pixel noise with policy-shaped losses","Policy-shaped losses make MBRL see past predictable distractions","Model-based RL focuses on what matters: policy-guided reconstruction","2x DreamerV3 on distracted Cheetah with policy-aware world models","Shaping world model loss by policy gradients beats distractors"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.00082,"raw_usage":{"total_tokens":3565,"prompt_tokens":893,"completion_tokens":2672,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":509,"completion_tokens_details":{"reasoning_tokens":2587}},"tokens_in":509,"tokens_out":2672,"duration_ms":17430,"temperature":1.0,"reasoning_tokens":2587,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-11T20:22:56.622884+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Train PSP on Reafferent Cheetah Run with per-segment loss weights replaced by random permutations of the true weights, keeping segmentation and the adversarial head intact: if the score stays near 383, the policy-gradient signal is not the cause of the gain, while a drop toward DreamerV3's 158 would confirm the salience weighting is doing the work.","supporting_citations":[],"review_version":1}