search

taichi-dev / taichi

Productive, portable, and performant GPU programming in Python.
https://taichi-lang.org
Apache License 2.0
25.09k stars 2.26k forks source link

Incorrect results from AutoDiff #4337

Closed kracon7 closed 2 years ago

kracon7 commented 2 years ago

The auto differentiation result seems incorrect.

I was trying to implement a simple differentiable particle-based 2D physics engine with spring contact. In this minimal example, I made 2 circles and the left one has an initial velocity towards the right one. A loss is defined on the final step position difference and the loss is used to regress the particle mass. However, I see that mass_est doesn't converge to mass_gt. I'm a bit confused about such behavior. Thanks in advance for helping!

To Reproduce

# two circles collision

import os
import sys
import numpy as np
import taichi as ti

DTYPE = ti.f64

ti.init(arch=ti.cpu, debug=True)

@ti.data_oriented
class ParticleSimulator:
    def __init__(self, dt=1e-3):
        self.dt = dt
        self.max_step = 200

        # world bound
        self.wx_min, self.wx_max, self.wy_min, self.wy_max = -30, 30, -30, 30

        # render resolution
        self.resol_x, self.resol_y = 800, 800
        self.gui = ti.GUI('billards', (self.resol_x, self.resol_y))

        self.ngeom = 2

        self.geom_mass = ti.field(DTYPE, self.ngeom, needs_grad=True)

        # contact parameters
        self.ks = 1e4     # spring constant for contact

        # pos, vel and force of each particle
        self.geom_pos = ti.Vector.field(2, DTYPE, shape=(self.max_step, self.ngeom), needs_grad=True)
        self.geom_vel = ti.Vector.field(2, DTYPE, shape=(self.max_step, self.ngeom), needs_grad=True)
        self.geom_force = ti.Vector.field(2, DTYPE, shape=(self.max_step, self.ngeom), needs_grad=True)

        # radius of the particles
        self.radius = ti.field(DTYPE, self.ngeom)
        self.radius.from_numpy(2*np.ones(self.ngeom))   # circle radius is 2

    @ti.kernel
    def collide(self, s: ti.i32):
        '''
        compute particle force based on pair-wise collision
        '''
        for i in range(self.ngeom):
            for j in range(self.ngeom):
                if i != j:
                    pi, pj = self.geom_pos[s, i], self.geom_pos[s, j]
                    r_ij = pj - pi
                    r = r_ij.norm()
                    n_ij = r_ij / r      # normal direction
                    if (self.radius[i] + self.radius[j] - r) > 0:
                        # spring force
                        fs = - self.ks * (self.radius[i] + self.radius[j] - r) * n_ij  
                        self.geom_force[s, i] += fs

    @ti.kernel
    def update(self, s: ti.i32):
        # update the next step velocity and position
        for i in range(self.ngeom):
            self.geom_vel[s+1, i] = self.geom_vel[s, i] + \
                                     self.dt * self.geom_force[s, i] / self.geom_mass[i]
            self.geom_pos[s+1, i] = self.geom_pos[s, i] + \
                                     self.dt * self.geom_vel[s+1, i]

    def initialize(self):
        # clear states and set the simulation scene
        self.clear_all()
        self.set_scene()

    @ti.kernel
    def clear_all(self):
        # clear all position velocity and gradients
        for s, i in ti.ndrange(self.max_step, self.ngeom):
            self.geom_pos[s, i] = [0., 0.]
            self.geom_vel[s, i] = [0., 0.]
            self.geom_force[s, i] = [0., 0.]
            self.geom_pos.grad[s, i] = [0., 0.]
            self.geom_vel.grad[s, i] = [0., 0.]
            self.geom_force.grad[s, i] = [0., 0.]

        for i in range(self.ngeom):
            self.geom_mass.grad[i] = 0

    @ti.kernel
    def set_scene(self):
        # place two circles in the scene and add initial velocity to the first circle
        self.geom_pos[0, 0] = [0, 0]
        self.geom_pos[0, 1] = [4, 0]
        self.geom_vel[0, 0][0] = 100  # vx
        self.geom_vel[0, 0][1] = 0  # vy

    def render(self, s):  # Render the scene on GUI
        np_pos = self.geom_pos.to_numpy()[s]
        # print(np_pos[:10])
        np_pos = (np_pos - np.array([self.wx_min, self.wy_min])) / \
                 (np.array([self.wx_max-self.wx_min, self.wy_max-self.wy_min]))

        # composite object
        r = self.radius[0] * self.resol_x / (self.wx_max-self.wx_min)
        self.gui.circles(np_pos, color=0xffffff, radius=r)

        self.gui.show()

@ti.data_oriented
class Loss():
    """docstring for Loss"""
    def __init__(self, engines):
        self.loss = ti.field(DTYPE, shape=(), needs_grad=True)
        self.engines = engines

    @ti.kernel
    def clear_loss(self):
        self.loss[None] = 0

    @ti.kernel
    def compute_loss(self, t: ti.i32):
        for i in range(self.engines[0].ngeom):
            self.loss[None] += (self.engines[0].geom_pos[t, i][0] - self.engines[1].geom_pos[t, i][0])**2 + \
                    (self.engines[0].geom_pos[t, i][1] - self.engines[1].geom_pos[t, i][1])**2

def run_world(sim):
    for s in range(sim.max_step-1):
        sim.collide(s)
        sim.update(s)

def forward():
    run_world(sim_est)
    run_world(sim_gt)
    loss.compute_loss(199)

# simulation engine with ground truth and estimated mass
sim_gt, sim_est = ParticleSimulator(), ParticleSimulator()
mass_gt, mass_est = np.ones(sim_est.ngeom), np.array([0.8, 1])
sim_gt.geom_mass.from_numpy(mass_gt)
sim_est.geom_mass.from_numpy(mass_est)

loss = Loss((sim_est, sim_gt))

lr = 1e-3
max_iter = 100

### regress the mass based on position difference
for i in range(max_iter):
    sim_gt.initialize()
    sim_est.initialize()
    loss.clear_loss()

    # forward sims to compute loss and gradient
    with ti.Tape(loss.loss):
        forward()

    print('Iteration %d loss: %12.4f   mass_est: %12.4f   mass_gt: %8.4f   gradient: %.4f'%(
            i, loss.loss[None], sim_est.geom_mass.to_numpy()[0], mass_gt[0], sim_est.geom_mass.grad[0]))

    grad = sim_est.geom_mass.grad

    # update estimated mass
    mass_est -= lr * grad.to_numpy()
    sim_est.geom_mass.from_numpy(mass_est)

# ######################   render simulation  ######################
# ###  render one episode

# sim = ParticleSimulator()
# mass = np.array([0.8, 1])
# sim.geom_mass.from_numpy(mass)

# sim.initialize()
# for s in range(sim.max_step-1):
#     sim.collide(s)
#     sim.update(s)
#     sim.render(s)

Log/Screenshots

$ python tests/test_sys_id_3.py 
[Taichi] version 0.8.7, llvm 10.0.0, commit 88d81df6, linux, python 3.6.9
[Taichi] Starting on arch=x64
[W 02/20/22 01:45:27.741 9879] [type_check.cpp:visit@147] [$8191] Local store may lose precision (target = f32, value = f64), at
8191
[W 02/20/22 01:45:27.882 9879] [type_check.cpp:visit@147] [$9386] Local store may lose precision (target = f32, value = f64), at
9386
[W 02/20/22 01:45:28.015 9879] [type_check.cpp:visit@147] [$10635] Local store may lose precision (target = f32, value = f64), at
10635
[W 02/20/22 01:45:28.154 9879] [type_check.cpp:visit@147] [$11830] Local store may lose precision (target = f32, value = f64), at
11830
Iteration 0 loss:       8.7899   mass_est:       0.8000   mass_gt:   1.0000   gradient: -98.7421
Iteration 1 loss:       0.1279   mass_est:       0.8987   mass_gt:   1.0000   gradient: -10.3719
Iteration 2 loss:       0.0078   mass_est:       0.9091   mass_gt:   1.0000   gradient: -1.4521
Iteration 3 loss:       0.0053   mass_est:       0.9106   mass_gt:   1.0000   gradient: -0.1959
Iteration 4 loss:       0.0053   mass_est:       0.9108   mass_gt:   1.0000   gradient: -0.0167
Iteration 5 loss:       0.0053   mass_est:       0.9108   mass_gt:   1.0000   gradient: 0.0090
Iteration 6 loss:       0.0053   mass_est:       0.9108   mass_gt:   1.0000   gradient: 0.0127
Iteration 7 loss:       0.0053   mass_est:       0.9108   mass_gt:   1.0000   gradient: 0.0132
Iteration 8 loss:       0.0053   mass_est:       0.9107   mass_gt:   1.0000   gradient: 0.0134
Iteration 9 loss:       0.0053   mass_est:       0.9107   mass_gt:   1.0000   gradient: 0.0134
Iteration 10 loss:       0.0053   mass_est:       0.9107   mass_gt:   1.0000   gradient: 0.0134
Iteration 11 loss:       0.0053   mass_est:       0.9107   mass_gt:   1.0000   gradient: 0.0135
Iteration 12 loss:       0.0053   mass_est:       0.9107   mass_gt:   1.0000   gradient: 0.0135
Iteration 13 loss:       0.0053   mass_est:       0.9107   mass_gt:   1.0000   gradient: 0.0136
Iteration 14 loss:       0.0053   mass_est:       0.9107   mass_gt:   1.0000   gradient: 0.0136
Iteration 15 loss:       0.0053   mass_est:       0.9106   mass_gt:   1.0000   gradient: 0.0136
Iteration 16 loss:       0.0053   mass_est:       0.9106   mass_gt:   1.0000   gradient: 0.0137
Iteration 17 loss:       0.0053   mass_est:       0.9106   mass_gt:   1.0000   gradient: 0.0137
Iteration 18 loss:       0.0053   mass_est:       0.9106   mass_gt:   1.0000   gradient: 0.0138
Iteration 19 loss:       0.0053   mass_est:       0.9106   mass_gt:   1.0000   gradient: 0.0138
Iteration 20 loss:       0.0053   mass_est:       0.9106   mass_gt:   1.0000   gradient: 0.0138
Iteration 21 loss:       0.0053   mass_est:       0.9106   mass_gt:   1.0000   gradient: 0.0139
Iteration 22 loss:       0.0053   mass_est:       0.9106   mass_gt:   1.0000   gradient: 0.0139
Iteration 23 loss:       0.0053   mass_est:       0.9105   mass_gt:   1.0000   gradient: 0.0140
Iteration 24 loss:       0.0053   mass_est:       0.9105   mass_gt:   1.0000   gradient: 0.0140
Iteration 25 loss:       0.0053   mass_est:       0.9105   mass_gt:   1.0000   gradient: 0.0140
Iteration 26 loss:       0.0053   mass_est:       0.9105   mass_gt:   1.0000   gradient: 0.0141
Iteration 27 loss:       0.0053   mass_est:       0.9105   mass_gt:   1.0000   gradient: 0.0141
Iteration 28 loss:       0.0053   mass_est:       0.9105   mass_gt:   1.0000   gradient: 0.0142
Iteration 29 loss:       0.0053   mass_est:       0.9105   mass_gt:   1.0000   gradient: 0.0142
Iteration 30 loss:       0.0053   mass_est:       0.9104   mass_gt:   1.0000   gradient: 0.0142
Iteration 31 loss:       0.0053   mass_est:       0.9104   mass_gt:   1.0000   gradient: 0.0143
Iteration 32 loss:       0.0053   mass_est:       0.9104   mass_gt:   1.0000   gradient: 0.0143
Iteration 33 loss:       0.0053   mass_est:       0.9104   mass_gt:   1.0000   gradient: 0.0144
Iteration 34 loss:       0.0053   mass_est:       0.9104   mass_gt:   1.0000   gradient: 0.0144
Iteration 35 loss:       0.0053   mass_est:       0.9104   mass_gt:   1.0000   gradient: 0.0144
Iteration 36 loss:       0.0053   mass_est:       0.9104   mass_gt:   1.0000   gradient: 0.0145
Iteration 37 loss:       0.0053   mass_est:       0.9103   mass_gt:   1.0000   gradient: 0.0145
Iteration 38 loss:       0.0052   mass_est:       0.9103   mass_gt:   1.0000   gradient: 0.0146
Iteration 39 loss:       0.0052   mass_est:       0.9103   mass_gt:   1.0000   gradient: 0.0146
Iteration 40 loss:       0.0052   mass_est:       0.9103   mass_gt:   1.0000   gradient: 0.0146
Iteration 41 loss:       0.0052   mass_est:       0.9103   mass_gt:   1.0000   gradient: 0.0147
Iteration 42 loss:       0.0052   mass_est:       0.9103   mass_gt:   1.0000   gradient: 0.0147
Iteration 43 loss:       0.0052   mass_est:       0.9103   mass_gt:   1.0000   gradient: 0.0148
Iteration 44 loss:       0.0052   mass_est:       0.9102   mass_gt:   1.0000   gradient: 0.0148
Iteration 45 loss:       0.0052   mass_est:       0.9102   mass_gt:   1.0000   gradient: 0.0149
Iteration 46 loss:       0.0052   mass_est:       0.9102   mass_gt:   1.0000   gradient: 0.0149
Iteration 47 loss:       0.0052   mass_est:       0.9102   mass_gt:   1.0000   gradient: 0.0149
Iteration 48 loss:       0.0052   mass_est:       0.9102   mass_gt:   1.0000   gradient: 0.0150
Iteration 49 loss:       0.0052   mass_est:       0.9102   mass_gt:   1.0000   gradient: 0.0150
Iteration 50 loss:       0.0052   mass_est:       0.9101   mass_gt:   1.0000   gradient: 0.0151
Iteration 51 loss:       0.0052   mass_est:       0.9101   mass_gt:   1.0000   gradient: 0.0151
Iteration 52 loss:       0.0052   mass_est:       0.9101   mass_gt:   1.0000   gradient: 0.0152
Iteration 53 loss:       0.0052   mass_est:       0.9101   mass_gt:   1.0000   gradient: 0.0152
Iteration 54 loss:       0.0052   mass_est:       0.9101   mass_gt:   1.0000   gradient: 0.0152
Iteration 55 loss:       0.0052   mass_est:       0.9101   mass_gt:   1.0000   gradient: 0.0153
Iteration 56 loss:       0.0052   mass_est:       0.9101   mass_gt:   1.0000   gradient: 0.0153
Iteration 57 loss:       0.0052   mass_est:       0.9100   mass_gt:   1.0000   gradient: 0.0154
Iteration 58 loss:       0.0052   mass_est:       0.9100   mass_gt:   1.0000   gradient: 0.0154
Iteration 59 loss:       0.0052   mass_est:       0.9100   mass_gt:   1.0000   gradient: 0.0155
Iteration 60 loss:       0.0052   mass_est:       0.9100   mass_gt:   1.0000   gradient: 0.0155
Iteration 61 loss:       0.0052   mass_est:       0.9100   mass_gt:   1.0000   gradient: 0.0156
Iteration 62 loss:       0.0052   mass_est:       0.9100   mass_gt:   1.0000   gradient: 0.0156
Iteration 63 loss:       0.0052   mass_est:       0.9099   mass_gt:   1.0000   gradient: 0.0156
Iteration 64 loss:       0.0052   mass_est:       0.9099   mass_gt:   1.0000   gradient: 0.0157
Iteration 65 loss:       0.0052   mass_est:       0.9099   mass_gt:   1.0000   gradient: 0.0157
Iteration 66 loss:       0.0052   mass_est:       0.9099   mass_gt:   1.0000   gradient: 0.0158
Iteration 67 loss:       0.0052   mass_est:       0.9099   mass_gt:   1.0000   gradient: 0.0158
Iteration 68 loss:       0.0052   mass_est:       0.9099   mass_gt:   1.0000   gradient: 0.0159
Iteration 69 loss:       0.0052   mass_est:       0.9099   mass_gt:   1.0000   gradient: 0.0159
Iteration 70 loss:       0.0052   mass_est:       0.9098   mass_gt:   1.0000   gradient: 0.0160
Iteration 71 loss:       0.0052   mass_est:       0.9098   mass_gt:   1.0000   gradient: 0.0160
Iteration 72 loss:       0.0052   mass_est:       0.9098   mass_gt:   1.0000   gradient: 0.0160
Iteration 73 loss:       0.0052   mass_est:       0.9098   mass_gt:   1.0000   gradient: 0.0161
Iteration 74 loss:       0.0052   mass_est:       0.9098   mass_gt:   1.0000   gradient: 0.0161
Iteration 75 loss:       0.0052   mass_est:       0.9098   mass_gt:   1.0000   gradient: 0.0162
Iteration 76 loss:       0.0052   mass_est:       0.9097   mass_gt:   1.0000   gradient: 0.0162
Iteration 77 loss:       0.0052   mass_est:       0.9097   mass_gt:   1.0000   gradient: 0.0163
Iteration 78 loss:       0.0052   mass_est:       0.9097   mass_gt:   1.0000   gradient: 0.0163
Iteration 79 loss:       0.0052   mass_est:       0.9097   mass_gt:   1.0000   gradient: 0.0164
Iteration 80 loss:       0.0052   mass_est:       0.9097   mass_gt:   1.0000   gradient: 0.0164
Iteration 81 loss:       0.0052   mass_est:       0.9097   mass_gt:   1.0000   gradient: 0.0165
Iteration 82 loss:       0.0052   mass_est:       0.9096   mass_gt:   1.0000   gradient: 0.0165
Iteration 83 loss:       0.0052   mass_est:       0.9096   mass_gt:   1.0000   gradient: 0.0166
Iteration 84 loss:       0.0052   mass_est:       0.9096   mass_gt:   1.0000   gradient: 0.0166
Iteration 85 loss:       0.0052   mass_est:       0.9096   mass_gt:   1.0000   gradient: 0.0167
Iteration 86 loss:       0.0052   mass_est:       0.9096   mass_gt:   1.0000   gradient: 0.0167
Iteration 87 loss:       0.0052   mass_est:       0.9096   mass_gt:   1.0000   gradient: 0.0167
Iteration 88 loss:       0.0052   mass_est:       0.9095   mass_gt:   1.0000   gradient: 0.0168
Iteration 89 loss:       0.0052   mass_est:       0.9095   mass_gt:   1.0000   gradient: 0.0168
Iteration 90 loss:       0.0052   mass_est:       0.9095   mass_gt:   1.0000   gradient: 0.0169
Iteration 91 loss:       0.0052   mass_est:       0.9095   mass_gt:   1.0000   gradient: 0.0169
Iteration 92 loss:       0.0052   mass_est:       0.9095   mass_gt:   1.0000   gradient: 0.0170
Iteration 93 loss:       0.0052   mass_est:       0.9095   mass_gt:   1.0000   gradient: 0.0170
Iteration 94 loss:       0.0052   mass_est:       0.9094   mass_gt:   1.0000   gradient: 0.0171
Iteration 95 loss:       0.0052   mass_est:       0.9094   mass_gt:   1.0000   gradient: 0.0171
Iteration 96 loss:       0.0052   mass_est:       0.9094   mass_gt:   1.0000   gradient: 0.0172
Iteration 97 loss:       0.0052   mass_est:       0.9094   mass_gt:   1.0000   gradient: 0.0172
Iteration 98 loss:       0.0052   mass_est:       0.9094   mass_gt:   1.0000   gradient: 0.0173
BisonLeo commented 2 years ago

If you change the 'update' with previous vel, it'll converge.

    @ti.kernel
    def update(self, s: ti.i32):
        # update the next step velocity and position
        for i in range(self.ngeom):
            self.geom_vel[s+1, i] = self.geom_vel[s, i] + \
                                     self.dt * self.geom_force[s, i] / self.geom_mass[i]
            self.geom_pos[s+1, i] = self.geom_pos[s, i] + \
                                     self.dt * self.geom_vel[s, i]     # <<<< change to use previous velocity
kracon7 commented 2 years ago

If you change the 'update' with previous vel, it'll converge.

    @ti.kernel
    def update(self, s: ti.i32):
        # update the next step velocity and position
        for i in range(self.ngeom):
            self.geom_vel[s+1, i] = self.geom_vel[s, i] + \
                                     self.dt * self.geom_force[s, i] / self.geom_mass[i]
            self.geom_pos[s+1, i] = self.geom_pos[s, i] + \
                                     self.dt * self.geom_vel[s, i]     # <<<< change to use previous velocity

Thanks for helping @BisonLeo . It works. Closing this issue now.