taichi 多个 simulator steps AutoDiff 问题

综合讨论区
#1

请问如果一次计算调用涉及多个 simulator 的过程，即反复调用多个相同的 taichi kernel 时应该怎么 AutoDiff ？

问题在于，多次调用同一个 taichi kernel ，AutoDiff 时应该怎么做？

示例伪代码如下：（其中 RigidBody 是一个预定义类，每一个obj.xxx() 函数都是一个 taichi kernel）

请问下面的两种求 gradient 的方法哪一个正确，如果都不正确，请问怎么才能得到正确的 loss 对 obj_px 的导数？

obj = RigidBody(obj_px)		# 初始化刚体仿真变量，会将 obj_px 赋值给 obj.p_x
floor = RigidBody(floor_px)

sim_steps = 10
dt = 0.01
count_list = []
for step in range(sim_steps):
	obj.pre_simulation()		# 每一个 obj.xxx() 函数都是一个 taichi kernel，下同
    obj.free_drop(dt)			# 重力加速度下自由下落
    
    count = 0
    while True:
        collision_detection(obj, floor)
        update_obj_v_w()
        count = count + 1
        if (obj.contact_points == 0 or count>100):
            break
    count_list.append(count)

loss = L2loss(obj.p_x - obj_px)		# obj.p_x 是更新后的 obj 点云位置


# 需要求解 经过simulator的 loss对 obj_px 的导数
    
# backward method 1
# 按照调用顺序 对每个taichi kernel 求导
for step_idx in range(sim_step):
    grad_step_idx = sim_step - step_idx - 1
    count = count_list[grad_step_idx]
    
    for i in range(count):
        update_obj_v_w.grad()
        collision_detection.grad(obj, floor)
    
    obj.free_drop.grad(dt)
    obj.pre_simulation.grad()
output_grad = obj.p_x.grad()

    

# backward method 2
# 一次调用kernel 导数，认为不同 step grad 已经叠加在一起
update_obj_v_w.grad()
collision_detection.grad(obj, floor)
obj.free_drop.grad(dt)
obj.pre_simulation.grad()
output_grad = obj.p_x.grad()