milescranmer / lagrangian_nns

Lagrangian Neural Networks

Lagrangian Neural Networks

Miles Cranmer, Sam Greydanus, Stephan Hoyer, Peter Battaglia, David Spergel, Shirley Ho Accepted to ICLR 2020 Workshop on Deep Differential Equations

• Paper
• Blog
• Self-Contained Tutorial
• Paper example notebook: double pendulum
• Paper example notebook: special relativity
• Paper example notebook: wave equation

Warning To use our implementation with more recent versions of JAX, you change jax.experimental.stax to jax.example_libraries.stax and jax.experimental.optimizers to jax.example_libraries.optimizers. Please raise an issue if there are other deprecated functionalities.

Summary

In this project we propose Lagrangian Neural Networks (LNNs), which can parameterize arbitrary Lagrangians using neural networks. In contrast to Hamiltonian Neural Networks, these models do not require canonical coordinates and perform well in situations where generalized momentum is difficult to compute (e.g., the double pendulum). This is particularly appealing for use with a learned latent representation, a case where HNNs struggle. Unlike previous work on learning Lagrangians, LNNs are fully general and extend to non-holonomic systems such as the 1D wave equation.

	Neural Networks	Neural ODEs	HNN	DLN (ICLR'19)	LNN (this work)
Learns dynamics	✔️	✔️	✔️	✔️	✔️
Learns continuous-time dynamics		✔️	✔️	✔️	✔️
Learns exact conservation laws			✔️	✔️	✔️
Learns from arbitrary coordinates	✔️	✔️		✔️	✔️
Learns arbitrary Lagrangians					✔️

Dependencies

• Jax
• NumPy
• MoviePy (visualization)
• celluloid (visualization)

This project is written in Python 3.

lagrangian_nns's Issues

Get time in analytical_fn

Hi! I was trying to access time values from the analytical_fn function. I would need this to perform operations such as setting a desired value of the pendulum joint after a period of time, like when I am applying a control action. The idea I have in mind can be represented as follows

def analytical_fn(state, t=0, m1=1, m2=1, l1=1, l2=1, g=9.8):
      t1, t2, w1, w2 = state
      if(t > 1.5):
          out = t1
      else:
         out = t2
      ...

However, when I try access t, which is supposed to be the value for odeint timestamps, the traced<shapedarray(float32[])> data type does not allow me to capture its values, even when applying the approaches in the JAX documentation. So I can't perform the boolean comparison. How this could work?

Thank you in advance for your availability!

Issues when I open the main branch on google colab and use the `base-nn-double-pendulum.ipynb`

Issues when I open the main branch on google colab and use the base-nn-double-pendulum.ipynb:

1. cannot import from jax.example_libraries import stax. Fixed it by replacing he line with: from jax.example_libraries import stax
2. issues with the HyperParameterSearch.py:
   • ValueError: Non-hashable static arguments are not supported, as this can lead to unexpected cache-misses. Static argument (index 3) of type <class 'jax.interpreters.partial_eval.DynamicJaxprTracer'> for function _odeint_wrapper is non-hashable. for the function odeint.

The above exception was the direct cause of the following exception:

TypeError                                 Traceback (most recent call last)
[/content/lagrangian_nns/experiment_dblpend/data.py](https://localhost:8080/#) in get_trajectory_analytic(y0, times, **kwargs)
     38 @partial(jax.jit, backend='cpu')
     39 def get_trajectory_analytic(y0, times, **kwargs):
---> 40     return odeint(analytical_fn, y0, t=times, rtol=1e-10, atol=1e-10, **kwargs)
     41 
     42 def get_dataset(seed=0, samples=1, t_span=[0,2000], fps=1, test_split=0.5, **kwargs):

TypeError: odeint() got an unexpected keyword argument 'mxsteps'

4. When I run loss(get_params(opt_state), batch_data, 0.0) I get:

[/content/lagrangian_nns__modified/hyperopt/HyperparameterSearch.py](https://localhost:8080/#) in dynamics(q, q_t)
     32 #     assert q.shape == (2,)
     33     state = wrap_coords(jnp.concatenate([q, q_t]))
---> 34     return jnp.squeeze(nn_forward_fn(params, state), axis=-1)
     35   return dynamics
     36 

[/usr/local/lib/python3.7/dist-packages/jax/numpy/lax_numpy.py](https://localhost:8080/#) in squeeze(a, axis)
   1165     axis = frozenset(_canonicalize_axis(i, ndim(a)) for i in axis)
   1166     if _any(shape_a[a] != 1 for a in axis):
-> 1167       raise ValueError("cannot select an axis to squeeze out which has size "
   1168                        "not equal to one")
   1169     newshape = [d for i, d in enumerate(shape_a)

ValueError: cannot select an axis to squeeze out which has size not equal to one

I managed to fix this by changing:
THIS COMMAND: return jnp.squeeze(nn_forward_fn(params, state), axis=-1)
TO THIS COMMAND: return nn_forward_fn(params, state)

requiremented environment

Hello, I am reproducing your code failed due to jax version issues and configured environment. Can you release the environment and corresponding version of the pips in this paper，such as numpy==....,jax==.....

Issue with gln loss

Hello! I was trying to use the code in the "experiment_dblpend" directory. As long as it is training the neural network with the loss "baseline_nn", everything is fine.

The moment I try to use the loss "gln" (that should direct me to learning like the paper) something wrong happens. The resulting error is as follows (of course, without touching the code)

TypeError: Gradient only defined for scalar-output functions. Output had shape: (4,).

This error refers to the application of gln_loss at the following line

preds = jax.vmap(partial(lagrangian_eom, learned_dynamics(params)))(state)

Specifically to the command

(jax.grad(lagrangian, 0)(q, q_t)

of the function lagrangian_eom.

I honestly cannot explain this, because everything seems correct