[Performance] Compile the DreamerV3 learner losses and make the two-hot decode reduction-safe

[theap06]

Summary

• add optimization.compile_learner to the DreamerV3 example: losses compiles the encoder, decoder, reward, continuation, actor and value networks together with the value and replay-value losses (numerically identical to eager); all additionally compiles the actor loss with its imagination rollout
• enable compile_learner=losses in the --fast reproduction
• make two_hot_decode exact for uniform probabilities under any reduction order by pairing mirrored bins through their probability difference
• cast the RSSM scan carry back to the incoming dtypes so the compiled scan can be traced by an outer torch.compile under bf16 autocast
• regression tests for both library changes

Why

After #4216 the fast path spends 95% of training wall time inside the captured learner update, and that update launches about 9,700 kernels with a 3.9 us mean. The compiled RSSM scan is already fused; the loss math around it and the imagination rollout still run as eager kernel streams. Compiling those regions removes most of the launches.

Measured on one NVIDIA GB10 (PyTorch 2.13.0, CUDA 13.0), DMC Walker preset, batch 16 x 64, BF16 compute, same protocol as benchmarks/ad_hoc/bench_dreamer_v3_learner.py:

Learner update	median ms	vs fast path
Fast path (compiled scan + CUDA graph)	27.8	baseline
+ compile_learner=losses	23.1	-17%
+ compile_learner=all	17.2	-38%
Upstream JAX DreamerV3, same GPU, same model	23.4

Unmodified train.py, 20,000 Walker steps on the same GPU: 33.5 updates/s for the fast path, 43.4 with losses (+30%), 52.0 with all (+55%).

Two bugs surfaced when compiling these regions:

• two_hot_decode relied on a fixed summation order of mirrored products so that uniform probabilities cancel to exactly zero over the +-4.85e8 symexp support. Inductor's fused reduction reorders it and zero logits decoded to 0.154 instead of 0, which shifted the imagined rewards and lambda targets at initialization. The new form (p_upper - p_lower) * b_upper + p_lower * (b_upper + b_lower) is exact under any order and general for any support.
• Under bf16 autocast the RSSM scan body returned bf16 state and belief while the initial carry was fp32. Dynamo's higher-order scan rejects that mismatch whenever an outer torch.compile inlines the compiled scan, which is what compiling the model loss does.

The actor network object is shared between the actor loss and the real-world policy, so with compile_learner set the collector and evaluation policies run the compiled actor MLP for their own batch shapes (dynamic=False, one graph per shape); the policy's encoder copy stays eager. This is the configuration the numbers above were measured with.

all moves the imagination's random draws inside the compiled region, like the scan backend, so it is not on by default; its actor loss after 60 updates differs from eager by float32 rounding at the tail bins (about -0.05 vs -0.0002). A full Walker learning curve with all is still owed before it becomes part of --fast.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/4290

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 09d4248 with merge base bedcc83 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.