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)