ik- newton

newton.py

@jit
def forward(state):
  return forward_ik(state["angles"], state["lengths"])

@jit
def checker(state):
  end_location = forward(state)
  distance = jnp.linalg.norm(end_location - state["target"])
  return distance > 0.00001

@jit
def step(state):
  d = forward(state) - state["target"]
  J_inv = jnp.linalg.pinv(J(state["angles"], state["lengths"]))
  state["angles"] -= jnp.dot(J_inv, d)
  state["history"]["iterations"] += 1
  return state

def inverse_kinematics(state):
  return while_loop(checker, step, state)