Skip to content

Add PPISP USD export, SPG controller rendering, and SH bake validation#243

Closed
moennen wants to merge 42 commits into
nv-tlabs:mainfrom
moennen:nicolasm/ppisp-controller-usd-export
Closed

Add PPISP USD export, SPG controller rendering, and SH bake validation#243
moennen wants to merge 42 commits into
nv-tlabs:mainfrom
moennen:nicolasm/ppisp-controller-usd-export

Conversation

@moennen
Copy link
Copy Markdown
Collaborator

@moennen moennen commented May 5, 2026

Summary

This PR adds PPISP-aware USD export support, including SPG-based post-processing assets, per-camera render products, PPISP controller export, and post-processing bake paths for SH output.
It also adds validation and diagnostic tooling for PPISP USD/SPG output, image comparison workflows, bake-mode benchmarking, NuRec USD import/export refinements, and LightField USD validation.

Key Changes

  • Add PPISP/SPG USD assets and writers for static, dynamic, and controller-driven post-processing.
  • Extend USD export with render product authoring, camera copy utilities, particle sorting hints, color-space controls, and SH post-processing bake support.
  • Add validation tools for PPISP SPG rendering, controller parity, trained exports, and image comparisons.
  • Add linear-to-srgb post-processing support and document/configure PPISP post-processing behavior.
  • Update export/import tests for color space handling, USDZ camera composition, and sorting hint validation.

Test Plan

  • Searched for debug/TODO/breakpoint leftovers in touched export and tool code.
  • Attempted focused test run:
    python -m pytest threedgrut/export/tests/test_export_import.py -v
  • Test run was blocked because pytest is not installed in the current environment.

@moennen moennen marked this pull request as draft May 5, 2026 20:37
moennen and others added 29 commits May 5, 2026 16:40
Decode NuRec USD/USDZ inputs into LightField export data so transcode can convert NuRec packages into standard ParticleField USD.

Made-with: Cursor
Run OpenUSD validation after standard USD export so composition and stage metadata errors are caught before returning success.

Made-with: Cursor
Add a no-parameter post-processing module for training and rendering when the renderer outputs linear RGB against sRGB targets.

Made-with: Cursor
Expose export_usd.linear_srgb so LightField exports can author lin_rec709_scene instead of the default display-referred color space.

Made-with: Cursor
Author per-camera RenderProducts and attach the PPISP SPG shader assets so standard USD exports can preserve the trained PPISP effect.

Made-with: Cursor
Add a dedicated writer that maps PPISP exposure and fitted post-processing parameters to Omniverse USD attributes for deployments without reliable SPG support.

Made-with: Cursor
Make export_ppisp the master PPISP export gate and use ov-post-processing only to select SPG, exposure fallback, fitted fallback, or combined SPG-plus-fallback output.

Made-with: Cursor
Keep cameras, render products, and PPISP post-processing authored where the packaged root stage composes them, and align standalone USD export with PPISP configuration.

Made-with: Cursor
Ensure checkpoint exports carry the PPISP SPG graph, camera timeline, and Gaussian render settings needed for Kit to render the post-processed LdrColor output.

Made-with: Cursor
Add an explicit omni-usd opt-in so default USD exports remain neutral while Kit-specific PPISP SPG and ParticleField MDL material authoring are only emitted when requested.

Made-with: Cursor
Default checkpoint exports to include loaded PPISP modules unless explicitly disabled, and scope neutral camera exposure to a hidden RenderProduct-local camera for PPISP SPG.

Made-with: Cursor
Add a generic post-processing SH bake path with a PPISP adapter so standard USD exports can bake a fixed post-processing transform into Gaussian coefficients.
Remove the separate omni-usd and ov-post-processing controls while keeping PPISP native export behind the explicit omni-native post-processing mode.
Add fixed camera/frame controls for PPISP omni-native export and extend bake validation with L1 fitting plus reference baseline outputs.
Add a Viser-based image comparison tool for paired images and matched folders, with visual diff modes and metrics support.
Clamp the power branch away from zero so validation image conversion remains finite.
Display metrics as readable text, add folder-level aggregates, and compute FLIP by default while removing LPIPS from the viewer.
Add simple fixed-frame PPISP baking variants to the validation workflow and keep the fitted bake path aligned with the current L2 objective.
Expose and validate ParticleField sortingModeHint values so USD exports can target ray-hit sorting while preserving the existing cameraDistance default.
Adds the per-camera PPISP controller (CNN + adaptive avg pool + 3-layer
MLP) to the Omniverse-native USD export so the runtime SPG pipeline
predicts exposure/colour at frame time instead of relying on
time-sampled USD attributes baked at training.

Pipeline authored on each RenderProduct:
    HdrColor -> PPISPController_<n> -> ControllerParams (1x9 float)
              -> PPISP (dyn) -> PPISPColor -> LdrColor
The controller weights flatten into a single 241,961-element float[]
attribute on the Shader prim; the SPG slang reads them via a
StructuredBuffer<float>. The dynamic PPISP variant
(ppisp_usd_spg_dyn.slang) reads exposure/colour from the controller's
1x9 output texture instead of from USD inputs.

Added a slangpy-based headless harness under tools/render_ppisp_spg/
that compiles and dispatches the same shaders without booting Kit, plus
several validators (synthetic, real ppisp module, trained checkpoint)
that compare the slang result against the in-process PyTorch PPISP
forward pass.

End-to-end on a trained bonsai checkpoint: PSNR 63 dB, controller
9-float drift ~3e-7 vs torch — within fp32 rounding of the
rgba8_unorm output format.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Lets a downstream consumer that doesn't want runtime controller dispatch
(e.g. for stricter portability or smaller assets) opt out of the
controller export even when the checkpoint contains trained controllers.
The exporter falls back to the static SPG path with time-sampled
exposure / color attributes derived from ppisp.exposure_params and
ppisp.color_params -- the same optimized per-frame parameters PPISP
would use when the controller branch is bypassed.

No effect on checkpoints trained without a controller. Available both as
a CLI flag (--ignore-ppisp-controller) and a YAML key
(export_usd.ignore-ppisp-controller).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
diagnose_controller.py runs three independent checks against a trained
checkpoint and pinpoints which layer is responsible when the controller-
driven render disagrees with the optimized-per-frame-params render:

  H1 -- PyTorch controller(rgb) vs trained exposure_params/color_params.
        Pure-python check; failure means the controller did not converge
        to the per-frame state during distillation.
  H2 -- slang controller vs PyTorch controller on the same HDR. Failure
        means the slang shader / weight flatten / buffer upload diverges
        from the trained module.
  H3 -- ppisp_usd_spg_dyn.slang (reads ControllerParams texture) vs
        ppisp_usd_spg.slang (USD attributes), both fed the same 9 floats.
        Failure means the SPG plumbing between the two shaders is broken.

Each section prints per-frame numbers plus a clear pass/fail interpretation
threshold, and the summary names the failing hypothesis. Saves having
to instrument the export end-to-end again the next time something
disagrees in Kit.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Previously the PPISP-dyn shader's inputs:ControllerParams was connected
directly to PPISPController_n.outputs:ControllerParams. That's valid
UsdShade -- the slangpy harness happily resolves it -- but Kit's SPG
runtime walks RenderProduct.orderedVars and resolves connections through
omni:rtx:aov on RenderVar prims, never directly between Shader prims. A
direct Shader -> Shader hop leaves the consumer unbound at dispatch time
(or bound to whatever Kit falls back to), which manifests as a much-too-
exposed render in Omniverse even though slangpy looks correct.

Mirror the existing HdrColor / LdrColor idiom: insert an intermediate
"ControllerParams" RenderVar with `opaque omni:rtx:aov.connect =
PPISPController_n.outputs:ControllerParams`, add it to orderedVars, and
point PPISP.inputs:ControllerParams.connect at the RenderVar's
omni:rtx:aov attribute.

Wiring is now:
  HdrColor (RenderVar+aov)
    -> PPISPController_n.inputs:HdrColor.connect = HdrColor.omni:rtx:aov
       PPISPController_n.outputs:ControllerParams
    -> ControllerParams (RenderVar) omni:rtx:aov.connect
                                    = PPISPController_n.outputs:ControllerParams
    -> PPISP.inputs:ControllerParams.connect = ControllerParams.omni:rtx:aov
       PPISP.outputs:PPISPColor
    -> LdrColor (RenderVar) omni:rtx:aov.connect = PPISP.outputs:PPISPColor

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Kit's SPG lua sandbox does a textual deny-list match for the lua
module-loader keyword 'require' and rejects any source containing it
-- including inside string literals. Both new lua launchers had assert
messages that read "... requires ...", which triggered the sandbox
validator and caused the entire shader graph build to fail with:

  LuaSandbox: Lua source rejected - forbidden pattern 'require' found
  Output outputs:ControllerParams of node PPISPController_0 is missing
  shape spec. Can't compute output shape.

(The shape error is a downstream effect: with the lua rejected, SPG
never executes outputs["ControllerParams"] = slang.empty(...) so the
shape can't be inferred.)

Replace "requires" with "needs" in both messages. Also did a quick
audit for other commonly denied tokens (dofile/loadfile/loadstring/
os./io./debug./setfenv/getfenv) -- both luas are clean.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Kit's SPG slang lua sandbox does not expose ``slang.StructuredBuffer``
(it's nil), so the previous launcher failed with "attempt to call a
nil value (field 'StructuredBuffer')". Replace the hard-coded call with
a small probe that tries every common HLSL/Slang buffer-resource name
in turn (StructuredBuffer, RWStructuredBuffer, Buffer, RWBuffer,
ByteAddressBuffer, RWByteAddressBuffer) and falls back to an explicit
error message that lists every key currently in the ``slang.*`` table.

If the probe doesn't find a match the user gets a precise log line
naming the available helpers, which we can use to settle on the right
one without further guessing.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The previous probe only listed pairs(slang), which under SPG's slang
table iterates only the directly-stored keys (the truncated 'short,
full, bo...' from the Kit log). The actual resource helpers like
Texture2D/dispatch live behind a __index metatable, so pairs() misses
them and the prior probe couldn't tell us anything useful.

Expand the candidate list to cover everything plausible (cased and
uncased forms: Buffer/buffer, Array/array, FloatArray/floatArray,
image/Image, uniform/Uniform, list/List, etc.) and test each via direct
field access (which goes through __index). Report both the direct
pairs() keys and the candidate list in the error if nothing matches.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Kit's SPG was emitting one warning per dispatch:

  [Warning] [rtx.spg.slang] SpgSlangNode: Failed to find parameter
  'params:weights' in shader reflection

because the slang shader declared the buffer as ``g_Weights`` while the
USD Shader prim exposes it as ``inputs:weights``. SPG's reflection-based
auto-binding looks up parameters by name; the mismatch caused the spam
even though the explicit lua bind in the dispatch worked.

Rename the slang variable to ``weights`` so the names match. Validation
still passes at PSNR 63 dB; max|Δ| in the controller's 9-float output
remains ~3.58e-7.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…rning

Kit's SPG resolves USD ``inputs:foo`` attributes against fields of the
slang ParameterBlock -- its reflection lookup is ``params:foo``. With
``weights`` declared as a sibling top-level ``StructuredBuffer<float>``
the lookup failed every dispatch:

  [Warning] [rtx.spg.slang] SpgSlangNode: Failed to find parameter
  'params:weights' in shader reflection

The static PPISP shader follows the same convention -- every attribute
(exposureOffset, vignetting*, crf*, ...) lives inside the PPISPParams
struct that's wrapped by ParameterBlock<PPISPParams>.

Move ``weights`` into PPISPControllerParams alongside priorExposure;
update both the lua bind list (now passes the buffer as a positional
argument to slang.ParameterBlock) and the slangpy harness's
ShaderCursor (cur["g_Params"]["weights"] instead of cur["weights"]).

Validation still passes at PSNR 63 dB; controller 9-float drift remains
~3.58e-7 vs torch.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The fit-by-bake export started Adam from the cloned checkpoint's SH
state, even though we have a closed-form one-shot bake (simple_bake)
that already lands close to the optimum. Use it as the initialization
so the fitting loop only refines the residual.

Adapter changes:

  - PostProcessingBakeAdapter.initialize_fit(model, pp) -- new hook,
    default no-op.
  - bake_post_processing_into_sh now calls adapter.initialize_fit on
    the cloned baked_model right before constructing the optimizer.
  - PPISPPostProcessingBakeAdapter.initialize_fit calls simple_bake
    with higher_order=True (so the spatial SH coefficients are
    Jacobian-projected too) and apply_srgb_to_linear=True.

The new srgb_to_linear flag on simple_bake matters for the colour-
space round-trip: PPISP outputs display-referred values (its CRF
folds in gamma-like encoding). Storing those directly as linear SH
coefs leaves the asset double-encoded once a downstream consumer
applies linear_to_srgb (the validator does, Kit's tonemap does).
Applying srgb_to_linear before RGB2SH puts SH back in linear scene-
referred space; downstream linear_to_srgb then undoes the encoding
exactly. Verified srgb_to_linear∘linear_to_srgb round-trips to fp32
epsilon (max|Δ| ≈ 2.4e-7) on [0, 1].

Existing 19 export tests pass.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Horde and others added 13 commits May 5, 2026 16:40
Adds two non-training view samplers to bake_post_processing_into_sh, on
top of the existing "iterate the train dataloader" mode:

* random-pair-slerp -- pick two distinct training views uniformly at
  random and slerp between them at random s in [0, 1]. No global
  structure but cheap.

* trajectory -- order the training views along an approximate
  Hamiltonian path (nearest-neighbour seed + 2-opt refinement) using
  a position+direction distance, arc-length-parameterise the path on
  [0, 1], and per step sample a random t and slerp inside the
  bracketing segment. Closer to the camera continuum a viewer would
  fly through; better generalisation than discrete training poses.

New module ``post_processing_view_interpolation``:
  - slerp_pose(pose_a, pose_b, s) -- quaternion slerp + translation lerp,
    fp32-clean to within fp32 epsilon at s=0 / s=1.
  - order_views_along_trajectory(poses, ...) -- NN + 2-opt with a
    weighted (position L2 + 1 - cos(forward angle)) metric. Position
    distances are mean-normalised so the rotation term lives on a
    comparable scale across scenes.
  - InterpolatedViewSampler -- wraps a template Batch and emits
    steps_per_epoch synthetic Batches with only T_to_world replaced.

bake_post_processing_into_sh now accepts:
  view_sampling_mode: "training" | "random-pair-slerp" | "trajectory"
  interpolated_views_seed: optional int RNG seed
  trajectory_weight_position / trajectory_weight_rotation: trajectory
    metric weights.

Wired through USDExporter constructor + from_config (YAML keys
post-processing-bake-view-mode / -view-seed / -trajectory-weight-{position,rotation})
and the export_usd.py CLI (--post-processing-bake-view-mode etc.).
Default remains "training" so existing exports are unchanged.

Sanity checked on a 6-pose synthetic circle: shuffled inputs reorder to
a Hamiltonian path with uniform 0.2 arc-length steps; existing 19 export
tests pass.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
PPISP forward is display-referred, so the fitted SH side must encode to
sRGB and clamp before MSE — without this the loss plateaued near 13 dB on
real scenes. Mirrors post_processing_sh_bake_validation.py::_fitBakedSh.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Two bugs in simple_bake's higher_order=True path under apply_srgb_to_linear=True:

1. The Jacobian was dPPISP/dX while the DC bake landed in linear space
   via srgb_to_linear; the chain rule was missing srgb_to_linear'(PPISP).
   The composed Jacobian now matches the DC color space.

2. ~0.06% of Gaussians at PPISP-saturation extremes have pathological
   Jacobians (cond > 1e8, |J|_F > 1e4) that pump features_specular norms
   from O(1) to O(10^4). These outliers dominate Adam's adaptive variance
   and stall the fit at ~29 dB. Clipping |J|_F > 5 (and any non-finite J)
   to identity preserves the rotation for the well-behaved 99.94% while
   leaving the trained specular intact for outliers.

On bonsai (1M Gaussians, 9-epoch fit), the higher-order PPISP warm-start
went from 25.4 dB to 36.4 dB -- now within 0.2 dB of the DC-only sRGB
warm-start (36.6 dB).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Empirical sweep showed random-pair sampling was always within noise of
training views (sometimes -0.8 dB) and never improved over the
trajectory sampler. The trajectory mode is retained for sparse-view
datasets where interpolating between adjacent training poses can help
generalisation.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Sweeps simple-bake and Adam-fit variants of the PPISP SH bake on a
trained checkpoint, reports mean / median / min / max PSNR (optional
SSIM, LPIPS) across the validation split, and writes per-frame numbers
to metrics.json. Used to root-cause the higher-order warm-start
regression (Jacobian chain rule + outlier clipping).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…ochs

Update PPISP SH-bake export defaults to the configuration that wins on
the bake-modes ablation:

* PPISPPostProcessingBakeAdapter.initialize_fit now uses
  higher_order=False -- DC-only simple_bake leaves the trained
  features_specular intact, which Adam fine-tunes faster than
  recovering from a Jacobian-rotated specular.
* post_processing_bake_view_mode default flipped from "training" to
  "trajectory" -- arc-length-sampled NN+2-opt path through training
  poses; +0.17 dB on bonsai at 15 epochs and expected to help more on
  sparse-view datasets.
* post_processing_bake_epochs default raised from 1 to 13 -- the fit
  is far from converged at 1, and 13 sits in the knee of the
  PSNR-vs-epochs curve (~36.5-37 dB on bonsai).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Two bugs in the PPISP SH bake fit, each found by ablation on caterpillar:

1. The SH was fit in linear scene-referred space against a display-referred
   target with a linear_to_srgb step in the loss path. The chain rule
   through linear_to_srgb is wildly non-uniform on [0, 1] (~13x at darks,
   ~0.4x at brights). Adam over-weights dark-region updates and the
   high-degree SH bands oscillate to fit amplified noise -- visible as
   rainbow fringing around silhouettes. Standard 3DGS doesn't see this
   because it trains in gamma space directly with identity gradient
   through the loss.

   Fix: warm-start albedo with simple_bake(apply_srgb_to_linear=False) and
   strip linear_to_srgb from apply_fit_transform. SH eval is now display-
   referred; gradient is identity through the loss; the asset format
   aligns with no-PPISP exports (gamma SH + no post-processing layer).

2. Fixed-geometry SH refit on a smooth synthetic target (PPISP forward of
   the trained model's prediction) gives the high-degree bands too much
   license to align coherently into aliasing patterns -- nothing in the
   loss surface breaks the symmetry. Standard 3DGS training avoids this
   via density / split-clone-prune adapting the placement to the target.

   Fix: open features_albedo, features_specular and density to Adam,
   each at its 3DGS-standard learning rate (2.5e-3, 1.25e-4, 5e-2).

On caterpillar (48 val frames, 13 epochs):
  before: mean 32.0 dB, worst 26.6 dB, strong rainbow on hard frames
  after:  mean 43.4 dB, worst 39.0 dB, rainbow gone

Default vignetting flipped from "achromatic-fit" to "none": the
chromatic-vs-achromatic mismatch was contributing ~1 dB of error and the
new asset format doesn't have a place to ship a runtime vignette anyway.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Match the new bake_post_processing_into_sh signature in the exporter and
CLI: separate albedo / specular / density learning rates, with the
3DGS-standard ratio (specular = albedo / 20) preserved by default.

* default --post-processing-bake-learning-rate flipped to 2.5e-3
* new --post-processing-bake-learning-rate-specular (default = albedo/20)
* new --post-processing-bake-learning-rate-density (default 5e-2)
* default --ppisp-bake-vignetting-mode flipped to "none"

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Drop the linear-SH + achromatic-vignette modes (fit-base, fit-base-srgb,
fit-init, ...) -- they no longer match the production fit path and were
only useful to debug the rainbow regression. Replace with a leaner
catalogue aligned with the new defaults:

  simple              one-shot DC-only bake (no fit, gamma SH)
  simple-higher-order one-shot DC + Jacobian-rotated specular (no fit)
  fit-color-only      Adam on albedo + specular only (density ablation)
  fit                 Adam on albedo + specular + density (production)
  fit-trajectory      fit + trajectory view sampling

Eval likewise simplifies: SH is already display-referred so the baked
side just clips to [0, 1]; reference uses no-vignette PPISP.

On caterpillar (48 val frames, 13 epochs) the 'fit' mode lands at
43.4 dB mean / 39.0 dB worst-case, vs 32.0 / 26.6 with the old fit-base.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The well-conditioned gamma-SH fit converges fast: per-step loss curves
on bonsai and caterpillar show diminishing returns past epoch 7. Cutting
the default ~halves wall-clock time per export with marginal PSNR cost
(rough estimate -0.5 dB mean, the asymptote is ~the same). Users wanting
the last fraction of a dB can still set --post-processing-bake-epochs.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
A single multiplicative knob applied to the asset's SH-evaluated RGB.
``features_specular`` is scaled linearly; ``features_albedo`` picks up an
extra ``(s - 1) * 0.5 / C0`` term to compensate for the constant offset
baked into ``RGB2SH`` so a forward eval reproduces ``s * original_rgb``.

Works uniformly across export modes (no-PPISP linear-SH, no-PPISP gamma,
PPISP baked-sh) -- the offset compensation is structural, not colour-
space-dependent. Default 1.0 keeps existing exports byte-identical.

Useful for matching downstream tonemap exposure or compensating for
runtime gain that the asset's consumer applies.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Adds a ``responsivityR/G/B`` float input on the PPISP SPG shader (both
the static and controller-aware variants) that is premultiplied to the
input HdrColor before exposure / vignetting / colour correction / CRF
run. Default 1.0 per channel keeps exported assets visually identical;
authors can scale per-channel sensitivity post-export by editing the USD
``inputs:responsivityR/G/B`` attributes on the per-camera shader prim.

Touches:
* both ``.slang`` shaders (struct field + premultiply at top of main)
* both ``.slang.lua`` launchers (bind the new fields, default 1.0)
* both ``.slang.usda`` schemas (declare inputs with default 1.0)
* ``ppisp_writer.py`` writes the inputs explicitly via
  ``_set_responsivity_params`` so the attributes round-trip cleanly.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Remove local Cursor agent metadata from the branch and fix stale PPISP plan references so the PR content matches the exported tooling.
@moennen moennen force-pushed the nicolasm/ppisp-controller-usd-export branch from 33e8f2a to 0e791a2 Compare May 5, 2026 20:41
@moennen moennen closed this May 6, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant