Different computed gradients between enabling/disabling ti.static

[conceace]

Different computed gradients between enabling/disabling ti.static

When ti.static is enabled or disabled, the forward computation results are the same, but the gradients are different.

How to reproduce

Configurations:

• OS: Ubuntu 22.04.2 LTS
• Taichi version: 1.5.0
• Python version: 3.9.16
• CPU: AMD Ryzen Threadripper 3970X 32-Core Processor
• GPU: None

Code:

In the code below, the function collide is used to check collision between balls and walls. The function collide uses ti.static in the outermost for-loop. The bug would disappear after avoiding using ti.static in the outermost for-loop. The collision-checking functions collide_ball_pairs and collide_with_wall are adapted from the billiards.py in the official difftaichi repository.

(Sorry, the code is lengthy, but the bug could not be seen after disabling any of the ball-ball and ball-wall collision-checking functions.)

import taichi as ti


real = ti.f32
ti.init(arch=ti.cpu, debug=True, default_fp=real, flatten_if=False,
        advanced_optimization=False, cpu_max_num_threads=1)

max_num_steps = 100 # Used for allocating the capacity of arrays
dt = 0.01           # Time step is 0.01

state_size = 1      # The balls are only allowed to move on the x-axis.
n_balls = 2         # The total number of balls is 2

scalar = lambda: ti.field(dtype=real)
vec = lambda: ti.Vector.field(state_size, dtype=real)

loss = scalar()
x = vec()
x_inc = vec()  # for TOI
v = vec()
impulse = vec()
ctrls = vec()


radius = 0.1        # The radius of the two balls
elasticity = 1.0    # Fully elastic collision

ti.root.dense(ti.i, max_num_steps + 1).dense(ti.j, n_balls).place(x, v, x_inc, impulse)
ti.root.dense(ti.i, max_num_steps).dense(ti.j, n_balls).place(ctrls)
ti.root.place(loss)
ti.root.lazy_grad()


n_walls = 2                    # Two walls in total
wall_pos = ti.Vector.field(state_size, dtype=real, shape=(n_walls,))
wall_pos[0] = ti.Vector([-1])  # Wall 1 is located at -1
wall_pos[1] = ti.Vector([1])   # Wall 2 is located at 1


@ti.func
def collide_with_wall(t, ball_id, wall_id):
    # A function that checks collision between the ball `ball_id` and the wall `wall_id`
    # This function is slightly modified from the billiards.py in the official repository
    dist = (x[t, ball_id] + dt * v[t, ball_id]) - wall_pos[wall_id]
    dist_norm = dist.norm()
    if dist_norm < radius:
        dir = ti.Vector.normalized(dist, 1e-6)
        projected_v = dir.dot(v[t, ball_id])

        if projected_v < 0:
            imp = -(1 + elasticity) * projected_v * dir
            toi = (dist_norm - radius) / ti.min(
                -1e-3, projected_v)  # Time of impact
            x_inc_contrib = ti.min(toi - dt, 0) * imp
            x_inc[t + 1, ball_id] += x_inc_contrib
            impulse[t + 1, ball_id] += imp


@ti.func
def collide_ball_pairs(t, s_id, o_id):
    # A function that checks collision between the ball `s_id` and the ball `o_id`
    # This function is copied from the billiards.py in the official repository
    dist = (x[t, s_id] + dt * v[t, s_id]) - (x[t, o_id] + dt * v[t, o_id])
    dist_norm = dist.norm()
    rela_v = v[t, s_id] - v[t, o_id]
    if dist_norm < 2 * radius:
        dir = ti.Vector.normalized(dist, 1e-6)
        projected_v = dir.dot(rela_v)

        if projected_v < 0:
            imp = -(1 + elasticity) * 0.5 * projected_v * dir
            toi = (dist_norm - 2 * radius) / ti.min(
                -1e-3, projected_v)  # Time of impact
            x_inc_contrib = ti.min(toi - dt, 0) * imp
            x_inc[t + 1, s_id] += x_inc_contrib
            impulse[t + 1, s_id] += imp


@ti.kernel
def collide(t: ti.i32):
    # Check collision for all pairs of balls except a ball with itself
    for i in ti.static(range(n_balls)):
        for j in range(0, i):
            collide_ball_pairs(t, i, j)
    for i in ti.static(range(n_balls)):
        for j in range(i + 1, n_balls):
            collide_ball_pairs(t, i, j)
    # Check collision for ball-wall pairs.
    for i in ti.static(range(n_balls)):
        for j in range(n_walls):
            collide_with_wall(t, i, j)


# The bug would disappear after avoiding using ti.static in the outermost for-loop
# Comment out the function `collide` above and uncomment the code below to see the effects.
#
# @ti.kernel
# def collide(t: ti.i32):
#     for i in range(n_balls):
#         for j in range(0, i):
#             collide_ball_pairs(t, i, j)
#     for i in range(n_balls):
#         for j in range(i + 1, n_balls):
#             collide_ball_pairs(t, i, j)
#     for i in range(n_balls):
#         for j in range(n_walls):
#             collide_with_wall(t, i, j)


@ti.kernel
def advance_w_toi(t: ti.i32):
    for i in ti.static(range(n_balls)):
        v[t, i] = v[t - 1, i] + impulse[t, i] + ctrls[t - 1, i] * dt
        x[t, i] = x[t - 1, i] + dt * v[t, i] + x_inc[t, i]


def forward(num_steps: int):
    for t in range(1, num_steps + 1):
        collide(t - 1)
        advance_w_toi(t)  # from t - 1 to t


@ti.kernel
def compute_loss(num_steps: ti.int32):
    # A simple loss function that just adds up all the position and velocity values.
    for i in range(n_balls):
        x_diff = ti.abs(x[num_steps, i])
        for k in ti.static(range(state_size)):
            loss[None] += x_diff[k]
    for i in range(n_balls):
        v_diff = ti.abs(v[num_steps, i])
        for k in ti.static(range(state_size)):
            loss[None] += v_diff[k]


if __name__ == '__main__':
    # Just an initial setting that will cause the two balls to collide with each other
    # and the right-side wall during the simulation process.
    x[0, 0] = ti.Vector([0.6])    # Initial position of ball 1 is 0.6
    x[0, 1] = ti.Vector([0.89])   # Initial position of ball 2 is 0.89
    v[0, 0] = ti.Vector([1.0])    # Initial velocity of ball 1 is 1.0
    v[0, 1] = ti.Vector([0.5])    # Initial velocity of ball 2 is 0.5
    ctrls[0, 0] = ti.Vector([1.]) # The control applied on ball 1 during the first step is 1
    # The control force by default is 0 if not explicitly assigned.

    # Run forward simulation for 10 time steps and back propagate
    with ti.ad.Tape(loss):
        forward(10)
        compute_loss(10)

    # The forward simulation results is seen to be the same no matter whether ti.static
    # is used in the function `collide` or not.
    print("Final output loss value:", loss[None])

    # Print the gradients w.r.t. the control force on the two balls in the first time step.
    print("Gradients w.r.t the initial control:", ctrls.grad[0, 0], ctrls.grad[0, 1])

Output:

When directly running the code above, taichi would give the following output:

[Taichi] version 1.5.0, llvm 15.0.4, commit 7b885c28, linux, python 3.9.16
[Taichi] Starting on arch=x64
Final output loss value: 3.068909168243408
Gradients w.r.t the initial control: [-0.01100038] [-0.00900008]

After commenting the function collide in the code given above and uncommenting the version of function collide that does not use ti.static, the output would be:

[Taichi] version 1.5.0, llvm 15.0.4, commit 7b885c28, linux, python 3.9.16
[Taichi] Starting on arch=x64
Final output loss value: 3.068909168243408
Gradients w.r.t the initial control: [0.00899888] [0.00939929]

It can be seen that the forward simulation results, i.e., the loss value, is the same across the two versions of function collide. However, the gradients w.r.t. the initial control force on the two balls are different. The gradients w.r.t. the initial control force on the two balls are supposed to be the same, since ti.static shouldn't affect the outputs of gradient computation.