| 9 | |
| 10 | |
| 11 | class Hopper(PipelineEnv): |
| 12 | def __init__(self): |
| 13 | path = epath.resource_path("brax") / "envs/assets/hopper.xml" |
| 14 | sys = mjcf.load(path) |
| 15 | |
| 16 | self._reset_noise_scale = 5e-3 |
| 17 | |
| 18 | super().__init__(sys=sys, backend="positional", n_frames=20) |
| 19 | |
| 20 | def reset(self, rng: jax.Array) -> State: |
| 21 | """Resets the environment to an initial state.""" |
| 22 | rng, rng1, rng2 = jax.random.split(rng, 3) |
| 23 | |
| 24 | low, hi = -self._reset_noise_scale, self._reset_noise_scale |
| 25 | qpos = self.sys.init_q + jax.random.uniform( |
| 26 | rng1, (self.sys.q_size(),), minval=low, maxval=hi |
| 27 | ) |
| 28 | qvel = jax.random.uniform(rng2, (self.sys.qd_size(),), minval=low, maxval=hi) |
| 29 | |
| 30 | pipeline_state = self.pipeline_init(qpos, qvel) |
| 31 | |
| 32 | obs = self._get_obs(pipeline_state) |
| 33 | reward, done = jp.zeros(2) |
| 34 | return State(pipeline_state, obs, reward, done, {}) |
| 35 | |
| 36 | def step(self, state: State, action: jax.Array) -> State: |
| 37 | """Runs one timestep of the environment's dynamics.""" |
| 38 | pipeline_state0 = state.pipeline_state |
| 39 | assert pipeline_state0 is not None |
| 40 | pipeline_state = self.pipeline_step(pipeline_state0, action) |
| 41 | |
| 42 | obs = self._get_obs(pipeline_state) |
| 43 | reward = self._get_reward(pipeline_state) |
| 44 | |
| 45 | return state.replace( |
| 46 | pipeline_state=pipeline_state, obs=obs, reward=reward, done=0.0 |
| 47 | ) |
| 48 | |
| 49 | def _get_obs(self, pipeline_state: base.State) -> jax.Array: |
| 50 | """Returns the environment observations.""" |
| 51 | position = pipeline_state.q |
| 52 | position = position.at[1].set(pipeline_state.x.pos[0, 2]) |
| 53 | velocity = jp.clip(pipeline_state.qd, -10, 10) |
| 54 | |
| 55 | return jp.concatenate((position, velocity)) |
| 56 | |
| 57 | def _get_reward(self, pipeline_state: base.State) -> jax.Array: |
| 58 | """Returns the environment reward.""" |
| 59 | return ( |
| 60 | pipeline_state.x.pos[0, 0] + |
| 61 | # pipeline_state.x.pos[1, 0] + |
| 62 | # pipeline_state.x.pos[2, 0] + |
| 63 | # pipeline_state.x.pos[3, 0] + |
| 64 | - jp.clip(jp.abs(pipeline_state.x.pos[0, 2] - 1.0), -1.0, 1.0) * 0.5 |
| 65 | ) |