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_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_lqr_feedback_coordinate_transform():
p = create_params()
n, k = p.n, p.k
dubins = p.system_dynamics_params.system(p.dt, params=p.simulation_params)
# # Robot starts from (0, 0, 0)
# # and does a small spiral
start_pos_n13 = tf.constant(np.array([[[0.0, 0.0, 0.0]]], dtype=np.float32))
speed_nk1 = np.ones((n, k - 1, 1), dtype=np.float32) * 2.0
angular_speed_nk1 = np.linspace(1.5, 1.3, k - 1, dtype=np.float32)[None, :, None]
u_nk2 = tf.constant(np.concatenate([speed_nk1, angular_speed_nk1], axis=2))
traj_ref_egocentric = dubins.simulate_T(start_pos_n13, u_nk2, k)
cost_fn = QuadraticRegulatorRef(traj_ref_egocentric, dubins, p)
lqr_solver = LQRSolver(T=k - 1, dynamics=dubins, cost=cost_fn)
start_config0 = SystemConfig(dt=p.dt, n=p.n, k=1,
position_nk2=start_pos_n13[:, :, :2],
heading_nk1=start_pos_n13[:, :, 2:3])
lqr_res = lqr_solver.lqr(start_config0, traj_ref_egocentric, verbose=False)
traj_lqr_egocentric = lqr_res['trajectory_opt']
K_array_egocentric = lqr_res['K_opt_nkfd']
k_array = lqr_res['k_opt_nkf1']
# The origin of the egocentric frame in world coordinates
start_pos_n13 = tf.constant(np.array([[[1.0, 1.0, np.pi / 2.]]], dtype=np.float32))
ref_config = SystemConfig(dt=p.dt, n=p.n, k=1,
position_nk2=start_pos_n13[:, :, :2],
heading_nk1=start_pos_n13[:, :, 2:3])
# Convert to world coordinates
traj_ref_world = dubins.to_world_coordinates(ref_config, traj_ref_egocentric, mode='new')
traj_lqr_world = dubins.to_world_coordinates(ref_config, traj_lqr_egocentric, mode='new')
K_array_world = dubins.convert_K_to_world_coordinates(ref_config, K_array_egocentric, mode='new')
# Apply K_array_world to the system from ref_config
traj_test_world = lqr_solver.apply_control(ref_config, traj_ref_world,
k_array, K_array_world)
# Check that the LQR Trajectory computed using K_array_egocentric
# then transformed to world (traj_lqr_world) matches the
# LQR Trajectory computed directly using K_array_world
assert((tf.norm(traj_lqr_world.position_nk2() - traj_test_world.position_nk2(), axis=2).numpy() < 1e-4).all())
assert((tf.norm(angle_normalize(traj_lqr_world.heading_nk1() - traj_test_world.heading_nk1()), axis=2).numpy() < 1e-4).all())
assert((tf.norm(traj_lqr_world.speed_nk1() - traj_test_world.speed_nk1(), axis=2).numpy() < 1e-4).all())
assert((tf.norm(traj_lqr_world.angular_speed_nk1() - traj_test_world.angular_speed_nk1(), axis=2).numpy() < 1e-4).all())
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')