[Python] Gradient Lost When Constructing a Matrix in Mitsuba 3

Discussão em 'Python' iniciado por Stack, Fevereiro 11, 2025.

  1. Stack

    Stack Membro Participativo

    I am trying to optimize the absolute scale of an area light in Mitsuba 3.

    I'm doing a simple toy example where I start with a scene like this:

    [​IMG]

    And my target image is like this:

    [​IMG]

    My loss is simply MSE between the two images.

    I was able to get this to work by using a latent variable (a single float) that is turned into a uniform scaling matrix. However since this was iteratively multiplicative, the scale had a compounding exponential effect, making it difficult for convergence (by the time the scale value got big enough, the latent variable which represented the scale factor was so much larger than one that a linear learning rate adjustment couldn't get it small enough in time). It worked though if I stopped optimization before the optimized light got bigger than the reference light. Here's a MWE to illustrate:


    scene = mi.load_file("my_scene.xml")
    params = mi.traverse(scene)
    params.update()

    # Generate a reference image (area light at correct size)
    reference = mi.render(scene, params, spp=1024)

    # Disturb the scale by a factor of 0.5
    def multiplicative_resize(matrix, factor):
    """Applies local scale around the centroid (multiplicative)."""
    centroid = mi.Vector3f(matrix[0, 3], matrix[1, 3], matrix[2, 3])
    to_origin = mi.Transform4f.translate(-centroid)
    scale_tf = mi.Transform4f.scale(factor)
    from_origin = mi.Transform4f.translate(centroid)
    return from_origin @ scale_tf @ to_origin @ mi.Transform4f(matrix)

    light_key = "Light.to_world"
    original_matrix = params[light_key].matrix
    params[light_key] = multiplicative_resize(original_matrix, 0.5)
    params.update()

    optimizer = mi.ad.Adam(lr=0.01)
    optimizer["latent_scale_factor"] = mi.Float(1.001) # couldn't be exactly 1 otherwise there would be no gradient

    for it in range(10):
    # Render
    image = mi.render(scene, params, spp=1024)
    loss = dr.mean(dr.sqr(image - reference))

    dr.backward(loss)
    optimizer.step()

    # Clamp
    scale_val = dr.clamp(optimizer["latent_scale_factor"], 0.1, 2.0)
    optimizer["latent_scale_factor"] = scale_val

    # Update the transform multiplicatively
    current_matrix = params[light_key].matrix
    new_matrix = multiplicative_resize(current_matrix, scale_val)
    params[light_key] = mi.Transform4f(new_matrix)
    params.update()

    print(f"[Multiplicative] Iter {it}, loss = {loss[0]:.6f}")


    So, I figured that setting the absolute scale would make optimization easier. I tried to get the optimizer to learn how to match the absolute scale of the area light to a reference image:

    def set_area_light_scale(matrix, scale_value):
    """
    Sets the (0,0) and (1,1) entries to `scale_value`
    directly, ignoring previous scale.
    """
    # Suppose for a rectangular area light, the matrix
    # defaults to something like:
    # [ 2 0 0 x ]
    # [ 0 2 0 y ]
    # [ 0 0 2 z ]
    # [ 0 0 0 1 ]
    # Forcibly place `scale_value` in positions (0,0) and (1,1).
    new_mat = dr.llvm.ad.Matrix4f(
    scale_value, matrix[0,1], matrix[0,2], matrix[0,3],
    matrix[1,0], scale_value, matrix[1,2], matrix[1,3],
    matrix[2,0], matrix[2,1], matrix[2,2], matrix[2,3],
    matrix[3,0], matrix[3,1], matrix[3,2], matrix[3,3]
    )
    return mi.Transform4f(new_mat)

    scene = mi.load_file("my_scene.xml")
    params = mi.traverse(scene)
    params.update()

    reference = mi.render(scene, params, spp=1024)

    # Disturb the scale absolutely (initially 0.5, e.g.)
    light_key = "Light.to_world"
    original_matrix = params[light_key].matrix

    params[light_key] = set_area_light_scale(original_matrix, 0.5)
    params.update()

    optimizer = mi.ad.Adam(lr=0.01)
    optimizer["latent_scale_factor"] = mi.Float(0.5)

    for it in range(10):
    image = mi.render(scene, params, spp=1024)
    loss = dr.mean(dr.sqr(image - reference))

    dr.backward(loss)
    optimizer.step()

    # clamp
    scale_val = dr.clamp(optimizer["latent_scale_factor"], 0.1, 2.0)
    optimizer["latent_scale_factor"] = scale_val

    # forcibly set the scale
    current_tf = params[light_key].matrix
    params[light_key] = set_area_light_scale(current_tf, scale_val)
    params.update()

    print(f"[Absolute Scale] Iter {it}, loss = {loss[0]:.6f}")



    The new function properly updates the size of the area light, as I would expect. My issue however, is that now the loss is exactly the same every iteration (it's not going down). Am I breaking the gradient somewhere? What am I doing wrong?

    TL;DR:

    • The multiplicative scale transforms approach is properly working with gradient flows, but the change is exponential over the iterations.
    • The absolute scale approach yields exactly the same loss every iteration
    • I’d like to figure out why the gradient w.r.t. a newly assigned absolute scale doesn’t reduce the loss, and how to fix it. Thanks!

    Continue reading...

Compartilhe esta Página