Getting NaN for gradients when computing the loss with jit compile #400

chan-mander · 2023-10-02T20:13:28Z

chan-mander
Oct 2, 2023

Hello,

I am fairly new to JAX and Brax, and I am running into a issue that I don't quite understand. I am trying to get the sensitivity of the loss with respect to the model params by running the following minimum reproducible example.

import jax
import jax.numpy as jnp
from jax.config import config
import brax
from brax.io import mjcf
from brax.generalized import pipeline
import equinox as eqx
import optax
from functools import partial
from jax import config


class SimpleModel(eqx.Module):
    layers: list          # list of the layers in the NN
    extra_bias: jax.Array # extra bias on the output

    def __init__(self, key):
        key1, key2, key3 = jax.random.split(key, 3)
        # layers contain trainable parameters
        self.layers = [
            eqx.nn.Linear(2, 4, key=key1),
            eqx.nn.Linear(4, 4, key=key2),
            eqx.nn.Linear(4, 1, key=key3)
        ]
        # extra_bias is also a trainable parameter
        self.extra_bias = jax.numpy.ones(1)

    @jax.jit
    def __call__(self, x):
        for layer in self.layers[:-1]:
            x = jax.nn.elu(layer(x))
        return self.layers[-1](x) + self.extra_bias
    

@partial(jax.jit, static_argnames=('timesteps'))
def loss(model: SimpleModel, system: brax.System, state: brax.State, timesteps):

    def step(i: int, carry: tuple):

        system_s, state_s, model_s, x_pos_s = carry

        x_pos_s = x_pos_s.at[i].add(state_s.x.pos[0][0])
        x = jnp.array([state_s.x.pos[0][0], state_s.qd[0]])
        force = model_s(x.transpose())
        state_s = pipeline.step(system_s, state_s, force)

        return (system_s, state_s, model_s, x_pos_s)

    x_pos = jnp.zeros((timesteps, 1), dtype=jnp.float32)

    carry = (system, state, model, x_pos)
    system, state, model, x_pos = jax.lax.fori_loop(0, timesteps, step, carry)

    targetposition = jnp.zeros((timesteps, 1))

    loss_value = jnp.sum(1.0*jnp.abs(x_pos - targetposition)) / timesteps

    return loss_value


def main():
    # config.update("jax_disable_jit", True) 
    # config.update("jax_debug_nans", True)

    key = jax.random.PRNGKey(seed=0)
    key, subkey = jax.random.split(key, 2)
    model = SimpleModel(subkey)

    lr = 0.001
    opt = optax.adam(lr)
    opt_state = opt.init(model)

    simTime = 0.01
    timeSteps = int(simTime / 0.001)

    system: brax.System = mjcf.loads(
                                    """<mujoco>
                                            <option timestep="0.001"/>
                                            <worldbody>
                                                <body pos="{x} {y} {z}">
                                                    <joint name="slider" type="slide" axis="1 0 0" limited="false"/>
                                                    <geom size="0.5 0.5 0.5" type="box" mass="1"/>
                                                </body>
                                                <geom name="floor" size="40 40 40" type="plane"/>
                                            </worldbody>
                                            <actuator>
                                                <motor ctrllimited="true" ctrlrange="-10 10" gear="1"  joint="slider" name="slide"/>
                                            </actuator>
                                        </mujoco>""".format(x=1.0, y=0.0, z=0.5))

    qd = jnp.array([0.0])
    state = jax.jit(pipeline.init)(system, system.init_q, qd)

    for i in range(2):

        l, grad = eqx.filter_value_and_grad(loss, has_aux=False)(model, system, state, timeSteps)

        print("Batch: ", i+1, " Loss: ", l)
        print("Gradients of layer 0 weights:\n", grad.layers[0].weight)

        updates, opt_state = opt.update(grad, opt_state)
        model = optax.apply_updates(model, updates)


if __name__ == "__main__":
    main()

When I run this with Jax jit compile I get NaN for the gradients, but when I implement config.update("jax_disable_jit", True) to disable the jit compile I am getting values for gradients, but it is very slow. Can anyone help me shed some light as to why this is happening, and how I can compute gradients with the jit compiler? I am currently raising this issue with the authors of JAX as well. I also brought this up with the authors of equinox (https://github.com/patrick-kidger/equinox/issues/523), and they instructed me to raise the issue here.

chan-mander · 2023-10-04T03:32:30Z

chan-mander
Oct 4, 2023
Author

I was able to solve the issue. It seems that the problem was in how I was initializing the initial state with Brax.

1 reply

btaba Oct 15, 2023
Maintainer

Glad you were able to resolve the issue, apologies for the late reply!

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Getting NaN for gradients when computing the loss with jit compile #400

{{title}}

Replies: 1 comment 1 reply

{{title}}

{{title}}

Select a reply

Getting NaN for gradients when computing the loss with jit compile #400

chan-mander Oct 2, 2023

Replies: 1 comment · 1 reply

chan-mander Oct 4, 2023 Author

btaba Oct 15, 2023 Maintainer

chan-mander
Oct 2, 2023

Replies: 1 comment 1 reply

chan-mander
Oct 4, 2023
Author

btaba Oct 15, 2023
Maintainer