| 9 | |
| 10 | |
| 11 | class Cartpole(PipelineEnv): |
| 12 | def __init__(self, backend="positional", **kwargs): |
| 13 | sys = mjcf.load(f"{mbd.__path__[0]}/assets/cartpole.xml") |
| 14 | |
| 15 | n_frames = 2 |
| 16 | |
| 17 | if backend in ["spring", "positional"]: |
| 18 | sys = sys.replace(dt=0.005) |
| 19 | n_frames = 4 |
| 20 | |
| 21 | kwargs["n_frames"] = kwargs.get("n_frames", n_frames) |
| 22 | |
| 23 | super().__init__(sys=sys, backend=backend, **kwargs) |
| 24 | |
| 25 | def reset(self, rng: jax.Array) -> State: |
| 26 | """Resets the environment to an initial state.""" |
| 27 | rng, rng1, rng2 = jax.random.split(rng, 3) |
| 28 | |
| 29 | q = self.sys.init_q + jax.random.uniform( |
| 30 | rng1, (self.sys.q_size(),), minval=-0.01, maxval=0.01 |
| 31 | ) + jp.array([0.0, jp.pi]) |
| 32 | qd = jax.random.uniform(rng2, (self.sys.qd_size(),), minval=-0.01, maxval=0.01) |
| 33 | pipeline_state = self.pipeline_init(q, qd) |
| 34 | obs = self._get_obs(pipeline_state) |
| 35 | reward, done = jp.zeros(2) |
| 36 | metrics = {} |
| 37 | |
| 38 | return State(pipeline_state, obs, reward, done, metrics) |
| 39 | |
| 40 | def step(self, state: State, action: jax.Array) -> State: |
| 41 | """Run one timestep of the environment's dynamics.""" |
| 42 | pipeline_state = self.pipeline_step(state.pipeline_state, action) |
| 43 | obs = self._get_obs(pipeline_state) |
| 44 | reward = jp.cos(pipeline_state.q[1]) - jp.abs(pipeline_state.qd[0]) |
| 45 | done = 0.0 |
| 46 | return state.replace( |
| 47 | pipeline_state=pipeline_state, obs=obs, reward=reward, done=done |
| 48 | ) |
| 49 | |
| 50 | @property |
| 51 | def action_size(self): |
| 52 | return 1 |
| 53 | |
| 54 | def _get_obs(self, pipeline_state: base.State) -> jax.Array: |
| 55 | """Observe cartpole body position and velocities.""" |
| 56 | return jp.concatenate([pipeline_state.q, pipeline_state.qd]) |