Neural Physics Attempt

robz
robz
Nov 4 · 5 min read

Inspired by this paper that was recently discussed by Two Minute Papers, I attempted to train a neural network to simulate the physics in a simple 2D example: a ball bouncing around inside a box, subject to gravity and some energy decay.

I thought this would be a particularly interesting thing to simulate because of the box boundary. I think it would be relatively easy for a function approximator to learn trajectories that are subject to a constant acceleration, but fairly difficult to handle collisions.

I first wrote a simple simulator to generate episodes like this one:

Each episode, the ball is initialized with a random position and velocity. It then bounces around inside a box until the ball’s energy drops to a certain threshold. The state transitions at each time step (x, y, dx, dy) -> (x’, y’, dx’, dy’) were recorded from each episode.

I created a fully connected neural network with 4 input and 4 output neurons (one for each state variable), and 2 internal layers of 64 neurons each. I trained it on all the state transitions from 100 episodes, where each start and end state from the episodes served as training example for the NN. This yielded about 17,000 example state transitions. Training ran for 200 epochs.

I did all of this in a Google Collaboratory Notebook.

After training was complete, I used the NN to simulate some example trajectories of the ball. In each example, the ball was randomly initialized just like in the training episodes. The blue trajectory shows the path of the ball as predicted by my own simulation. The orange trajectory was created from first passing the start state into the NN to get the next state, and then passing that state back into the model to generate the next state, and so on.

The results are… terrible:

Sometimes the NN respects the bounds of the box (shown on the left), but 9 out of 10 times the NN made the ball fly off into infinity (as shown by the graph on the right). In any case, the NN’s trajectory never even comes close to the one my simulator calculates.

Reducing to a 1 dimensional world

I decided to simplify this (already very simple) example by only simulating the y axis (removing the x, or side-to-side component of motion). So instead of taking in and outputting a four-variable state (x, y, dx, dy), the NN now will only deal with a two-variable state (y, dy). This reduces the problem to just a 2 input-2 output function approximator.

The same simulator I wrote for the 2D ball was used to generate test data for this 1D system by always having the initial x and dx values set to 0. The NN architecture had 2x64x64x2 layers of neurons, fully connected, and was trained over 100 episodes of state transitions. Here’s a notebook showing the code.

Instead of plotting the x y position of the ball, trajectories below are renders with the x axis being the simulated time step (or time) and y axis is the height of the ball. As before, the blue line indicates the simulation I wrote, and the orange line represents the NN’s prediction.

The results for this system are somewhat more promising:

As with the 2D case, it seems like the NN can do a decent job of at least bounding the ball to the floor. This is probably due to how many examples it had where the ball collided with the floor; practically every episode ends with it resting on the floor.

Unlike the 2D case, I found no test runs in the 1D case where the ball flew off to infinity. Looking back at the times where that happened in the 2D case, I notice that all of the cases where the NN simulated the ball flying to infinity also included the ball having a sidewise velocity. Maybe because there weren’t enough examples where the ball hit the side walls, the NN wan’t able to properly model them, which allowed the ball to fly out of bounds.

I decided to push this 1D example further by training the same architecture using 1000 episodes instead of just 100. The results are quite a bit better. Here’s a montage of examples:

I’ve heard that more layers is better than a wider layers, so I attempted to train the system on a 2x8x8x8x8x2 network instead. Even with just 100 episodes, the results seem better than then 2x64x64x2 network above, especially at the beginning of the episode:

In summary, it seems like having more training data and a deeper network work well for the 2-input 2-output case, so those principles might generalize for the 2D system too. It’s also possible that having more examples where the ball smacks against the walls more could help develop a better model that includes the walls and prevents the balls from flying away.

Returning to 2D

With renewed motivation, I went back to working the 2D case, where the NN must simulate both the x and y components of the position and velocity of the ball.

I first applied the idea of having a deeper network by replacing the 4x64x64x4 network with a 4x8x8x8x8x4 one, still using just 100 episodes. The results were still quite bad, but at least I couldn’t find any cases where the ball flew to infinity:

This is as far as I’ll go for now. I’ll update this post if I experiment with different architectures or use a larger training set.

Welcome to a place where words matter. On Medium, smart voices and original ideas take center stage - with no ads in sight. Watch
Follow all the topics you care about, and we’ll deliver the best stories for you to your homepage and inbox. Explore
Get unlimited access to the best stories on Medium — and support writers while you’re at it. Just $5/month. Upgrade