def test_quad_cost_with_wrapping():
p = create_params()
n, k = p.n, p.k
x_dim, u_dim = 3, 2
a, b = p.a, p.b
goal_x, goal_y = 10., 10.
dubins = DubinsV1(p.dt, create_system_dynamics_params())
goal = np.array([goal_x, goal_y, 0.], dtype=np.float32)
x_ref_nk3 = np.array(np.tile(goal, (n, k, 1)))
u_ref_nk2 = np.array(np.zeros((n, k, u_dim), dtype=np.float32))
trajectory_ref = dubins.assemble_trajectory(x_ref_nk3, u_ref_nk2)
cost_fn = QuadraticRegulatorRef(trajectory_ref, dubins, p)
x_nk3 = np.array(np.zeros((n, k, x_dim), dtype=np.float32))
u_nk2 = np.array(np.zeros((n, k, u_dim), dtype=np.float32))
trajectory = dubins.assemble_trajectory(x_nk3, u_nk2)
cost_nk, _ = cost_fn.compute_trajectory_cost(trajectory)
cost = .5 * (a * goal_x**2 + a * goal_y**2)
assert (cost_nk.shape[0] == n and cost_nk.shape[1] == k)
assert ((cost_nk == cost).all())
H_xx_nkdd, H_xu_nkdf, H_uu_nkff, J_x_nkd, J_u_nkf = cost_fn.quad_coeffs(
trajectory)
assert (np.equal(H_xx_nkdd[0, 0], np.eye(x_dim)).all())
assert ((H_xu_nkdf == 0.).all())
assert (np.equal(H_uu_nkff[0, 0], np.eye(u_dim) * b).all())
def test_custom_dubins_v1(plot=False):
# Visualize One Trajectory for Debugging
dt = .1
n, k = 5, 50
x_dim, u_dim = 3, 2
# Generate Dubin's Model
db = DubinsV1(dt, create_system_dynamics_params())
# Generate Start State (x0, y0, t0)
state_113 = np.array([[[0, 0, 0]]], dtype=np.float32)
# Generate linear velocity controls
v0 = 0.8
t0 = 0.5
v_1k = np.ones((k, 1)) * v0
# v_1k = rand_array(from_range=(0, 10), size=(k,1))
# Generate angular velocity controls
# Randomly generate directions to control
# w_1k = rand_array(from_range=(-10, 10), size=(k,1))
w_1k = np.linspace(1, 10, k)[:, None]
# w_1k = np.ones((k, 1)) * t0 # np.linspace(0.9, 1.1, k)[:, None]
ctrl_1k2 = np.array(np.concatenate([v_1k, w_1k], axis=1)[None],
dtype=np.float32)
trajectory = db.simulate_T(state_113, ctrl_1k2, T=k)
state_1k3, _ = db.parse_trajectory(trajectory)
if plot:
fig = plt.figure()
ax = fig.add_subplot(111)
xs, ys, ts = state_1k3[0, :, 0], state_1k3[0, :, 1], state_1k3[0, :, 2]
ax.plot(xs, ys, 'r--')
ax.quiver(xs, ys, np.cos(ts), np.sin(ts))
# plt.show()
fig.savefig("./tests/dynamics/custom_dubins.png",
bbox_inches='tight',
pad_inches=0)
def visualize_coordinate_transform():
"""Visual sanity check that coordinate transforms
are working. """
fig, _, axs = utils.subplot2(plt, (2, 2), (8, 8), (.4, .4))
axs = axs[::-1]
n, k = 1, 30
dt = .1
dubins_car = DubinsV1(dt=dt)
traj_egocentric = Trajectory(dt=dt, n=n, k=k, variable=True)
traj_world = Trajectory(dt=dt, n=n, k=k, variable=True)
# Form a trajectory in global frame
# convert to egocentric and back
start_pos_global_n12 = tf.constant([[[1.0, 1.0]]], dtype=tf.float32)
start_heading_global_n11 = tf.constant([[[np.pi / 2.]]], dtype=tf.float32)
start_config_global = SystemConfig(dt=dt, n=n, k=1, position_nk2=start_pos_global_n12,
heading_nk1=start_heading_global_n11)
start_n13 = tf.concat([start_pos_global_n12, start_heading_global_n11], axis=2)
u_n12 = np.array([[[.01, .1]]], dtype=np.float32)
u_nk2 = tf.constant(np.broadcast_to(u_n12, (n, k, 2)), dtype=tf.float32)
trajectory_world = dubins_car.simulate_T(start_n13, u_nk2, T=k - 1)
trajectory_world.render([axs[0]], batch_idx=0, freq=4, name='World')
# Convert to egocentric
dubins_car.to_egocentric_coordinates(start_config_global, trajectory_world, traj_egocentric)
traj_egocentric.render([axs[1]], batch_idx=0, freq=4, name='Egocentric')
dubins_car.to_world_coordinates(start_config_global, traj_egocentric, traj_world)
traj_world.render([axs[2]], batch_idx=0, freq=4, name='World #2')
plt.savefig('./tmp/coordinate_transform.png', bbox_inches='tight')
def test_lqr1(visualize=False):
p = create_params()
np.random.seed(seed=p.seed)
n, k = p.n, 50
dt = p.dt
db = DubinsV1(dt, params=p.system_dynamics_params.simulation_params)
x_dim, u_dim = db._x_dim, db._u_dim
x_n13 = tf.constant(np.zeros((n, 1, x_dim)), dtype=tf.float32)
v_1k = np.ones((k - 1, 1)) * .1
w_1k = np.linspace(.5, .3, k - 1)[:, None]
u_1k2 = tf.constant(np.concatenate([v_1k, w_1k], axis=1)[None],
dtype=tf.float32)
u_nk2 = tf.zeros((n, k - 1, 2), dtype=tf.float32) + u_1k2
trajectory_ref = db.simulate_T(x_n13, u_nk2, T=k)
cost_fn = QuadraticRegulatorRef(trajectory_ref, db, p)
x_nk3 = tf.constant(np.zeros((n, k, x_dim), dtype=np.float32))
u_nk2 = tf.constant(np.zeros((n, k, u_dim), dtype=np.float32))
trajectory = db.assemble_trajectory(x_nk3, u_nk2)
lqr_solver = LQRSolver(T=k - 1, dynamics=db, cost=cost_fn)
start_config = db.init_egocentric_robot_config(dt=dt, n=n)
lqr_res = lqr_solver.lqr(start_config, trajectory, verbose=False)
trajectory_opt = lqr_res['trajectory_opt']
assert (np.abs(lqr_res['J_hist'][1][0] - .022867) < 1e-4)
assert (np.abs(lqr_res['J_hist'][0][0] - 38.17334) < 1e-4)
if visualize:
pos_ref = trajectory_ref.position_nk2()[0]
pos_opt = trajectory_opt.position_nk2()[0]
heading_ref = trajectory_ref.heading_nk1()[0]
heading_opt = trajectory_opt.heading_nk1()[0]
fig = plt.figure()
ax = fig.add_subplot(121)
ax.plot(pos_ref[:, 0], pos_ref[:, 1], 'r-', label='ref')
ax.quiver(pos_ref[:, 0], pos_ref[:, 1], tf.cos(heading_ref),
tf.sin(heading_ref))
ax.plot(pos_opt[:, 0], pos_opt[:, 1], 'b-', label='opt')
ax.quiver(pos_opt[:, 0], pos_opt[:, 1], tf.cos(heading_opt),
tf.sin(heading_opt))
ax.legend()
# plt.show()
fig.savefig('./tests/lqr/test_lqr1.png',
bbox_inches='tight',
pad_inches=0)
else:
print('rerun test_lqr1 with visualize=True to visualize the test')
def test_lqr0(visualize=False):
p = create_params()
np.random.seed(seed=p.seed)
n, k = 1, p.k
dt = p.dt
db = DubinsV1(dt, params=p.system_dynamics_params.simulation_params)
x_dim, u_dim = db._x_dim, db._u_dim
goal_x, goal_y = 4.0, 0.0
goal = np.array([goal_x, goal_y, 0.], dtype=np.float32)
x_ref_nk3 = tf.constant(np.tile(goal, (n, k, 1)))
u_ref_nk2 = tf.constant(np.zeros((n, k, u_dim), dtype=np.float32))
trajectory_ref = db.assemble_trajectory(x_ref_nk3, u_ref_nk2)
cost_fn = QuadraticRegulatorRef(trajectory_ref, db, p)
x_nk3 = tf.constant(np.zeros((n, k, x_dim), dtype=np.float32))
u_nk2 = tf.constant(np.zeros((n, k, u_dim), dtype=np.float32))
# Initiate a blank trajectory (all 0's)
trajectory = db.assemble_trajectory(x_nk3, u_nk2)
lqr_solver = LQRSolver(T=k - 1, dynamics=db, cost=cost_fn)
cost = lqr_solver.evaluate_trajectory_cost(trajectory)
expected_cost = .5 * goal_x**2 * k + .5 * goal_y**2 * k
assert ((cost.numpy() == expected_cost).all())
start_config = db.init_egocentric_robot_config(dt=dt, n=n)
lqr_res = lqr_solver.lqr(start_config, trajectory, verbose=False)
trajectory_opt = lqr_res['trajectory_opt']
J_opt = lqr_res['J_hist'][-1]
assert ((J_opt.numpy() == 8.).all())
assert (np.allclose(trajectory_opt.position_nk2()[:, 1:, 0], 4.0))
if visualize:
pos_ref = trajectory_ref.position_nk2()[0]
pos_opt = trajectory_opt.position_nk2()[0]
fig = plt.figure()
ax = fig.add_subplot(111)
ax.scatter(pos_ref[:, 0], pos_ref[:, 1])
ax.plot(pos_opt[:, 0], pos_opt[:, 1], 'b--', label='opt')
ax.legend()
# plt.show()
fig.savefig('./tests/lqr/test_lqr0.png',
bbox_inches='tight',
pad_inches=0)
else:
print('rerun test_lqr0 with visualize=True to visualize the test')
def test_coordinate_transform():
n, k = 1, 30
dt = .1
p = create_params()
dubins_car = DubinsV1(dt=dt, params=p.system_dynamics_params)
ref_config = dubins_car.init_egocentric_robot_config(dt=dt, n=n)
pos_nk2 = np.ones((n, k, 2), dtype=np.float32) * np.random.rand()
traj_global = Trajectory(dt=dt, n=n, k=k, position_nk2=pos_nk2)
traj_egocentric = Trajectory(dt=dt, n=n, k=k, variable=True)
traj_global_new = Trajectory(dt=dt, n=n, k=k, variable=True)
dubins_car.to_egocentric_coordinates(ref_config, traj_global,
traj_egocentric)
# Test 0 transform
assert ((pos_nk2 == traj_egocentric.position_nk2()).all())
ref_config_pos_112 = np.array([[[5.0, 5.0]]], dtype=np.float32)
ref_config_pos_n12 = np.repeat(ref_config_pos_112, repeats=n, axis=0)
ref_config = SystemConfig(dt=dt, n=n, k=1, position_nk2=ref_config_pos_n12)
traj_egocentric = dubins_car.to_egocentric_coordinates(
ref_config, traj_global, traj_egocentric)
# Test translation
assert ((pos_nk2 - 5.0 == traj_egocentric.position_nk2()).all())
ref_config_heading_111 = np.array([[[3. * np.pi / 4.]]], dtype=np.float32)
ref_config_heading_nk1 = np.repeat(ref_config_heading_111,
repeats=n,
axis=0)
ref_config = SystemConfig(dt=dt,
n=n,
k=1,
position_nk2=ref_config_pos_n12,
heading_nk1=ref_config_heading_nk1)
traj_egocentric = dubins_car.to_egocentric_coordinates(
ref_config, traj_global, traj_egocentric)
traj_global_new = dubins_car.to_world_coordinates(ref_config,
traj_egocentric,
traj_global_new)
assert (np.allclose(traj_global.position_nk2(),
traj_global_new.position_nk2()))
def test_lqr2(visualize=False):
p = create_params()
np.random.seed(seed=p.seed)
n, k = 2, 50
dt = p.dt
db = DubinsV1(dt, params=p.system_dynamics_params.simulation_params)
x_dim, u_dim = db._x_dim, db._u_dim
x_n13 = tf.constant(np.zeros((n, 1, x_dim)), dtype=tf.float32)
v_1k, w_1k = np.ones((k - 1, 1)) * .1, np.linspace(.5, .3, k - 1)[:, None]
u_1k2 = tf.constant(np.concatenate([v_1k, w_1k], axis=1)[None],
dtype=tf.float32)
u_nk2 = tf.zeros((n, k - 1, 2), dtype=tf.float32) + u_1k2
trajectory_ref = db.simulate_T(x_n13, u_nk2, T=k)
x_nk3, u_nk2 = db.parse_trajectory(trajectory_ref)
# stack two different reference trajectories together
# to verify that batched LQR works across the batch dim
goal_x, goal_y = 4.0, 0.0
goal = np.array([goal_x, goal_y, 0.], dtype=np.float32)
x_ref_nk3 = tf.constant(np.tile(goal, (1, k, 1)))
u_ref_nk2 = tf.constant(np.zeros((1, k, u_dim), dtype=np.float32))
x_nk3 = tf.concat([x_ref_nk3, x_nk3[0:1]], axis=0)
u_nk2 = tf.concat([u_ref_nk2, u_nk2[0:1]], axis=0)
trajectory_ref = db.assemble_trajectory(x_nk3, u_nk2)
cost_fn = QuadraticRegulatorRef(trajectory_ref, db, p)
x_nk3 = tf.constant(np.zeros((n, k, x_dim), dtype=np.float32))
u_nk2 = tf.constant(np.zeros((n, k, u_dim), dtype=np.float32))
trajectory = db.assemble_trajectory(x_nk3, u_nk2)
lqr_solver = LQRSolver(T=k - 1, dynamics=db, cost=cost_fn)
start_config = db.init_egocentric_robot_config(dt=dt, n=n)
lqr_res = lqr_solver.lqr(start_config, trajectory, verbose=False)
trajectory_opt = lqr_res['trajectory_opt']
assert (np.abs(lqr_res['J_hist'][1][0] - 8.0) < 1e-4)
assert (np.abs(lqr_res['J_hist'][0][0] - 400.0) < 1e-4)
if visualize:
pos_ref = trajectory_ref.position_nk2()[0]
pos_opt = trajectory_opt.position_nk2()[0]
heading_ref = trajectory_ref.heading_nk1()[0]
heading_opt = trajectory_opt.heading_nk1()[0]
fig = plt.figure()
ax = fig.add_subplot(121)
ax.plot(pos_ref[:, 0], pos_ref[:, 1], 'r-', label='ref')
ax.quiver(pos_ref[:, 0], pos_ref[:, 1], tf.cos(heading_ref),
tf.sin(heading_ref))
ax.plot(pos_opt[:, 0], pos_opt[:, 1], 'b-', label='opt')
ax.quiver(pos_opt[:, 0], pos_opt[:, 1], tf.cos(heading_opt),
tf.sin(heading_opt))
ax.set_title('Goal [4.0, 0.0]')
ax.legend()
pos_ref = trajectory_ref.position_nk2()[1]
pos_opt = trajectory_opt.position_nk2()[1]
heading_ref = trajectory_ref.heading_nk1()[1]
heading_opt = trajectory_opt.heading_nk1()[1]
ax = fig.add_subplot(122)
ax.plot(pos_ref[:, 0], pos_ref[:, 1], 'r-', label='ref')
ax.quiver(pos_ref[:, 0], pos_ref[:, 1], tf.cos(heading_ref),
tf.sin(heading_ref))
ax.plot(pos_opt[:, 0], pos_opt[:, 1], 'b-', label='opt')
ax.quiver(pos_opt[:, 0], pos_opt[:, 1], tf.cos(heading_opt),
tf.sin(heading_opt))
ax.set_title('Nonlinear Traj')
ax.legend()
# plt.show()
fig.savefig('./tests/lqr/test_lqr2.png',
bbox_inches='tight',
pad_inches=0)
else:
print('rerun test_lqr2 with visualize=True to visualize the test')
def test_dubins_v1(visualize=False):
np.random.seed(seed=1)
dt = .1
n, k = 5, 20
x_dim, u_dim = 3, 2
# Test that All Dimensions Work
db = DubinsV1(dt, create_system_dynamics_params())
state_nk3 = np.array(np.zeros((n, k, x_dim), dtype=np.float32))
ctrl_nk2 = np.array(np.random.randn(n, k, u_dim), dtype=np.float32)
trajectory = db.assemble_trajectory(state_nk3, ctrl_nk2)
state_tp1_nk3 = db.simulate(state_nk3, ctrl_nk2)
assert (state_tp1_nk3.shape == (n, k, x_dim))
jac_x_nk33 = db.jac_x(trajectory)
assert (jac_x_nk33.shape == (n, k, x_dim, x_dim))
jac_u_nk32 = db.jac_u(trajectory)
assert (jac_u_nk32.shape == (n, k, x_dim, u_dim))
A, B, c = db.affine_factors(trajectory)
# Test that computation is occurring correctly
n, k = 2, 3
ctrl = 1
state_n13 = np.array(np.zeros((n, 1, x_dim)), dtype=np.float32)
ctrl_nk2 = np.array(np.ones((n, k, u_dim)) * ctrl, dtype=np.float32)
trajectory = db.simulate_T(state_n13, ctrl_nk2, T=k)
state_nk3 = trajectory.position_and_heading_nk3()
x1, x2, x3, x4 = (state_nk3[0, 0], state_nk3[0, 1], state_nk3[0, 2],
state_nk3[0, 3])
assert ((x1 == np.zeros(3)).all())
assert (np.allclose(x2, [.1, 0., .1]))
assert (np.allclose(x3, [.1 + .1 * np.cos(.1), .1 * np.sin(.1), .2]))
assert (np.allclose(x4, [.2975, .0298, .3], atol=1e-4))
trajectory = db.assemble_trajectory(state_nk3[:, :-1], ctrl_nk2)
A, B, c = db.affine_factors(trajectory)
A0, A1, A2 = A[0, 0], A[0, 1], A[0, 2]
A0_c = np.array([[1., 0., 0.], [0., 1., .1], [0., 0., 1.]])
A1_c = np.array([[1., 0., -.1 * np.sin(.1)], [0., 1., .1 * np.cos(.1)],
[0., 0., 1.]])
A2_c = np.array([[1., 0., -.1 * np.sin(.2)], [0., 1., .1 * np.cos(.2)],
[0., 0., 1.]])
assert (np.allclose(A0, A0_c))
assert (np.allclose(A1, A1_c))
assert (np.allclose(A2, A2_c))
B0, B1, B2 = B[0, 0], B[0, 1], B[0, 2]
B0_c = np.array([[.1, 0.], [0., 0.], [0., .1]])
B1_c = np.array([[.1 * np.cos(.1), 0.], [.1 * np.sin(.1), 0.], [0., .1]])
B2_c = np.array([[.1 * np.cos(.2), 0.], [.1 * np.sin(.2), 0.], [0., .1]])
assert (np.allclose(B0, B0_c))
assert (np.allclose(B1, B1_c))
assert (np.allclose(B2, B2_c))
if visualize:
# Visualize One Trajectory for Debugging
k = 50
state_113 = np.array(np.zeros((1, 1, x_dim)), dtype=np.float32)
v_1k, w_1k = np.ones((k, 1)) * .2, np.linspace(1.1, .9, k)[:, None]
ctrl_1k2 = np.array(np.concatenate([v_1k, w_1k], axis=1)[None],
dtype=np.float32)
trajectory = db.simulate_T(state_113, ctrl_1k2, T=k)
state_1k3, _ = db.parse_trajectory(trajectory)
fig = plt.figure()
ax = fig.add_subplot(111)
xs, ys, ts = state_1k3[0, :, 0], state_1k3[0, :, 1], state_1k3[0, :, 2]
ax.plot(xs, ys, 'r--')
ax.quiver(xs, ys, np.cos(ts), np.sin(ts))
# plt.show()
fig.savefig("./tests/dynamics/test_dynamics1.png",
bbox_inches='tight',
pad_inches=0)