| 10 | |
| 11 | |
| 12 | class HumanoidStandup(PipelineEnv): |
| 13 | |
| 14 | def __init__(self): |
| 15 | sys = mjcf.load(f"{mbd.__path__[0]}/assets/humanoidstandup.xml") |
| 16 | |
| 17 | super().__init__(sys=sys, backend="positional", n_frames=7) |
| 18 | |
| 19 | def reset(self, rng: jax.Array) -> State: |
| 20 | """Resets the environment to an initial state.""" |
| 21 | rng, rng1, rng2 = jax.random.split(rng, 3) |
| 22 | |
| 23 | low, hi = -0.01, 0.01 |
| 24 | qpos = self.sys.init_q + jax.random.uniform( |
| 25 | rng1, (self.sys.q_size(),), minval=-0.01, maxval=0.01 |
| 26 | ) |
| 27 | qvel = jax.random.uniform(rng2, (self.sys.qd_size(),), minval=low, maxval=hi) |
| 28 | |
| 29 | pipeline_state = self.pipeline_init(qpos, qvel) |
| 30 | obs = self._get_obs(pipeline_state, jnp.zeros(self.sys.act_size())) |
| 31 | reward, done, zero = jnp.zeros(3) |
| 32 | metrics = { |
| 33 | "reward_linup": zero, |
| 34 | "reward_quadctrl": zero, |
| 35 | } |
| 36 | return State(pipeline_state, obs, reward, done, metrics) |
| 37 | |
| 38 | def step(self, state: State, action: jax.Array) -> State: |
| 39 | """Runs one timestep of the environment's dynamics.""" |
| 40 | pipeline_state = self.pipeline_step(state.pipeline_state, action) |
| 41 | |
| 42 | obs = self._get_obs(pipeline_state, action) |
| 43 | reward = self._get_reward(pipeline_state) |
| 44 | |
| 45 | return state.replace(pipeline_state=pipeline_state, obs=obs, reward=reward) |
| 46 | |
| 47 | def _get_obs(self, pipeline_state: base.State, action: jax.Array) -> jax.Array: |
| 48 | return jnp.concatenate([pipeline_state.q, pipeline_state.qd], axis=-1) |
| 49 | |
| 50 | def _get_reward(self, pipeline_state: base.State) -> jax.Array: |
| 51 | return ( |
| 52 | 1.5 |
| 53 | - jnp.clip(jnp.abs(pipeline_state.x.pos[0, 2] - 1.3), -2.0, 1.0) |
| 54 | - jnp.abs(pipeline_state.x.pos[0, 0]) * 0.1 |
| 55 | - jnp.abs(pipeline_state.x.pos[0, 1]) * 0.1 |
| 56 | ) |