def reset()

in gym_hil/envs/panda_arrange_boxes_gym_env.py [0:0]

• 15 lines of code
• 3 McCabe index (conditional complexity)

    def reset(self, seed=None, **kwargs) -> Tuple[Dict[str, np.ndarray], Dict[str, Any]]:
        """Reset the environment."""
        # Ensure gymnasium internal RNG is initialized when a seed is provided
        super().reset(seed=seed)

        mujoco.mj_resetData(self._model, self._data)

        # Reset the robot to home position
        self.reset_robot()

        positions_coords = np.linspace(-self.block_range, self.block_range, self.no_blocks)
        np.random.shuffle(positions_coords)

        # Sample a new block position
        blocks = [f"block{i}" for i in range(1, self.no_blocks + 1)]
        np.random.shuffle(blocks)

        for block, pos in zip(blocks, positions_coords, strict=False):
            block_x_coord = self._data.joint(block).qpos[0]
            block_coords = np.array([block_x_coord, pos])
            self._data.joint(block).qpos[:3] = (*block_coords, self._block_z)

        mujoco.mj_forward(self._model, self._data)
        obs = self._compute_observation()

        return obs, {}