def _initialize_observation_space()

in gym_xarm/tasks/base.py [0:0]

• 21 lines of code
• 4 McCabe index (conditional complexity)

    def _initialize_observation_space(self):
        image_shape = (self.observation_height, self.observation_width, 3)
        obs = self.get_obs()
        if self.obs_type == "state":
            observation_space = gym.spaces.Box(-1000.0, 1000.0, shape=obs.shape, dtype=np.float64)
        elif self.obs_type == "pixels":
            observation_space = gym.spaces.Box(low=0, high=255, shape=image_shape, dtype=np.uint8)
        elif self.obs_type == "pixels_agent_pos":
            observation_space = gym.spaces.Dict(
                {
                    "pixels": gym.spaces.Box(low=0, high=255, shape=image_shape, dtype=np.uint8),
                    "agent_pos": gym.spaces.Box(
                        low=-1000.0, high=1000.0, shape=obs["agent_pos"].shape, dtype=np.float64
                    ),
                }
            )
        else:
            raise ValueError(
                f"Unknown obs_type {self.obs_type}. Must be one of [pixels, state, pixels_agent_pos]"
            )

        return observation_space