Optimizing position of points in a loop

[h-OUS-e]

I made a little sample function, similar to the one you would find in an integrator, that shoots rays from a defined set of points, and traces them through the scene. Calling sample(ray_origins, ray_dir, paths) works, and it even works if I call it multiple times inside a for loop without an issue.

The problem arises when I try to apply a displacement texture map similar to the caustics optimization tutorial. It is basically the exact same function as the tutorial. I am trying to optimize the vertex of some mesh such that all the points at a particular depth value come closer to a point reference value. When I run:

for it in range(config['max_iterations']):
    apply_displacement()
    sampled = sample(ray_origins, ray_dir, paths)

I get the following error : Exception: loop_process_state(): one of the supplied loop state variables of type Float is attached to the AD graph (i.e., grad_enabled(..) is true). However, propagating derivatives through multiple iterations of a recorded loop is not supported (and never will be). Please see the documentation on differentiating loops for details and suggested alternatives.

The error disappears when I remove dr.enable_grad(paramsOpt['data']) in

def apply_displacement(amplitude = 1.):
    # Enforce reasonable range. For reference, the receiving plane
    # is 7 scene units away from the lens.
    vmax = 1 / 100.
    paramsOpt['data'] = dr.clamp(paramsOpt['data'], -vmax, vmax)
    dr.enable_grad(paramsOpt['data'])

    height_values = heightmap_texture.eval_1(lens_si)
    new_positions = (height_values * normals_initial * amplitude + positions_initial)
    params_scene['output.vertex_positions'] = dr.ravel(new_positions)
    params_scene.update()

How do I go about keeping the gradients inside a custom mi.loop to optimize a the vertices of a mesh? Seems like mi.render does that for us.

[njroussel]

I think the error message is pretty explicit: this is not something that is supported.

I'm not sure I understand your analogy to mi.render. This function has no loops: we're tracing all rays in parallel, not one after the other.

On the other hand sample() does usually have a recorded loop. However, in that loop we never update a parameter that is being differentiated. For example, you can differentiate a BSDF's texture, but we're definitely not changing its value as our ray's depth increases (the sample() loop).

To be clear, the error message here is about your loop in sample(), not your optimization loop.

[h-OUS-e]

Thank you for your quick answers as usual. Sorry this might be going over my head, but yes, it is was about my sample function, to which I was a bit confused since I am not doing any operations on the texture map inside the mi.loop of the sample function as the ray's depth increases. As far as I can tell, I wasn't updating any parameters of the texture inside the mi.loop.

Here is my sample function for reference:

def sample(ray_origins, ray_dir, paths):
    # Construct Ray Vectors 
    rays = mi.Ray3f(o=ray_origins, d=ray_dir)
    paths.reset() # reset paths to remove all previous recordings

    # Getting Sampler from Camera
    loaded_sensor = mi.load_dict(sensor)
    sampler = loaded_sensor.sampler().clone()

    # Primary Surface Intersection
    si = scene.ray_intersect(rays)

    # BSDF context
    bsdf_ctx = mi.BSDFContext()    
    depth_count = mi.UInt32(0)
    
    active = mi.Bool(True)

    loop = mi.Loop(name='Path Tracing', state=lambda: (sampler, rays, depth_count, active, si))
    loop.set_max_iterations(max_depth)
    print("Starting loop")

    while loop(si.is_valid() & (depth_count < max_depth) & active):
        # 1. Get BSDF from Surface Interaction and context
        bsdf = si.bsdf(rays)        
        # 2. Sample new ray directions
        bsdf_sample, bsdf_weight = bsdf.sample(bsdf_ctx, si, sampler.next_1d(),
                                                sampler.next_2d() )
         # 3. Spawn a new ray starting at the surface interactions
        rays2 = si.spawn_ray(si.to_world(bsdf_sample.wo))
        # # 4. Record points
        paths[depth_count] = PathType(rays2.o, rays2.d)
        # # 5. Calculate new surface interactions
        si = scene.ray_intersect(rays2)
        # # 6. Increase loop Iteration Count
        depth_count+=1
        
        active = (depth_count < max_depth) & si.is_valid()


    print("Loop Done")

I am only updating the parameters of the texture in the optimization loop using apply_displacement()

[h-OUS-e]

After some debugging, it seems that si is the reason for this error. I am assuming that is because I am updating a mesh in the scene when I do params_scene.update() in apply_displacement(). However, I need to put si in the loop to intersect the rays with the scene. How could I do that?
With the caustic optimization tutorial, since we choose a pre-built integrator, this step doesn't seem to be an issue since we also update a mesh in the scene.

Just to note, I am differentiating the BSDF texture, and updating the parameters of the scene in the optimization loop (outside the sample loop)