| 11 | |
| 12 | |
| 13 | class HumanoidTrack(PipelineEnv): |
| 14 | |
| 15 | def __init__(self, mode="jog"): |
| 16 | sys = mjcf.load(f"{mbd.__path__[0]}/assets/humanoidtrack.xml") |
| 17 | self.H = 50 # traj time 1.5s |
| 18 | body_names = [ |
| 19 | "torso", |
| 20 | "left_thigh", |
| 21 | "right_thigh", |
| 22 | "left_shin", |
| 23 | "right_shin", |
| 24 | ] |
| 25 | self.track_body_names = body_names |
| 26 | self.track_body_idx = jnp.array( |
| 27 | [sys.link_names.index(name) for name in self.track_body_names] |
| 28 | ) |
| 29 | self.ref_body_names = [name + "_ref" for name in body_names] |
| 30 | self.ref_body_idx = jnp.array( |
| 31 | [sys.link_names.index(name) for name in self.ref_body_names] |
| 32 | ) |
| 33 | with open(f"{mbd.__path__[0]}/assets/jog_xref.pkl", "rb") as f: |
| 34 | xs_demo_dict = pickle.load(f) |
| 35 | xref = [] |
| 36 | for name in body_names: |
| 37 | x = xs_demo_dict[name] |
| 38 | if len(x) < self.H: |
| 39 | x = jnp.concatenate([x, jnp.tile(x[-1:], (self.H - len(x), 1))], axis=0) |
| 40 | else: |
| 41 | x = x[70 : (self.H + 70)] |
| 42 | xref.append(x) |
| 43 | self.xref = jnp.stack(xref, axis=0) |
| 44 | self.rew_xref = 1.0 |
| 45 | |
| 46 | super().__init__(sys=sys, backend="positional", n_frames=5) |
| 47 | |
| 48 | def reset(self, rng: jax.Array) -> State: |
| 49 | """Resets the environment to an initial state.""" |
| 50 | |
| 51 | qpos = self.sys.init_q |
| 52 | qvel = jnp.zeros(self.sys.qd_size()) |
| 53 | |
| 54 | pipeline_state = self.pipeline_init(qpos, qvel) |
| 55 | obs = self._get_obs(pipeline_state) |
| 56 | reward, done, zero = jnp.zeros(3) |
| 57 | metrics = { |
| 58 | "reward_linup": zero, |
| 59 | "reward_quadctrl": zero, |
| 60 | } |
| 61 | return State(pipeline_state, obs, reward, done, metrics) |
| 62 | |
| 63 | def step(self, state: State, action: jax.Array) -> State: |
| 64 | """Runs one timestep of the environment's dynamics.""" |
| 65 | pipeline_state = self.pipeline_step(state.pipeline_state, action) |
| 66 | # set reference state for visualization |
| 67 | for i, idx in enumerate(self.ref_body_idx): |
| 68 | pipeline_state = pipeline_state.replace( |
| 69 | x=pipeline_state.x.replace( |
| 70 | pos=pipeline_state.x.pos.at[idx].set( |