PaTaRM: Autoregressive Reward Rollouts

• PaTaRM is a method that combines autoregressive modeling with reward-based rollouts to generate high-valued sequential outputs and simulate future decision outcomes.
• It employs a two-stage training process, starting with teacher-guided ODE initialization followed by self-rollout optimization using direct reward feedback.
• Empirical results demonstrate competitive performance in video generation and bandit tasks, improving dynamic degree, aesthetic quality, and uncertainty estimation.

Autoregressive Generative Reward Rollouts (PaTaRM) are a class of methods that combine autoregressive modeling with reward-based rollouts for applications ranging from high-fidelity video generation to structured decision-making problems. These approaches utilize the sequential sampling properties of autoregressive models to either (i) autoregressively construct high-valued outputs guided explicitly by reward feedback, or (ii) simulate plausible, reward-weighted future outcomes for decision policies, generalizing the concept of Thompson sampling and model-based planning. PaTaRM methods replace or augment teacher supervision and explicit prior specification with a pipeline in which future reward sequences are generated conditionally, using the generative model both for inference and optimization (Zhang et al., 23 Jan 2026, Cai et al., 2024).

1. Mathematical Foundations

In the core paradigm, an autoregressive generative model $G_\theta$ parameterizes a factorized predictive distribution over sequential outputs. Given context—latent noise $z$ for generation, side information $x_{1:t}$ and actions $a_{1:t}$ for decision problems—the model defines

$$p_\theta(x_{1:T}\mid z) = \prod_{t=1}^T p_\theta(x_t\mid x_{<t}, z) \ p_\theta(y_{t+1:T}\mid x_{1:t}, a_{1:t}) = \prod_{k=t+1}^T p_\theta(y_k\mid x_{1:k-1}, a_{1:k}, y_{1:k-1})$$

In reward-driven training, as implemented in Reward-Forcing (Zhang et al., 23 Jan 2026), the objective transitions from maximum likelihood (MLE) or teacher-distillation to direct maximization of a reward function $R$ applied to the generated output (typically the last frame or outcome):

$$L_{\rm reward}(\theta) = -\,\mathbb{E}_z[R(G_\theta(z)_T)]$$

In decision-policy settings, the PaTaRM recipe (Cai et al., 2024) involves simulating hypothetical completions of missing rewards via autoregressive sampling. A population-mean estimate is formed by combining observed rewards and imputed future rewards to drive action selection.

2. Reward Function Specification and Application

Reward models in PaTaRM pipelines serve as differentiable critics or evaluators. In video generation (Zhang et al., 23 Jan 2026), the pretrained ImageReward network is used to assign scalar scores in $[0,1]$ to the last frame of a generated sequence. This approach avoids training auxiliary reward models and promotes consistent optimization. Empirical findings indicate that restricting $R$ to the final frame preserves temporal coherence, while applying it to earlier frames induces motion collapse ("freezing").

In bandit or sequential decision settings (Cai et al., 2024), the reward function coincides with task-specific outcomes (e.g., user clicks or binary feedback). The model learns to simulate missing future rewards, and the empirical mean of observed and simulated outcomes is used as the selection criterion:

$$\hat\mu_t^{(a)} = \frac{1}{T} \Big[\sum_{i=1}^{N^{(a)}} y_i^{(a)} + \sum_{i=N^{(a)}+1}^{T} R(\hat y_i^{(a)})\Big]$$

3. Rollout Mechanisms and Training Procedures

PaTaRM frameworks typically exhibit a two-stage training process. In Reward-Forcing (Zhang et al., 23 Jan 2026), the initial ODE initialization phase uses a bidirectional teacher to provide reverse trajectories, which the student autoregressive generator matches via supervised regression. Subsequent training ("self-rollout") involves sampling autoregressive sequences from the generator itself and optimizing reward directly via gradient-based updates, forming a fully differentiable, teacher-free pipeline. No Monte Carlo tree search or REINFORCE step is used.

For decision problems (Cai et al., 2024), historical sequence data is used to train the autoregressive model by next-outcome prediction. At inference, for each possible action (arm), the model generates hypothetical future reward sequences by ancestral (stepwise) sampling conditional on observed outcomes, then estimates action values from simulated rollouts.

Illustrative pseudocode is given in both works for the two-phase procedure:

Stage	Operation	Objective/Action
ODE Initialization	Sample latents, use teacher ODE to generate targets	Student regresses initial frame (video); bootstraps dynamics
Reward Feedback	Self-rollout, compute R on last frame, SGD update	Optimize expected reward over generator output
Rollout (decision)	For each arm, generate future rewards, select best	Implements Thompson sampling via model-based simulation

4. Model Architecture and Implementation Details

Architectures compatible with PaTaRM include any autoregressive sequence model. In video generation (Zhang et al., 23 Jan 2026), the backbone is a latent-space flow-matching U-Net (Wan2.1 T2V-1.3B) with multi-scale spatiotemporal attention, chunked autoregressive factorization, and relative positional encoding. The U-Net is equipped with causal masking and a KV-cache for efficient frame-chunk generation (≈17 FPS on H100 hardware).

For decision tasks (Cai et al., 2024), the model is an autoregressive network (transformer, RNN, or feed-forward) that incorporates side information (e.g., DistilBERT-encoded text) and a summary of past outcomes. Training incorporates bootstrap resampling, label smoothing, temperature scaling, and regularization to ensure both sharp predictions and reliable uncertainty estimates.

Sampling strategies include ancestral sampling for generative rollouts and truncated rollouts for computational efficiency.

5. Theoretical Properties and Regret Guarantees

PaTaRM in bandit settings admits precise characterization. Lemma 3.1 in (Cai et al., 2024) shows that, if the autoregressive model $p_\theta$ recovers the data-generating $p^*$, then the sampling-based estimate $\hat\mu_t^{(a)}$ and derived policy exactly implement population-level Thompson sampling. A finite-sample regret bound (Theorem 4.1) states that per-round regret is bounded by the sum of a standard Thompson sampling term and a penalty proportional to the model misspecification KL divergence.

$$\text{Regret} \leq \sqrt{\frac{|A| \log |A|}{2T}} + \sqrt{\frac{|A|}{2} \varepsilon}$$

where $\ell_T(p_\theta) - \ell_T(p^*) \leq \varepsilon$ denotes the next-outcome prediction loss gap.

6. Empirical Performance

In video generation benchmarks (Zhang et al., 23 Jan 2026), Reward-Only PaTaRM achieves competitive or superior results compared to state-of-the-art autoregressive and bidirectional models at comparable parameter scales. On VBench (832×480, 1.3B parameters, 0.69s latency), Reward-Only achieves a total score of 84.92, exceeding Self Forcing (84.31) and matching or exceeding bidirectional methods without heterogeneous distillation. Notably, Dynamic Degree (81.94 vs. 72.22) and Aesthetic Quality (70.51 vs. 65.75) are improved versus established baselines.

In synthetic and semi-real bandit tasks (Cai et al., 2024), flexible PS-AR (PaTaRM with neural networks) matches oracle regret performance and provides uncertainty intervals with correct coverage. Text-embedding based models outperform feature-agnostic or ensemble baselines on cumulative regret and interval calibration in news recommendation settings.

Model	Total (VBench)	Dynamic Degree	Aesthetic Quality
Reward Only	84.92	81.94	70.51
Self Forcing	84.31	72.22	65.75

7. Limitations and Future Directions

PaTaRM methods as presently instantiated possess several limitations. In video generation, reliance on image-level reward models lacking temporal awareness (e.g., ImageReward) can lead to motion collapse if applied to intermediate frames; use of fully video-based or VLM-based critics remains an open direction (Zhang et al., 23 Jan 2026). Stability issues inherent to direct reward optimization (potential for reward hacking, lack of formal RL guarantees) are unresolved. The two-stage training—teacher ODE pre-initialization followed by reward-driven rollouts—may be further optimized by curriculum scheduling or by integrating or replacing teacher initialization.

In meta-bandit contexts (Cai et al., 2024), model calibration and exchangeable sequence generation remain critical for robust uncertainty estimation. Truncation and computational limitations constrain rollout length in practice. A plausible implication is that leveraging richer side information and hybrid training pipelines could further enhance policy performance.

PaTaRM continues to provide a flexible, black-box, and scalable family of approaches for sequential generative modeling and decision-policy simulation, enabling efficient optimization of complex objectives without explicit posterior construction or heavy reliance on teacher models (Zhang et al., 23 Jan 2026, Cai et al., 2024).