如何扩展OpenAI Gym中Frozen-Lake v0环境的观测空间？

Absolutely! You can totally extend the observation space of FrozenLake-v0 using Discrete, Box, MultiDiscrete, or even custom Gym space types—here's a breakdown of how to implement each, with practical code examples to get you started:

1. Using Discrete (Simplest Extension) #

The vanilla FrozenLake-v0 uses a Discrete(16) space, representing the 16 grid positions. If you want to add binary or finite-state info (like whether the agent picked up a key), you can expand this into a larger discrete space by combining the original state with your new variables.

For example, let's add a "has key" state:

import gym
from gym.envs.toy_text.frozen_lake import FrozenLakeEnv
from gym.spaces import Discrete

class FrozenLakeWithKeyEnv(FrozenLakeEnv):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        # Combine 16 grid positions with 2 key states (0=no, 1=yes)
        self.observation_space = Discrete(16 * 2)
        self.has_key = False

    def reset(self, **kwargs):
        base_obs = super().reset(**kwargs)
        self.has_key = False
        # Encode combined state into a single integer
        return base_obs * 2 + (1 if self.has_key else 0)

    def step(self, action):
        base_obs, reward, done, info = super().step(action)
        # Example: Pick up key when stepping on grid position 5
        if base_obs == 5 and not self.has_key:
            self.has_key = True
        # Return the encoded combined observation
        combined_obs = base_obs * 2 + (1 if self.has_key else 0)
        return combined_obs, reward, done, info

2. Using MultiDiscrete (For Separate Discrete Dimensions) #

If you prefer to keep state variables separate instead of encoding them into one integer, MultiDiscrete is perfect. It lets you define multiple discrete dimensions, each with their own range.

Let's add both a "has key" flag and a step counter:

import gym
from gym.envs.toy_text.frozen_lake import FrozenLakeEnv
from gym.spaces import MultiDiscrete

class FrozenLakeMultiDiscreteEnv(FrozenLakeEnv):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        # Define discrete dimensions:
        # - Grid position: 0-15
        # - Has key: 0 (no) or 1 (yes)
        # - Steps taken: 0-100 (arbitrary max limit)
        self.observation_space = MultiDiscrete([16, 2, 101])
        self.has_key = False
        self.steps_taken = 0

    def reset(self, **kwargs):
        base_obs = super().reset(**kwargs)
        self.has_key = False
        self.steps_taken = 0
        # Return a list of discrete values for each dimension
        return [base_obs, 1 if self.has_key else 0, self.steps_taken]

    def step(self, action):
        base_obs, reward, done, info = super().step(action)
        self.steps_taken += 1
        if base_obs == 5 and not self.has_key:
            self.has_key = True
        return [base_obs, 1 if self.has_key else 0, self.steps_taken], reward, done, info

3. Using Box (For Continuous State Information) #

If you want to include continuous values (like distance to the goal, or a simulated "fatigue" metric), Box is the way to go. This space defines a n-dimensional continuous range with low and high bounds.

Let's add grid coordinates and Euclidean distance to the goal:

import gym
import numpy as np
from gym.envs.toy_text.frozen_lake import FrozenLakeEnv
from gym.spaces import Box

class FrozenLakeBoxEnv(FrozenLakeEnv):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        # Define Box space with 4 dimensions:
        # - X coordinate (0-3)
        # - Y coordinate (0-3)
        # - Distance to goal (0 ~ 4.24)
        # - Has key (0.0 or 1.0, treated as continuous)
        self.observation_space = Box(
            low=np.array([0, 0, 0.0, 0.0]),
            high=np.array([3, 3, 4.25, 1.0]),
            dtype=np.float32
        )
        self.has_key = False

    def _get_xy(self, obs):
        # Convert flat grid index to (x, y) coordinates
        return (obs // 4, obs % 4)

    def reset(self, **kwargs):
        base_obs = super().reset(**kwargs)
        self.has_key = False
        x, y = self._get_xy(base_obs)
        goal_x, goal_y = self._get_xy(self.nS - 1)  # Goal is the last grid state
        distance = np.linalg.norm(np.array([x - goal_x, y - goal_y]))
        return np.array([x, y, distance, 1.0 if self.has_key else 0.0], dtype=np.float32)

    def step(self, action):
        base_obs, reward, done, info = super().step(action)
        x, y = self._get_xy(base_obs)
        goal_x, goal_y = self._get_xy(self.nS - 1)
        distance = np.linalg.norm(np.array([x - goal_x, y - goal_y]))
        if base_obs == 5 and not self.has_key:
            self.has_key = True
        return np.array([x, y, distance, 1.0 if self.has_key else 0.0], dtype=np.float32), reward, done, info

Quick Notes #

• Always make sure the observations returned by reset() and step() strictly match the bounds/rules of your defined observation_space—Gym will throw errors if you don't.
• For even more flexibility, you can use Tuple or Dict spaces to combine different space types (e.g., a Dict with a Discrete position and a Box distance metric).

内容的提问来源于stack exchange，提问作者user11602789