def __init__(self, T=150):
# parameters
self.nDoF = 2
self.dt = 0.01 # s
""" arm model parameters """
self.x0 = np.array([np.deg2rad(5),
np.deg2rad(10), 0., 0.]) # initial state vector
self.l = np.array([0.3, 0.33]) # m. link length
self.m = np.array([1.4, 1.1]) # kg. link mass
self.I = np.array(
[0.025,
0.045]) # kg m^2. MOI about the CG of each of the link, Izz
self.s = np.array([0.11,
0.16]) # m. distance of CG from each of the joints
""" pre-compute some constants """
self.d1 = self.I[0] + self.I[1] + self.m[1] * self.l[0]**2
self.d2 = self.m[1] * self.l[0] * self.s[1]
self.dt_ = 0.01
self.ilqr = None
self.res = None
self.T = T
self.Q = 100 * np.eye(self.nDoF)
self.R = .01 * np.eye(self.nDoF)
# terminal Q to regularize final speed
self.Q_T = .1 * np.eye(self.nDoF)
self.target = np.array(
[0.4, 0.2]
) # cartesian position of target. At the target, the robot has to stop and thus zero velocity
self.target_thetas = self.inv_kin(
self.target
) # provides the joint angles associated with cart. target position
def reset(self):
"""Reset the state to start a new episode."""
self.rng, rng_input = random.split(self.rng)
angle = jnp.pi + random.uniform(
rng_input, minval=-jnp.deg2rad(10), maxval=jnp.deg2rad(10))
angular_velocity = 0.0
return jnp.array([angle, angular_velocity]) # state
def reward_func(self, state, action, next_state):
"""Return the reward for the given transition."""
distance_to_upright = abs(next_state[0] - jnp.pi)
if distance_to_upright < jnp.deg2rad(10):
reward = 1.0
else:
reward = -0.1
return reward
def test_inv_observe_range_bearing_when_robot_is_at_origin_and_landmark_is_on_right_side(
):
x_r = np.array([0, 0, 0])
p_local_polar = np.array([2., np.deg2rad(-90.)])
x, y = rbs.inv_observe_range_bearing(x_r, p_local_polar)
assert assert_almost_equal(x, 0.) is None
assert assert_almost_equal(y, -2.) is None
def test_inv_observe_range_bearing_when_robot_is_at_origin_and_landmark_is_directly_in_front(
):
x_r = np.array([0, 0, 0])
p_local_polar = np.array([4.5, np.deg2rad(0.)])
x, y = rbs.inv_observe_range_bearing(x_r, p_local_polar)
assert assert_almost_equal(x, 4.5) is None
assert assert_almost_equal(y, 0.) is None
def test_observe_range_bearing_when_robot_is_not_at_origin_and_landmark_is_on_right_side(
):
x_r = np.array([1., 0., 0.])
p_world_rect = np.array([1., -2.])
range_, bearing = rbs.observe_range_bearing(x_r, p_world_rect)
assert range_ == 2.
assert bearing == np.deg2rad(-90.)
def test_jacobian_H_L_i_at_specific_test_points_compared_to_auto_diff():
X = np.array([23.5, -14.6, np.deg2rad(45.)])
p_world_rect = np.array([39., 10.4]) # landmark position
f = rbs.observe_range_bearing
H_l = rbs.jacobian_H_L_i(x_r=X[0],
y_r=X[1],
alpha_r=X[2],
l_i_x=p_world_rect[0],
l_i_y=p_world_rect[1])
J = jax.jacfwd(f, argnums=1)(X, p_world_rect)
assert assert_array_almost_equal(J, H_l) is None
def test_jacobian_G_y_i_at_specific_test_points_compared_to_auto_diff():
X = np.array([23.5, -14.6, np.deg2rad(45.)])
p_local_polar = np.array([11.6, .25]) # range-bearing measurement
f = rbs.inv_observe_range_bearing
G_x = rbs.jacobian_G_y_i(x_r=X[0],
y_r=X[1],
alpha_r=X[2],
rho=p_local_polar[0],
psi=p_local_polar[1])
J = jax.jacfwd(f, argnums=1)(X, p_local_polar)
assert assert_array_almost_equal(J, G_x) is None
def deg2rad(x): if isinstance(x, JaxArray): x = x.value return JaxArray(jnp.deg2rad(x))
def correct_logsp_params(A, phi, q, psi, dpsi, theta):
Ap = jnp.exp(-dpsi * jnp.tan(jnp.deg2rad(phi)))
return A * Ap, phi, q, psi, theta + dpsi
def __lsp(A, phi, theta):
return (A * jnp.exp(theta * jnp.tan(jnp.deg2rad(phi))) * jnp.stack(
(jnp.cos(theta), jnp.sin(theta)))).T
def get_reparametrized_errors(agg_res):
if isinstance(agg_res.model, dict):
model = agg_res.params
else:
model = agg_res.model
disk = model['disk']
if 'bulge' in model:
bulge = model['bulge']
else:
bulge = EMPTY_SERSIC
if 'bar' in model:
bar = model['bar']
else:
bar = EMPTY_SERSIC
disk_e = agg_res.errors['disk']
if 'bulge' in agg_res.errors:
bulge_e = agg_res.errors['bulge']
else:
bulge_e = EMPTY_SERSIC_ERR
if 'bar' in agg_res.errors:
bar_e = agg_res.errors['bar']
else:
bar_e = EMPTY_SERSIC_ERR
errs = pd.DataFrame([],
columns=['disk', 'bulge', 'bar', 'spiral'],
dtype=jnp.float64)
errs['disk'] = disk_e.copy()
# it is possible that we have zero error for ellipticity, which will
# cause problems. Instead, fix it as a small value
errs.loc['q', 'disk'] = max(0.001, errs.loc['q', 'disk'])
errs.loc['L', 'disk'] = jnp.inf
errs.loc['I', 'disk'] = jnp.nan
errs.loc['n', 'disk'] = jnp.nan
errs.loc['c', 'disk'] = jnp.nan
errs['bulge'] = bulge_e.copy()
errs.loc['q', 'bulge'] = max(0.001, errs.loc['q', 'bulge'])
errs.loc['scale', 'bulge'] = bulge.Re / disk.Re * jnp.sqrt(
bulge_e.Re**2 / bulge.Re**2 + disk_e.Re**2 / disk.Re**2)
errs.loc['frac', 'bulge'] = jnp.inf
errs.loc['I', 'bulge'] = jnp.nan
errs.loc['Re', 'bulge'] = jnp.nan
errs.loc['c', 'bulge'] = jnp.nan
errs['bar'] = bar_e.copy()
errs.loc['q', 'bar'] = max(0.001, errs.loc['q', 'bar'])
errs.loc['scale',
'bar'] = bar.Re / disk.Re * jnp.sqrt(bar_e.Re**2 / bar.Re**2 +
disk_e.Re**2 / disk.Re**2)
errs.loc['frac', 'bar'] = jnp.inf
errs.loc['I', 'bar'] = jnp.nan
errs.loc['Re', 'bar'] = jnp.nan
errs.loc['mux', 'centre'] = jnp.sqrt(
jnp.nansum(jnp.array((disk_e.mux**2, bulge_e.mux**2))))
errs.loc['muy', 'centre'] = jnp.sqrt(
jnp.nansum(jnp.array((disk_e.muy**2, bulge_e.muy**2))))
for i in range(len(agg_res.spiral_arms)):
errs.loc['I.{}'.format(i), 'spiral'] = jnp.inf
errs.loc['falloff.{}'.format(i), 'spiral'] = jnp.inf
errs.loc['spread.{}'.format(i), 'spiral'] = jnp.inf
errs.loc['A.{}'.format(i), 'spiral'] = 0.01
errs.loc['phi.{}'.format(i), 'spiral'] = 1
errs.loc['t_min.{}'.format(i), 'spiral'] = jnp.deg2rad(0.5)
errs.loc['t_max.{}'.format(i), 'spiral'] = jnp.deg2rad(0.5)
return df_to_dict(errs)
