Jax Cliffwalking Env

[balisujohn]

Description

This adds a CliffWalkingJaxEnv to the tabular environment suite, as well as the underlying FuncEnv subclass CliffWalkingFunctional.

CliffWalkingJaxEnv is registered with the name "Jax-CliffWalking-v0" This PR still needs a bit of polish work, but you can try it out for yourself in GUI mode with

python3 ./gymnasium/envs/tabular/cliffwalking.py

Type of change

Please delete options that are not relevant.

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• This change requires a documentation update

Checklist:

• I have run the pre-commit checks with pre-commit run --all-files (see CONTRIBUTING.md instructions to set it up)
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[pseudo-rnd-thoughts]

Would we want the RenderStateType to be a NamedTuple?

[balisujohn]

changed

[pseudo-rnd-thoughts]

This looks good, and I think close to being merged. Thanks for your hard work.

I was able to get the type issue solved using from __future__ import annotations and if TYPE_CHECKING: import pygame

from __future__ import annotations

from os import path
from typing import TYPE_CHECKING, NamedTuple

import jax
import jax.numpy as jnp
import numpy as np
from jax.random import PRNGKey

from gymnasium import spaces
from gymnasium.error import DependencyNotInstalled
from gymnasium.experimental.functional import ActType, FuncEnv, StateType
from gymnasium.experimental.functional_jax_env import FunctionalJaxEnv
from gymnasium.utils import EzPickle
from gymnasium.wrappers import HumanRendering


if TYPE_CHECKING:
    import pygame


class RenderStateType(NamedTuple):
    """A named tuple which contains the full render state of the Cliffwalking Env. This is static during the episode."""

    screen: pygame.surface
    shape: tuple[int, int]
    nS: int
    cell_size: tuple[int, int]
    cliff: np.ndarray
    elf_images: tuple[pygame.Surface, pygame.Surface, pygame.Surface, pygame.Surface]
    start_img: pygame.Surface
    goal_img: pygame.Surface
    bg_imgs: tuple[str, str]
    mountain_bg_img: tuple[pygame.Surface, pygame.Surface]
    near_cliff_imgs: tuple[str, str]
    near_cliff_img: tuple[pygame.Surface, pygame.Surface]
    cliff_img: pygame.Surface

[balisujohn]

added that typing fix