Source-linked AI summary
ConfAL-WM: Confidence-Guided Active Learning for Action-Conditioned World Models
Xiang Liu, Sen Cui, Changshui Zhang
TL;DR
Action-conditioned world models can fail in localized spatiotemporal regions under distribution shifts, motivating more targeted post-training. ConfAL-WM adds a dense confidence probe to EVAC and uses its scores for staged data selection and localized retraining. On RoboTwin2.0, confidence-guided selection and frame/patch weighting improve reported prediction quality and embodied trajectory consistency over scalar scoring baselines.
Problem
Action-conditioned world-model errors under new task and scene distributions often concentrate in localized regions, limiting the value of globally adding post-training data.
Method
ConfAL-WM attaches a lightweight probe to EVAC UNet decoder features, aggregates dense confidence maps into task-, frame-, and patch-level scores, and uses them for staged selection and weighted retraining.
Results
Confidence-guided selection and frame/patch-weighted retraining improve prediction quality and embodied trajectory consistency over scalar reward, progress, and judge-based scoring baselines.
Takeaways & Limitations
Confidence can serve both as a risk signal for allocating post-training data and as a localized supervision weight for unreliable regions.
Takeaways & Limitations
The probe is tied to EVAC's UNet decoder features and trained only on RoboTwin2.0, limiting direct transfer across architectures and domains.
Abstract
from arXiv · showhide
Action-conditioned world models have become an important foundation for embodied prediction, planning, and synthetic data generation, but their errors under new task and scene distributions are often concentrated in localized spatiotemporal regions such as robot arms, manipulated objects, contact areas, and occluded objects. This paper presents ConfAL-WM, a confidence-guided active learning framework for post-training embodied world models. Built upon EVAC, we attach a lightweight confidence probe to UNet decoder features and predict dense confidence maps in the latent space. These maps are aggregated into task-, frame-, and patch-level scores, enabling both efficient data selection and localized training enhancement. Our pipeline first retrains the confidence probe and warms up EVAC with a small subset of target-domain data, then performs task-level prescreening to allocate sampling budgets, and finally applies selected-data retraining with optional frame or patch weighted data enhancement. Experiments on RoboTwin2.0 show that confidence-guided selection improves post-training efficiency, while dense frame and patch weighting further enhances prediction quality and embodied trajectory consistency compared with scalar reward, progress, and judge-based scoring baselines. A quick visual overview of this work is available at https://ConfAL-WM.github.io.
1 Introduction
ConfAL-WM targets localized prediction failures in action-conditioned world models with confidence-guided selection and dense retraining enhancement. Built on EVAC, it uses confidence at task, frame, and patch levels to improve post-training efficiency and prediction quality.
- Motivation: Prediction errors concentrate around robot arms, manipulated objects, contact regions, occlusions, and long-horizon interactions rather than being uniformly distributed.This motivates selecting and enhancing data according to likely failure regions instead of adding data globally.
- Method: ConfAL-WM attaches a lightweight confidence probe to EVAC UNet decoder features and predicts dense latent-space confidence maps.Decoder features preserve spatial locality while retaining global contextual information.
- Method: Confidence maps are aggregated into task-, frame-, and patch-level scores for data selection and localized retraining.The staged pipeline uses a small target-domain subset for probe retraining and EVAC warm-up before task-level prescreening and selected-data retraining.
- Method: Confidence-guided data enhancement weights frames and patches more strongly where supervision is needed.This focuses retraining on unreliable spatiotemporal regions rather than treating every video region equally.
- Results: Confidence improves data selection and weighted retraining compared with scalar reward, progress, and judge-based scoring methods.Experiments report improved prediction quality and embodied trajectory consistency under the EWMBench-style evaluation protocol.
2 Related Work
Related work spans action-conditioned world models, embodied evaluation, dense confidence and scalar scoring signals, localized supervision, and uncertainty-aware active learning. ConfAL-WM combines these directions by using confidence for both acquisition and patch-level retraining.
- Action-conditioned world models and evaluation: EVAC is a UNet-based action-conditioned world-model backbone pretrained on AgiBot World and post-trained on RoboTwin2.0 under new tasks, scenes, and embodiments.Related alternatives include IRASim, DINO-WM, and C3 for dense confidence estimation in DiT-based video models.
- Action-conditioned world models and evaluation: EWMBench evaluates embodied world models with reconstruction, scene, motion, and semantic metrics.It provides the main evaluation protocol used by this work.
- Dense confidence and scoring signals: Dense confidence methods estimate localized uncertainty, while ConfAL-WM additionally uses confidence as an acquisition score and local training weight.This connects confidence estimation to both data selection and retraining.
- Dense confidence and scoring signals: Scalar scoring methods estimate trajectory or frame quality through progress, rewards, preferences, or process-completion signals.The compared methods include GVL, RoboReward, Robometer, and LRMs, whereas ConfAL-WM also localizes patch errors.
- Localized supervision: Localized supervision methods emphasize action-relevant or physics-informative interaction regions through reconstruction, action-aware, or physics-consistency objectives.ConfAL-WM identifies such regions through world-model confidence and uses them for active selection and patch-weighted retraining.
- Active learning: Active learning and uncertainty-aware robotics methods select views, rollouts, or demonstrations using uncertainty, disagreement, or coverage.ConfAL-WM applies these principles to embodied world-model post-training.
3 Method
ConfAL-WM estimates dense latent-space prediction confidence from UNet decoder features, then uses it for multi-level acquisition and localized retraining. Its pipeline combines task-aware selection with optional frame- and patch-weighted loss enhancement.
- 3.1 Dense Confidence Probe for UNet World Models: The confidence probe estimates patch-level prediction reliability from intermediate UNet decoder features and produces dense confidence maps.Decoder features preserve spatial locality while retaining global contextual information.
- 3.1 Dense Confidence Probe for UNet World Models: Local prediction error is computed in latent space over channel-spatial entries and used to define binary confidence targets for future frames.The probe is trained with binary cross-entropy while the EVAC backbone remains frozen.
- 3.1 Dense Confidence Probe for UNet World Models: Decoder-layer selection trades coarse semantic context against finer spatial resolution when producing confidence predictions.The probe uses diffusion-timestep and sampled error-threshold embeddings together with channel, spatial, and temporal processing.
- 3.2 Confidence-Guided Active Learning: The pipeline warms up EVAC and the confidence probe on a small subset, prescreens tasks to allocate quotas, and then performs selection-only or confidence-weighted retraining.Selected episodes can receive additional confidence scoring before localized retraining.
- 3.2 Confidence-Guided Active Learning: Mean, tail, and persistent risk aggregate dense maps into acquisition signals emphasizing overall difficulty, severe localized failures, or recurring high-risk regions.Persistent risk averages the highest-risk patches within high-risk frames.
- 3.2 Confidence-Guided Active Learning: Frame weighting spreads spatially averaged risk across a frame, whereas patch weighting preserves dense local risk; α = 0 gives frame weighting and α = 1 recovers patch weighting.Risk is normalized, clipped, and converted into local loss multipliers whose strength increases during early retraining steps.
4 Experiments
Experiments evaluate whether confidence tracks prediction error and whether confidence-guided selection and localized retraining improve world-model quality under a fixed data budget. On RoboTwin2.0, confidence provides multi-scale risk signals and outperforms scalar selection baselines across most component metrics.
- Experimental setup: Experiments use RoboTwin2.0 to test confidence validity and active-learning gains in a transfer setting with new tasks, scenes, and robot embodiments.EVAC is pretrained on AgiBot World; the evaluation uses 50 tasks, 24,992 videos, and videos ranging from 98 to 578 frames.
- Why Confidence?: Risk maps localize failures around moving manipulators, object interactions, contacts, and temporary occlusions rather than reflecting only global video quality.Agreement with latent prediction-error maps indicates that the probe captures local failures.
- Why Confidence?: 0.540, 0.590, and 0.595 Spearman correlations are obtained at patch-, frame-, and task-levels, respectively, in latent space.For top-5% highest-error patches, AUROC is 0.761 and AUPRC is 0.146, versus random baselines of 0.5 and 0.05.
- Why Confidence?: Risk maps are temporally stable, with adjacent-frame top-region IoU of 0.740, flicker score of 0.005, and peak risk-error correlation of 0.602 near zero lag.Top-5% spatial IoU is 0.130, so confidence identifies error-prone regions without precisely reproducing their boundaries.
- Why Active Learning?: Confidence-guided mean-risk selection achieves the best result on eight of nine component metrics under selection-only retraining.These include PSNR, SSIM, Scene Consistency, Sem.-CLIP, Sem.-BLEU, and all three trajectory metrics.
- Why Active Learning?: Confidence with frame-and-patch weighting achieves the best PSNR, SSIM, Logics, Sem.-CLIP, Sem.-BLEU, Traj-HSD, Traj-Dyn, and Traj-nDTW among weighted methods.Frame weighting achieves the best Scene Consistency, while frame-and-patch weighting is slightly lower on that metric.
- Why Active Learning?: Mean risk achieves the best performance on seven of nine component metrics in the aggregation ablation and is used as the default.Tail risk has the highest Sem.-BLEU, while persistent risk has the highest Sem.-CLIP.
5 Conclusion
ConfAL-WM uses a decoder-feature confidence probe to support both task-level selection and localized retraining. Experiments on RoboTwin2.0 report improved reconstruction, scene consistency, and semantic quality over scalar scoring baselines.
- Conclusion: ConfAL-WM attaches a lightweight confidence probe to UNet decoder features for task-level selection and frame- and patch-level weighted retraining.The framework uses confidence as both an acquisition signal and a localized supervision signal.
- Conclusion: Experiments on RoboTwin2.0 show improved reconstruction, scene consistency, and semantic quality over scalar scoring baselines.The conclusion summarizes confidence as an effective risk signal for post-training.
- Limitations: The confidence probe is tied to a UNet diffusion backbone and trained only with EVAC on RoboTwin2.0, limiting direct transfer across architectures and domains.The authors suggest training a general-purpose world confidence model across multiple backbones and embodied datasets.
Impact Statement
The work targets more data-efficient and reliable embodied world-model post-training by focusing computation on informative data and unreliable regions.
- Impact Statement: ConfAL-WM focuses post-training computation on informative data and unreliable regions to improve data efficiency and reliability.The stated potential applications include robotic simulation, planning, and synthetic-data generation.
- Impact Statement: World-model predictions should not be treated as guaranteed physical outcomes, especially under distribution shifts or safety-critical deployment.This is an explicit scope boundary for applying the paper's improved predictions.
A.1 Patch-level Error Construction
The method constructs patch-level latent prediction errors from velocity predictions, aggregates channel and spatial errors, and converts them into thresholded confidence targets. EMA-adapted randomized thresholds stabilize supervision across changing error scales.
- Latent and patch error construction: Velocity-prediction errors are converted into equivalent clean-latent reconstruction errors without executing the full reverse diffusion process.At a fixed diffusion timestep, the two errors differ only by a scalar factor.
- Latent and patch error construction: Channel-wise absolute latent errors are averaged over channels and spatially pooled within regions aligned to confidence-probe patches.Each probe output location corresponds to a latent spatial region used to compute its patch-level error.
- Binary confidence target: Each local error is compared with an adaptive threshold, producing a binary target where 1 denotes a reliable patch and 0 denotes high error.The thresholded target supervises dense confidence predictions aligned with local patches.
- EMA-calibrated thresholding: Batch-wise error percentiles define lower and upper threshold bounds, which are updated with a two-stage EMA schedule.Faster warmup updates adapt to target-domain error scale, while later updates provide stronger smoothing.
- EMA-calibrated thresholding: Randomly sampling thresholds within the EMA band makes binary targets stochastic and induces a continuous confidence ordering that decreases with prediction error.Thresholds are sampled separately for future frames but shared across patches within each frame.
- EMA-calibrated thresholding: The threshold band rapidly contracts during early training and later tracks the target-domain error scale more slowly and stably.For RoboTwin2.0, the implementation initializes the band at 0.20–0.70 and updates it throughout probe training.
B.1 Additional Confidence Evaluation Details
Additional evaluations characterize how confidence aggregates risk, localizes errors, calibrates under operational thresholds, and responds to aggregation and conditioning choices. Mean aggregation is the most consistent global ranking signal, while spatial confidence is better suited to soft weighting than exact hard selection.
- Task-level aggregation: Mean risk has the strongest global Spearman correlation, while tail and persistent risk emphasize severe or recurring failure patterns.The comparison covers 50 prescreened tasks in latent and pixel spaces.
- Spatial localization: At top-5%, latent-space localization reaches approximately 0.13 IoU and 0.22 overlap, with slightly lower pixel-space values.Moderate IoU alongside higher overlap indicates recovery of relevant regions without exact boundary matching.
- Operational calibration: Calibration depends strongly on the operational correctness threshold because changing θeval alters correctness labels but not confidence maps.Latent-space ECE and Brier improve with looser thresholds, yet the probe remains overconfident near the main operating point.
- Parameter sensitivity: Mean frame and task aggregation gives the strongest Spearman correlations in both latent and pixel spaces, whereas extreme-value aggregation amplifies isolated noisy patches.This supports the mean aggregation used in the main evaluation.
- Parameter sensitivity: Low-to-moderate θcond values provide the strongest spatial discrimination, while large thresholds make most patches highly confident and reduce localization contrast.The default threshold lies near the best-performing range, and the overall ranking variation is moderate.
B.2 Additional Numerical Results for Active Learning
This section reports aggregated, bootstrap, and seed-wise active-learning results, with learned scoring methods averaged across three seeds and improvements evaluated relative to EVAC-v1.
- Aggregated results: Aggregated active-learning results report reconstruction, scene, motion, and semantics dimensions, with learned scoring values averaged over seeds 42, 3407, and 123.Table 4 presents each aggregate together with its relative change from EVAC-v1.
- Bootstrap evidence: Paired-bootstrap comparisons pool episode-level differences from all three seeds and report 95% confidence intervals with paired means.The bootstrap procedure uses 10,000 resamples for learned scoring methods.
- Seed-wise results: Tables 6–8 provide complete normalized component metrics separately for seeds 42, 3407, and 123, with seed 42 serving as the default reference.Base EVAC, EVAC-v1, and Random are repeated across the seed-wise tables.
B.3 Qualitative Evolution from Base EVAC to EVAC-v2
Qualitative and per-episode analyses show progression from zero-shot cross-embodiment failure through warmup adaptation to confidence-guided refinement. Frame-and-patch weighting most consistently improves local interactions, object persistence, and reconstruction quality.
- Per-episode progression: The qualitative comparisons cover object placement, cabinet interaction, block stacking, bowl stacking, and switch manipulation.The selected episodes expose a progression from cross-embodiment failure to adaptation and confidence-guided refinement.
- Per-episode progression: Across six representative episodes, Base EVAC performs substantially worse under zero-shot transfer, while EVAC-v1 recovers much of the target-domain quality.Confidence-guided EVAC-v2 further improves most dimensions, and frame-and-patch weighting achieves the strongest Reconstruction in all six episodes.
- Representative visual evolution: Frame weighting improves overall rollouts, but frame-and-patch weighting more clearly improves manipulated objects, gripper geometry, and fine interaction details.The latter concentrates learning on difficult local regions that remain problematic after frame weighting.
- Representative visual evolution: Frame-and-patch weighting correctly preserves tray contents and grasping in burger placement, unlike EVAC-v1 and frame-only weighting.The figure shows frame-only hallucination and deformation errors that are resolved by frame-and-patch weighting.
- Representative visual evolution: Frame-and-patch weighting maintains coherent boundary grasping and scene structure for empty-cup placement, while frame-only prediction deteriorates.The cup is grasped from the image boundary and placed on the plate.
- Representative visual evolution: For switch manipulation, frame-and-patch weighting further improves local arm geometry after frame weighting reduces color bias.Its predicted embodiment is substantially closer to ground truth.
- Detailed measurements: Detailed EWMBench measurements show further gains from confidence-guided EVAC-v2 across reconstruction, scene, semantic, and trajectory dimensions.Frame-and-patch weighting gives the strongest overall progression, particularly for semantic quality and Motion when valid trajectory signals exist.
B.4 More Results of Qualitative Confidence Visualization
Qualitative visualizations show that the confidence probe often localizes dominant prediction-error regions, while performance weakens when oracle errors are spatially diffuse.
- Patch-level Spearman correlations reached 0.713 and 0.695 in the two best-ranked episodes.Predicted risk maps concentrated around manipulators, operated objects, and interaction regions, agreeing broadly with pixel- and latent-space oracle errors.
- The predicted risk maps provide spatially localized signals rather than merely reflecting global video quality.
- Patch-level Spearman correlations were 0.241 and 0.265 in two challenging low-ranked episodes.The maps still responded to dominant high-error regions, although they did not precisely reproduce diffuse oracle-error boundaries.
- The best and worst cases illustrate both confidence localization capability and remaining limitations of the probe.Overall, confidence captures relative error severity and spatial concentration even when local oracle-error boundaries are not exactly reproduced.