Method __init__

mbd/envs/humanoidtrack.py:15–46

(self, mode="jog")

class HumanoidTrack(PipelineEnv):

    def __init__(self, mode="jog"):
        sys = mjcf.load(f"{mbd.__path__[0]}/assets/humanoidtrack.xml")
        self.H = 50  # traj time 1.5s
        body_names = [
            "torso",
            "left_thigh",
            "right_thigh",
            "left_shin",
            "right_shin",
        ]
        self.track_body_names = body_names
        self.track_body_idx = jnp.array(
            [sys.link_names.index(name) for name in self.track_body_names]
        )
        self.ref_body_names = [name + "_ref" for name in body_names]
        self.ref_body_idx = jnp.array(
            [sys.link_names.index(name) for name in self.ref_body_names]
        )
        with open(f"{mbd.__path__[0]}/assets/jog_xref.pkl", "rb") as f:
            xs_demo_dict = pickle.load(f)
        xref = []
        for name in body_names:
            x = xs_demo_dict[name]
            if len(x) < self.H:
                x = jnp.concatenate([x, jnp.tile(x[-1:], (self.H - len(x), 1))], axis=0)
            else:
                x = x[70 : (self.H + 70)]
            xref.append(x)
        self.xref = jnp.stack(xref, axis=0)
        self.rew_xref = 1.0

        super().__init__(sys=sys, backend="positional", n_frames=5)

    def reset(self, rng: jax.Array) -> State:
        """Resets the environment to an initial state."""