Sets the state of the environment at the start of each episode.
(self, physics)
| 175 | self.max_reward = 4 |
| 176 | |
| 177 | def initialize_episode(self, physics): |
| 178 | """Sets the state of the environment at the start of each episode.""" |
| 179 | # TODO Notice: this function does not randomize the env configuration. Instead, set BOX_POSE from outside |
| 180 | # reset qpos, control and box position |
| 181 | with physics.reset_context(): |
| 182 | physics.named.data.qpos[:16] = START_ARM_POSE |
| 183 | np.copyto(physics.data.ctrl, START_ARM_POSE) |
| 184 | assert BOX_POSE[0] is not None |
| 185 | physics.named.data.qpos[-7*2:] = BOX_POSE[0] # two objects |
| 186 | # print(f"{BOX_POSE=}") |
| 187 | super().initialize_episode(physics) |
| 188 | |
| 189 | @staticmethod |
| 190 | def get_env_state(physics): |
no outgoing calls
no test coverage detected