Learn to Throw¶
In this notebook, we will train a fully-connected neural network to solve an inverse ballistics problem. We will compare supervised training to differentiable physics training, both numerically and visually.
In this notebook, we will train a fully-connected neural network to solve an inverse ballistics problem. We will compare supervised training to differentiable physics training, both numerically and visually.
# !pip install phiflow
# from phi.tf.flow import *
from phi.torch.flow import *
# from phi.jax.stax.flow import *
To define the physics problem, we write a function to simulate the forward physics. This function takes the initial position, height, speed and angle of a thrown object and computes where the object will land, assuming the object follows a parabolic trajectory. Neglecting friction, can compute this by solving a quadratic equation.
def simulate_hit(pos, height, vel, angle, gravity=1.):
vel_x, vel_y = math.cos(angle) * vel, math.sin(angle) * vel
height = math.maximum(height, .5)
hit_time = (vel_y + math.sqrt(vel_y**2 + 2 * gravity * height)) / gravity
return pos + vel_x * hit_time, hit_time, height, vel_x, vel_y
simulate_hit(10, 1, 1, 0)[0]
float64 11.41421353816986
Let's plot the trajectory! We define y(x) and sample it on a grid. We have to drop invalid values, such as negative flight times and below-ground points.
def sample_trajectory(pos, height, vel, angle, gravity=1.):
hit, hit_time, height, vel_x, vel_y = simulate_hit(pos, height, vel, angle, gravity)
t = math.linspace(0, hit_time, spatial(t=100))
return vec(x=pos + vel_x * t, y=height + vel_y * t - gravity / 2 * t ** 2)
vis.plot(sample_trajectory(tensor(10), 1, 1, math.linspace(-PI/4, 1.5, channel(angle=7))), title="Varying Angle")
Before we train neural networks on this problem, let's perform a classical optimization using gradient descent in the initial velocity vel. We need to define a loss function to optimize. Here we desire the object to hit at x=0.
vel = tensor(1)
def loss_function(vel):
return math.l2_loss(simulate_hit(10, 1, vel, 0)[0] - 0)
loss_function(0)
float64 50.0
ΦFlow uses the selected library (PyTorch/TensorFlow/Jax) to derive analytic derivatives. By default, the gradient function also returns the function value.
grad_fn = gradient(loss_function, get_output=True)
grad_fn(vel)
(65.14213, [16.142136])
trj = [vel]
for i in range(10):
loss, (grad,) = grad_fn(vel)
vel = vel - .2 * grad
trj.append(vel)
print(f"vel={vel} - loss={loss}")
trj = math.stack(trj, channel('opt'))
vis.plot(sample_trajectory(tensor(10), 1, trj, 0))
vel=-2.2284272 - loss=65.14213 vel=-4.1654835 - loss=23.451168 vel=-5.3277173 - loss=8.442421 vel=-6.0250573 - loss=3.0392709 vel=-6.4434614 - loss=1.0941381 vel=-6.694504 - loss=0.39388976 vel=-6.8451295 - loss=0.14180061 vel=-6.935505 - loss=0.051048037 vel=-6.9897304 - loss=0.018377367 vel=-7.0222654 - loss=0.0066157645
Now we can just subtract the gradient times a learning rate $\eta = 0.2$ until we converge.
Next, we generate a training set and test by sampling random values.
def generate_data(shape):
pos = math.random_normal(shape)
height = math.random_uniform(shape) + .5
vel = math.random_uniform(shape)
angle = math.random_uniform(shape) * PI/2
return math.stack(dict(pos=pos, height=height, vel=vel, angle=angle), channel('vector'))
x_train = generate_data(batch(examples=1000))
x_test = generate_data(batch(examples=1000))
y_train = simulate_hit(*x_train.vector)[0]
y_test = simulate_hit(*x_test.vector)[0]
Now, let's create a fully-connected neural network with three hidden layers. We can reset the seed to make the weight initialization predictable.
math.seed(0)
net_sup = dense_net(1, 4, [32, 64, 32])
net_sup
--------------------------------------------------------------------------- TypeError Traceback (most recent call last) Cell In[9], line 1 ----> 1 math.seed(0) 2 net_sup = dense_net(1, 4, [32, 64, 32]) 3 net_sup File /opt/hostedtoolcache/Python/3.12.11/x64/lib/python3.12/site-packages/phiml/math/_ops.py:137, in seed(seed) 127 """ 128 Sets the current seed of all backends and the built-in `random` package. 129 (...) 134 seed: Seed to use. 135 """ 136 for backend in BACKENDS: --> 137 backend.seed(seed) 138 import random 139 random.seed(0) File numpy/random/mtrand.pyx:4820, in numpy.random.mtrand.seed() TypeError: seed() takes at most 1 positional argument (2 given)
# net_sup.trainable_weights[0] # TensorFlow
net_sup.state_dict()['linear0.weight'].flatten() # PyTorch
# net_sup.parameters[0][0] # Stax
--------------------------------------------------------------------------- NameError Traceback (most recent call last) Cell In[10], line 2 1 # net_sup.trainable_weights[0] # TensorFlow ----> 2 net_sup.state_dict()['linear0.weight'].flatten() # PyTorch 3 # net_sup.parameters[0][0] # Stax NameError: name 'net_sup' is not defined
For the differentiable physics network we do the same thing again.
math.seed(0)
net_dp = dense_net(1, 4, [32, 64, 32])
# net_dp.trainable_weights[0] # TensorFlow
net_dp.state_dict()['linear0.weight'].flatten() # PyTorch
# net_dp.parameters[0][0] # Stax
--------------------------------------------------------------------------- TypeError Traceback (most recent call last) Cell In[11], line 1 ----> 1 math.seed(0) 2 net_dp = dense_net(1, 4, [32, 64, 32]) 3 # net_dp.trainable_weights[0] # TensorFlow File /opt/hostedtoolcache/Python/3.12.11/x64/lib/python3.12/site-packages/phiml/math/_ops.py:137, in seed(seed) 127 """ 128 Sets the current seed of all backends and the built-in `random` package. 129 (...) 134 seed: Seed to use. 135 """ 136 for backend in BACKENDS: --> 137 backend.seed(seed) 138 import random 139 random.seed(0) File numpy/random/mtrand.pyx:4820, in numpy.random.mtrand.seed() TypeError: seed() takes at most 1 positional argument (2 given)
Indeed, the We see that the networks were initialized identically! Alternatively, we could have saved and loaded the state.
Now we can train the network. We feed the desired hit position into the network and predict a possible initial state. For supervised training, we compare the network prediction to the ground truth from our training set.
opt_sup = adam(net_sup)
def supervised_loss(x, y, net=net_sup):
prediction = math.native_call(net, y)
return math.l2_loss(prediction - x)
print(f"Supervised loss (training set): {supervised_loss(x_train, y_train)}")
print(f"Supervised loss (test set): {supervised_loss(x_test, y_test)}")
--------------------------------------------------------------------------- NameError Traceback (most recent call last) Cell In[12], line 1 ----> 1 opt_sup = adam(net_sup) 3 def supervised_loss(x, y, net=net_sup): 4 prediction = math.native_call(net, y) NameError: name 'net_sup' is not defined
for i in range(100):
update_weights(net_sup, opt_sup, supervised_loss, x_train, y_train)
print(f"Supervised loss (training set): {supervised_loss(x_train, y_train)}")
print(f"Supervised loss (test set): {supervised_loss(x_test, y_test)}")
--------------------------------------------------------------------------- NameError Traceback (most recent call last) Cell In[13], line 2 1 for i in range(100): ----> 2 update_weights(net_sup, opt_sup, supervised_loss, x_train, y_train) 4 print(f"Supervised loss (training set): {supervised_loss(x_train, y_train)}") 5 print(f"Supervised loss (test set): {supervised_loss(x_test, y_test)}") NameError: name 'net_sup' is not defined
What? That's almost no progress! Feel free to run more iterations but there is a deeper problem at work here. Before we get into that, let's train a the network again but with a differentiable physics loss function.
For the differentiable physics loss, we simulate the trajectory given the initial conditions predicted by the network. Then we can measure how close to the desired location the network got.
def physics_loss(y, net=net_dp):
prediction = math.native_call(net, y)
y_sim = simulate_hit(*prediction.vector)[0]
return math.l2_loss(y_sim - y)
opt_dp = adam(net_dp)
for i in range(100):
update_weights(net_dp, opt_dp, physics_loss, y_train)
print(f"Supervised loss (training set): {supervised_loss(x_train, y_train, net=net_dp)}")
print(f"Supervised loss (test set): {supervised_loss(x_test, y_test, net=net_dp)}")
--------------------------------------------------------------------------- NameError Traceback (most recent call last) Cell In[14], line 1 ----> 1 def physics_loss(y, net=net_dp): 2 prediction = math.native_call(net, y) 3 y_sim = simulate_hit(*prediction.vector)[0] NameError: name 'net_dp' is not defined
This looks even worse! The differentiable physics network seems to stray even further from the ground truth. Well, we're not trying to match the ground truth, though. Let's instead measure how close to the desired location the network threw the object.
print(f"Supervised network (test set): {physics_loss(y_test, net=net_sup)}")
print(f"Diff.Phys. network (test set): {physics_loss(y_test, net=net_dp)}")
--------------------------------------------------------------------------- NameError Traceback (most recent call last) Cell In[15], line 1 ----> 1 print(f"Supervised network (test set): {physics_loss(y_test, net=net_sup)}") 2 print(f"Diff.Phys. network (test set): {physics_loss(y_test, net=net_dp)}") NameError: name 'physics_loss' is not defined
Now this is much more promissing. The diff.phys. network seems to hit the desired location very accurately considering it was only trained for 100 iterations. With more training steps, this loss will go down even further, unlike the supervised network.
So what is going on here? Why does the supervised network perform so poorly? The answer lies in the problem itself. The task is multi-modal, i.e. there are many initial states that will hit the same target. The network only gets the target position and must decide on a single initial state. With supervised training, there is no way to know which ground truth solution occurs in the test set. The best the network can do is average nearby solutions from the training set. But since the problem is non-linear, this will give only a rough guess.
The diff.phys. network completely ignores the ground truth solutions which are not even passed to the physics_loss function. Instead, it learns to hit the desired spot, which is exactly what we want.
We can visualize the difference by looking at a couple of trajectories.
predictions = math.stack({
"Ground Truth": x_test.examples[:4],
"Supervised": math.native_call(net_sup, y_test.examples[:4]),
"Diff.Phys": math.native_call(net_dp, y_test.examples[:4]),
}, channel('curves'))
vis.plot(sample_trajectory(*predictions.vector), size=(16, 4))
--------------------------------------------------------------------------- NameError Traceback (most recent call last) Cell In[16], line 3 1 predictions = math.stack({ 2 "Ground Truth": x_test.examples[:4], ----> 3 "Supervised": math.native_call(net_sup, y_test.examples[:4]), 4 "Diff.Phys": math.native_call(net_dp, y_test.examples[:4]), 5 }, channel('curves')) 6 vis.plot(sample_trajectory(*predictions.vector), size=(16, 4)) NameError: name 'net_sup' is not defined
We can see that the differentiable physics network matches the hit point much more precisely than the supervised network, as expected from the loss values.