from contextlib import closing from io import StringIO from os import path from typing import Optional import numpy as np from gym import Env, logger, spaces from gym.envs.toy_text.utils import categorical_sample from gym.error import DependencyNotInstalled UP = 0 RIGHT = 1 DOWN = 2 LEFT = 3 class CliffWalkingEnv(Env): """ This is a simple implementation of the Gridworld Cliff reinforcement learning task. Adapted from Example 6.6 (page 106) from [Reinforcement Learning: An Introduction by Sutton and Barto](http://incompleteideas.net/book/bookdraft2018jan1.pdf). With inspiration from: [https://github.com/dennybritz/reinforcement-learning/blob/master/lib/envs/cliff_walking.py] (https://github.com/dennybritz/reinforcement-learning/blob/master/lib/envs/cliff_walking.py) ### Description The board is a 4x12 matrix, with (using NumPy matrix indexing): - [3, 0] as the start at bottom-left - [3, 11] as the goal at bottom-right - [3, 1..10] as the cliff at bottom-center If the agent steps on the cliff, it returns to the start. An episode terminates when the agent reaches the goal. ### Actions There are 4 discrete deterministic actions: - 0: move up - 1: move right - 2: move down - 3: move left ### Observations There are 3x12 + 1 possible states. In fact, the agent cannot be at the cliff, nor at the goal (as this results in the end of the episode). It remains all the positions of the first 3 rows plus the bottom-left cell. The observation is simply the current position encoded as [flattened index](https://numpy.org/doc/stable/reference/generated/numpy.unravel_index.html). ### Reward Each time step incurs -1 reward, and stepping into the cliff incurs -100 reward. ### Arguments ``` gym.make('CliffWalking-v0') ``` ### Version History - v0: Initial version release """ metadata = { "render_modes": ["human", "rgb_array", "ansi"], "render_fps": 4, } def __init__(self, render_mode: Optional[str] = None): self.shape = (4, 12) self.start_state_index = np.ravel_multi_index((3, 0), self.shape) self.nS = np.prod(self.shape) self.nA = 4 # Cliff Location self._cliff = np.zeros(self.shape, dtype=bool) self._cliff[3, 1:-1] = True # Calculate transition probabilities and rewards self.P = {} for s in range(self.nS): position = np.unravel_index(s, self.shape) self.P[s] = {a: [] for a in range(self.nA)} self.P[s][UP] = self._calculate_transition_prob(position, [-1, 0]) self.P[s][RIGHT] = self._calculate_transition_prob(position, [0, 1]) self.P[s][DOWN] = self._calculate_transition_prob(position, [1, 0]) self.P[s][LEFT] = self._calculate_transition_prob(position, [0, -1]) # Calculate initial state distribution # We always start in state (3, 0) self.initial_state_distrib = np.zeros(self.nS) self.initial_state_distrib[self.start_state_index] = 1.0 self.observation_space = spaces.Discrete(self.nS) self.action_space = spaces.Discrete(self.nA) self.render_mode = render_mode # pygame utils self.cell_size = (60, 60) self.window_size = ( self.shape[1] * self.cell_size[1], self.shape[0] * self.cell_size[0], ) self.window_surface = None self.clock = None self.elf_images = None self.start_img = None self.goal_img = None self.cliff_img = None self.mountain_bg_img = None self.near_cliff_img = None self.tree_img = None def _limit_coordinates(self, coord: np.ndarray) -> np.ndarray: """Prevent the agent from falling out of the grid world.""" coord[0] = min(coord[0], self.shape[0] - 1) coord[0] = max(coord[0], 0) coord[1] = min(coord[1], self.shape[1] - 1) coord[1] = max(coord[1], 0) return coord def _calculate_transition_prob(self, current, delta): """Determine the outcome for an action. Transition Prob is always 1.0. Args: current: Current position on the grid as (row, col) delta: Change in position for transition Returns: Tuple of ``(1.0, new_state, reward, terminated)`` """ new_position = np.array(current) + np.array(delta) new_position = self._limit_coordinates(new_position).astype(int) new_state = np.ravel_multi_index(tuple(new_position), self.shape) if self._cliff[tuple(new_position)]: return [(1.0, self.start_state_index, -100, False)] terminal_state = (self.shape[0] - 1, self.shape[1] - 1) is_terminated = tuple(new_position) == terminal_state return [(1.0, new_state, -1, is_terminated)] def step(self, a): transitions = self.P[self.s][a] i = categorical_sample([t[0] for t in transitions], self.np_random) p, s, r, t = transitions[i] self.s = s self.lastaction = a if self.render_mode == "human": self.render() return (int(s), r, t, False, {"prob": p}) def reset(self, *, seed: Optional[int] = None, options: Optional[dict] = None): super().reset(seed=seed) self.s = categorical_sample(self.initial_state_distrib, self.np_random) self.lastaction = None if self.render_mode == "human": self.render() return int(self.s), {"prob": 1} def render(self): if self.render_mode is None: logger.warn( "You are calling render method without specifying any render mode. " "You can specify the render_mode at initialization, " f'e.g. gym("{self.spec.id}", render_mode="rgb_array")' ) elif self.render_mode == "ansi": return self._render_text() else: return self._render_gui(self.render_mode) def _render_gui(self, mode): try: import pygame except ImportError: raise DependencyNotInstalled( "pygame is not installed, run `pip install gym[toy_text]`" ) if self.window_surface is None: pygame.init() if mode == "human": pygame.display.init() pygame.display.set_caption("CliffWalking") self.window_surface = pygame.display.set_mode(self.window_size) else: # rgb_array self.window_surface = pygame.Surface(self.window_size) if self.clock is None: self.clock = pygame.time.Clock() if self.elf_images is None: hikers = [ path.join(path.dirname(__file__), "img/elf_up.png"), path.join(path.dirname(__file__), "img/elf_right.png"), path.join(path.dirname(__file__), "img/elf_down.png"), path.join(path.dirname(__file__), "img/elf_left.png"), ] self.elf_images = [ pygame.transform.scale(pygame.image.load(f_name), self.cell_size) for f_name in hikers ] if self.start_img is None: file_name = path.join(path.dirname(__file__), "img/stool.png") self.start_img = pygame.transform.scale( pygame.image.load(file_name), self.cell_size ) if self.goal_img is None: file_name = path.join(path.dirname(__file__), "img/cookie.png") self.goal_img = pygame.transform.scale( pygame.image.load(file_name), self.cell_size ) if self.mountain_bg_img is None: bg_imgs = [ path.join(path.dirname(__file__), "img/mountain_bg1.png"), path.join(path.dirname(__file__), "img/mountain_bg2.png"), ] self.mountain_bg_img = [ pygame.transform.scale(pygame.image.load(f_name), self.cell_size) for f_name in bg_imgs ] if self.near_cliff_img is None: near_cliff_imgs = [ path.join(path.dirname(__file__), "img/mountain_near-cliff1.png"), path.join(path.dirname(__file__), "img/mountain_near-cliff2.png"), ] self.near_cliff_img = [ pygame.transform.scale(pygame.image.load(f_name), self.cell_size) for f_name in near_cliff_imgs ] if self.cliff_img is None: file_name = path.join(path.dirname(__file__), "img/mountain_cliff.png") self.cliff_img = pygame.transform.scale( pygame.image.load(file_name), self.cell_size ) for s in range(self.nS): row, col = np.unravel_index(s, self.shape) pos = (col * self.cell_size[0], row * self.cell_size[1]) check_board_mask = row % 2 ^ col % 2 self.window_surface.blit(self.mountain_bg_img[check_board_mask], pos) if self._cliff[row, col]: self.window_surface.blit(self.cliff_img, pos) if row < self.shape[0] - 1 and self._cliff[row + 1, col]: self.window_surface.blit(self.near_cliff_img[check_board_mask], pos) if s == self.start_state_index: self.window_surface.blit(self.start_img, pos) if s == self.nS - 1: self.window_surface.blit(self.goal_img, pos) if s == self.s: elf_pos = (pos[0], pos[1] - 0.1 * self.cell_size[1]) last_action = self.lastaction if self.lastaction is not None else 2 self.window_surface.blit(self.elf_images[last_action], elf_pos) if mode == "human": pygame.event.pump() pygame.display.update() self.clock.tick(self.metadata["render_fps"]) else: # rgb_array return np.transpose( np.array(pygame.surfarray.pixels3d(self.window_surface)), axes=(1, 0, 2) ) def _render_text(self): outfile = StringIO() for s in range(self.nS): position = np.unravel_index(s, self.shape) if self.s == s: output = " x " # Print terminal state elif position == (3, 11): output = " T " elif self._cliff[position]: output = " C " else: output = " o " if position[1] == 0: output = output.lstrip() if position[1] == self.shape[1] - 1: output = output.rstrip() output += "\n" outfile.write(output) outfile.write("\n") with closing(outfile): return outfile.getvalue()