Introduces a version number, requirements files, a setup definition, brief installation instructions, and a package-level init file. Also removes an unnecessary package dependency.

When trying to run a 1D quadratic control affine nonlinear system of the form as shown below, the ILQR implementation is unable to handle scalar valued systems and results in a dimensionality mismatch error. Please find error and code below.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import functools
import os

import jax
from jax import device_put
from jax import vmap
from jax.config import config
import jax.numpy as np
import numpy as onp

from trajax import optimizers
from trajax.integrators import euler
from trajax.integrators import rk4
import matplotlib.pyplot as plt

def quadratic_nonlinear(x, u, t, params=(5, 10)):
    """
    Simple quadratic nonlinear system where we introduce reference trajectory as input
    :param x: 1D scalar state
    :param u: 1D input
    :param params: Kp, Kd gains for PD control law
    :return xdot: 1D array of shape 1
    """
    del t
    Kp, Kd = params
    r = np.squeeze(u)
    # xdot = (x ** 2 + Kp * (x - r) - Kd * rdot)/(1 - Kd)
    xdot = x ** 2 + Kp * (x - r)
    return np.array([xdot])

class ILQR_test():
    """
    Testing ILQR implementation in trajax for simple nonlinear systems
    """
    def __init__(self):
        pass


    def discretize(self, type='euler', dynamics=None):
        if dynamics is not None:
            self.dynamics = dynamics

        self.dynamics = euler(self.dynamics, dt=0.01)
        if type != 'euler':
            self.dynamics = rk4(self.dynamics, dt=0.01)


    def testQuadNonLinear(self, maxiter):
        """
        Calling ilqr on quadratic nonlinear system with input as reference trajectory
        :param maxiter: maximum number of iterations to take in ilqr
        :return: list of ilqr fn output
        """
        horizon = 100
        dynamics = rk4(quadratic_nonlinear, dt=0.01)

        true_params = (100.0, 10.0, 1.0)

        def cost(params, state, action, t):
            final_weight, stage_weight, action_weight = params

            state_err = state - action
            state_cost = stage_weight * (state_err ** 2 + action ** 2)
            # action_cost = action_weight * np.squeeze(action) ** 2
            return np.where(t == horizon, final_weight * state_cost,
                            state_cost)

        x0 = np.array([-0.9])
        U0 = np.zeros((horizon, 1))
        X, U, obj, grad, adj, lqr_val, total_iter = optimizers.ilqr(
            functools.partial(cost, true_params), dynamics, x0, U0, maxiter)
        return [X, U, obj, grad, adj, lqr_val, total_iter]


test = ILQR_test()
traj_cost = []
num_iter = [2, 30, 40, 50]

for i in num_iter:
    print(i)
    # traj = test.apply_ilqr(x0=onp.random.randn(2), U=onp.random.randn(2), maxiter=i, dynamics=rk4(quadratic_nonlinear, dt=0.01))
    # traj = test.testPendulumReadmeExample(maxiter=i)
    traj = test.testQuadNonLinear(maxiter=i)
    traj_cost.append(traj[2])

X = traj[0]
U = traj[1]

print(traj_cost)

Has anyone tried the solvers on BRAX environments? Here's what I have:

import trajax
import jax
from jax import numpy as jnp
from jax.flatten_util import ravel_pytree
import brax
from brax import envs

def get_f_and_c(env):
    key = jax.random.PRNGKey(0)
    state = env.reset(key)
    _, x2qp = ravel_pytree(state.qp)
    def f(x, u, t):
        qp = x2qp(x)
        nqp, _ = env.sys.step(qp, u)
        return ravel_pytree(nqp)[0]
    def c(x, u, t):
        qp = x2qp(x)
        dstate = state.replace(qp=qp)
        nstate = env.step(dstate, u)
        return -nstate.reward
    return f, c

env = envs.create('inverted_pendulum')
key = jax.random.PRNGKey(0)
state = env.reset(key)
x_init, x2qp = ravel_pytree(state.qp)

f, c = get_f_and_c(env)

x, u, cost, *outputs = trajax.optimizers.ilqr(c, f, x_init, jnp.zeros([1, env.action_size]))

which gives:

TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Traced<ShapedArray(float32[3])>with<DynamicJaxprTrace(level=1/0)>
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError

Hi, I'd like to contribute to iCEM extensions of current CEM implementation. I'd like to ask:

• if it's not on your list?
• if is it ok, if I extend current cem() function with additional (optional) features and hyperparameters?

Regards,

Hi, I noticed there are some problems with running CEM.

1. Here the last index should be 7: https://github.com/google/trajax/blob/67fd5ed4867914f1ce1ab78819551c12206a7773/trajax/optimizers.py#L753 (similar in random shooting: https://github.com/google/trajax/blob/67fd5ed4867914f1ce1ab78819551c12206a7773/trajax/optimizers.py#L818)

2. Default hyperparameters should be frozendict https://github.com/google/trajax/blob/main/trajax/optimizers.py#L685-L692

Would you like a PR for this?

Regards,

Hi, do we have in total performance benefit from injecting jax args to scipy.optimize.minimize()? Are there any plans to extend jax.scipy.optimize.minimize() to constrained problems? Regards