| 41 | |
| 42 | |
| 43 | class Car2d: |
| 44 | def __init__(self): |
| 45 | self.dt = 0.1 |
| 46 | self.H = 50 |
| 47 | r_obs = 0.3 |
| 48 | self.obs_center = jnp.array( |
| 49 | [ |
| 50 | [-r_obs * 3, r_obs * 2], |
| 51 | [-r_obs * 2, r_obs * 2], |
| 52 | [-r_obs * 1, r_obs * 2], |
| 53 | [0.0, r_obs * 2], |
| 54 | [0.0, r_obs * 1], |
| 55 | [0.0, 0.0], |
| 56 | [0.0, -r_obs * 1], |
| 57 | [-r_obs * 3, -r_obs * 2], |
| 58 | [-r_obs * 2, -r_obs * 2], |
| 59 | [-r_obs * 1, -r_obs * 2], |
| 60 | [0.0, -r_obs * 2], |
| 61 | ] |
| 62 | ) |
| 63 | self.obs_radius = r_obs # Radius of the obstacle |
| 64 | self.x0 = jnp.array([-0.5, 0.0, jnp.pi*3/2]) |
| 65 | self.xg = jnp.array([0.5, 0.0, 0.0]) |
| 66 | self.xref = jnp.load(f"{mbd.__path__[0]}/assets/car2d_xref.npy") |
| 67 | # self.xref = jnp.load(f"{mbd.__path__[0]}/../figure/car2d_xref.npy") |
| 68 | xref_diff = jnp.diff(self.xref, axis=0) |
| 69 | theta = jnp.arctan2(xref_diff[:, 0], xref_diff[:, 1]) |
| 70 | self.thetaref = jnp.append(theta, theta[-1]) |
| 71 | self.rew_xref = jax.vmap(self.get_reward)(self.xref).mean() |
| 72 | |
| 73 | def reset(self, rng: jax.Array): |
| 74 | """Resets the environment to an initial state.""" |
| 75 | return State(self.x0, self.x0, 0.0, 0.0) |
| 76 | |
| 77 | @partial(jax.jit, static_argnums=(0,)) |
| 78 | def step(self, state: State, action: jax.Array) -> State: |
| 79 | """Run one timestep of the environment's dynamics.""" |
| 80 | action = jnp.clip(action, -1.0, 1.0) |
| 81 | q = state.pipeline_state |
| 82 | q_new = rk4(car_dynamics, state.pipeline_state, action, self.dt) |
| 83 | collide = check_collision(q_new, self.obs_center, self.obs_radius) |
| 84 | q = jnp.where(collide, q, q_new) |
| 85 | reward = self.get_reward(q) |
| 86 | return state.replace(pipeline_state=q, obs=q, reward=reward, done=0.0) |
| 87 | |
| 88 | @partial(jax.jit, static_argnums=(0,)) |
| 89 | def get_reward(self, q): |
| 90 | reward = ( |
| 91 | 1.0 - (jnp.clip(jnp.linalg.norm(q[:2] - self.xg[:2]), 0.0, 0.2) / 0.2) ** 2 |
| 92 | ) |
| 93 | return reward |
| 94 | |
| 95 | @partial(jax.jit, static_argnums=(0,)) |
| 96 | def eval_xref_logpd(self, xs): |
| 97 | xs_err = xs[:, :2] - self.xref[:, :2] |
| 98 | # theta_err = xs[:, 3] - self.thetaref |
| 99 | logpd = 0.0-( |
| 100 | (jnp.clip(jnp.linalg.norm(xs_err, axis=-1), 0.0, 0.5) / 0.5) ** 2 |