SO3XR3 light pose optimization in combination with tabulated light yields NAN gradients with SPP !=1

[dnakath]

Dear Mitsuba Community,

I am looking for a solution for several weeks now, but cannot find the source of my problem.

I implemented a new light, which basically captures the radiant intensity distribution in form of a set of points which define the falloff curve, obtained by linear interpolation between them. It can be found here.

Float falloff_curve(const Vector3f &d, Mask /*active*/) const {
        using Index = dr::uint32_array_t<Float>;

        Vector3f local_dir = dr::normalize(d);
        Float theta    = dr::acos(local_dir.z());
        
        /*
        // find interval in ordered set
        UInt32 i = math::find_interval<UInt32>(
                    (uint32_t) m_tab_angles.size(),
                    [&](UInt32 idx) DRJIT_INLINE_LAMBDA {
                        return dr::gather<Float>(m_tab_angles, idx) <= theta;
                    });
        */
        Index i = math::find_interval<Index>(
                    (uint32_t) m_tab_angles.size(),
                    [&](Index idx) DRJIT_INLINE_LAMBDA {
                        return dr::gather<Float>(m_tab_angles, idx) <= theta;
                    });

        // linerarly interpolate RID
        Float weight = (theta - dr::gather<Float>(m_tab_angles, i)) / (dr::gather<Float>(m_tab_angles, i+1) - dr::gather<Float>(m_tab_angles, i));
        Float falloff =  dr::lerp(dr::gather<Float>(m_tab_values, i), dr::gather<Float>(m_tab_values, i+1), weight);
        //std::cout << "min falloff " <<  dr::min(falloff) << " ++ max falloff " <<  dr::max(falloff) << std::endl;
        return dr::select(theta <= m_cutoff_angle, falloff, 0.f);
}

In addition, I would like to optimize the pose of the light in an SO3XR3 optimization scheme. To this end, I do a latent optimization of the translation and the x,y-components of the log-rotation. I only optimize two components as the z-component is meaningless wrt an isotropic light. After the update by Adam, I use a SO3XR3 exponential map to write the 4X4 Transform4f back into the to_world parameters of the light.

If I do this with a spotlight, everything works fine. However, if I combine the pose optimization with the tablight, I can only use 1 spp. If I use a higher SPP, the gradients of the latent pose parameters become NAN eventually. Unfortunately, I cannot find the reason for that. Anybody has an idea? E.g., is the vectorization in find_interval problematic?

Please find code, which reproduces the problem in a minimum working example here. To make the MWE work, the whole repository has to be cloned and Mitsuba has to be recompiled, as the tablight is written in c++.

Thank you very much for your help!

[njroussel]

Hi @dnakath

Just to confirm, without doing the SO3XR3 scheme on your tablight, do you get NaNs too ?

Right before getting the NaNs in your foward rendering, I imagine that the gradient value of your latent variables are NaNs too, right? You can check this by printing dr.grad(opt[key]) before opt.step().

If your gradients on the latent variables are NaNs, the next step would be to call dr.backward(err, flags=dr.ADFlag.ClearNone). This will keep the gradients at intermediate values. So now you would print dr.grad(var) for any that is used to compute the to_world from the latent variable.
This should help you identify which operation in blowing up once differentiated.
If the to_world is actually already the source of the NaN then things get more complicated, but the same strategy could still be applied: trace back using ClearNone and dr.grad() until you find the operation that no longer has a NaN gradient.

[dnakath]

Hi @njroussel,

thank you very much for your reply.

Well, as shown in my minimal example, I render the ground truth images with the tablight and no error occurs here. Of course I don‘t do that much iterations in the forward pass. Also, when I do pose optimization with tablight and 1spp ( I do not have an intuition here, why this happens) or with the plain spotlight the error does not occur.

And you are right, the error is a NaN-gradient on the latent translational and rotational values R3 and so3, respectively. Those get propagated to to_world by the exponential mapping eventually.

To confirm, you basically recommend to do a manual torch.detect_anomaly by keeping all intermediate gradients and step through them until the first non-NaN gradient. And this is then the exploding operation, right?

I am currently not at my machine, but will walk through this process in ~~ 2 weeks and report back.

[merlinND]

One faster thing you could try first is to walk through the implementation of your new plugin and check for any potentially "unsafe" math, e.g.: sqrt(-1), acos(1.01), exp(very_large), etc.
By replacing each of those calls with a safe version (e.g. clamping the argument to a safe range), you might already find the source of your NaN gradients.

[dnakath]

Hi @merlinND ,

luckily your approach was sufficient to solve the problem.

A simple change from

Float theta = dr::acos(local_dir.z());

to

Float theta = dr::safe_acos(local_dir.z());

fixes the gradient problems. I still don't understand, why they occurred in the pose estimation and not in the light-parameters themselves - but its working now.

@njroussel thank you also for your suggestions -- I will apply them some time in the future for sure.

I wrongly suspected a vectorization problem in the find_interval function, but both versions seem to work (the commented as well as the not-commented one):

        using Index = dr::uint32_array_t<Float>;  
       [...]
        /*
        // find interval in ordered set
        UInt32 i = math::find_interval<UInt32>(
                    (uint32_t) m_tab_angles.size(),
                    [&](UInt32 idx) DRJIT_INLINE_LAMBDA {
                        return dr::gather<Float>(m_tab_angles, idx) <= theta;
                    });
        */
        Index i = math::find_interval<Index>(
                    (uint32_t) m_tab_angles.size(),
                    [&](Index idx) DRJIT_INLINE_LAMBDA {
                        return dr::gather<Float>(m_tab_angles, idx) <= theta;
                    });

Could you kindly elaborate as to why UInt32 as well as dr::uint32_array_t<Float> work as a vectorized index?

[njroussel]

Hi again @dnakath

Great to hear that you got it work!

The two types are equivalent. The helper dr::uint32_array_t<T> simply replaces the underlying scalar type of T by uint32_t. For example, T here is a Float which is a DiffArray<JitArray<Backend, float>> whose scalar type therefore is float. The trait above would create DiffArray<JitArray<Backend, uint32_t>> which is exactly UInt32.