Method reset

mbd/envs/pushT.py:22–37

(self, rng: jnp.ndarray)

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

    def reset(self, rng: jnp.ndarray) -> State:
        rng, rng_goal_xy = jax.random.split(rng)

        q = self.sys.init_q
        q = q.at[:2].set(jnp.array([0.1, -0.15]))
        q = q.at[5:].set(
            jax.random.uniform(rng_goal_xy, (3,), minval=-1.0, maxval=1.0)
            * jnp.array([0.2, 0.2, jnp.pi/4]) + jnp.array([-0.4, 0.4, jnp.pi])
        )
        qd = jnp.zeros(self.sys.qd_size())
        pipeline_state = self.pipeline_init(q, qd)
        obs = self._get_obs(pipeline_state)
        reward = self._get_reward(pipeline_state)
        done = self._get_done(pipeline_state)
        metrics = {}
        return State(pipeline_state, obs, reward, done, metrics)

    def step(self, state: State, action: jnp.ndarray) -> State:
        pipeline_state = self.pipeline_step(state.pipeline_state, action)