Repository: Farama-Foundation/miniwob-plusplus (language: python)

Issue:
[Proposal] Port MiniWoBEnvironment to the gym.Env interface
### Proposal 

Port MiniWoBEnvironment to the gym.Env interface.

### Motivation

This is an effort toward making the MiniWoB codebase compatible with Gymnasium.

### Pitch

Perform the following steps:
- Convert MiniWoBEnvironment to the gym.Env interface, following the instructions on the [Make your own custom environment
](https://www.gymlibrary.dev/content/environment_creation/) page.
- Make sure it passes the test in `gymnasium.utils.env_checker.check_env`.

### Alternatives

None. This is a required [project standard](https://farama.org/project_standards).

### Additional context

None.

### Checklist

- [x] I have checked that there is no similar [issue](https://github.com/Farama-Foundation/miniwob-plusplus/issues) in the repo (**required**)



Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 50a202f0..a292861b 100644
--- a/README.md
+++ b/README.md
@@ -50,6 +50,9 @@ There are 2 ways to access the tasks:
 
 # Python Environment
 
+The Python environment implements the
+[Gymnasium](https://github.com/Farama-Foundation/Gymnasium/) interface.
+
 ## Setup
 
 - Python dependencies
@@ -71,9 +74,10 @@ There are 2 ways to access the tasks:
     conda install -c conda-forge selenium
     ```
 
-- Environment variable:
-  The tests depend on the variable `MINIWOB_BASE_URL` to point to the base URL
-  of the benchmark. Please see the main README on how to set up the tasks.
+- Environment variable (optional):
+  One can set the shell variable `MINIWOB_BASE_URL` to point to the base URL
+  of the benchmark. If the variable is not specified, the `file://` URL will be
+  inferred from the module path.
   - For the server setup:
     ```
     export MINIWOB_BASE_URL=http://localhost:8080/
@@ -89,13 +93,14 @@ There are 2 ways to access the tasks:
 pytest
 ```
 
-The tests in the `tests` directory also illustrate how the code can be used.
+The tests in the `miniwob/tests/` directory also illustrate how the code can be used.
 
 ---
 
 # JavaScript Utilities
 
-This version of MiniWoB incorporates a few additional JavaScript utilities.
+This version of MiniWoB environments incorporates a few JavaScript utilities,
+many of which are used by the Python interface.
 
 ## `Math.seedrandom([seed])`
 
diff --git a/miniwob/__init__.py b/miniwob/__init__.py
index e69de29b..5b32fdc8 100644
--- a/miniwob/__init__.py
+++ b/miniwob/__init__.py
@@ -0,0 +1,3 @@
+from miniwob.registration import register_miniwob_envs
+
+register_miniwob_envs()
diff --git a/miniwob/action.py b/miniwob/action.py
index 0ed7cf08..b129968e 100644
--- a/miniwob/action.py
+++ b/miniwob/action.py
@@ -1,257 +1,153 @@
-# MiniWoB actions
-import abc
+"""MiniWoB action space."""
 import logging
+from enum import IntEnum
 
+import numpy as np
+from gymnasium import spaces
 from selenium.webdriver.common.action_chains import ActionChains
 from selenium.webdriver.common.by import By
 
 
-class MiniWoBAction(metaclass=abc.ABCMeta):
-    """Defines an action in its __call__ method."""
-
-    @abc.abstractmethod
-    def __call__(self, driver):
-        """Performs the action defined by this class on the driver.
-
-        Args:
-            driver (Selenium WebDriver)
-        """
-        raise NotImplementedError()
-
-    def to_dict(self):
-        """Dict representation for JSON serialization."""
-        raise NotImplementedError()
-
-
-class MiniWoBTerminate(MiniWoBAction):
-    """Immediately fails the task.
-
-    This is done via a JavaScript call.
-    """
-
-    def __call__(self, driver):
-        driver.execute_script('return core.endEpisode(-1,false,"terminate");')
-
-    def __str__(self):
-        return "MiniWoBTerminate"
-
-    __repr__ = __str__
-
-    def __eq__(self, other):
-        return isinstance(other, MiniWoBTerminate)
-
-    def __hash__(self):
-        return hash(self.__class__.__name__)
-
-    def to_dict(self):
-        return {"type": "Terminate"}
-
-
-class MiniWoBCoordClick(MiniWoBAction):
-    """Defines a click action left pixels from the left of the screen and top
-    pixels from the top of the screen.
-
-    This is done via Selenium.
-
-    Args:
-        left (int): number of pixels from the left of the screen
-        top (int): number of pixels from the top of the screen
-    """
-
-    def __init__(self, left, top):
-        self._left = left
-        self._top = top
-
-    def __call__(self, driver):
-        """Clicks at coordinates (left, top)"""
-        body = driver.find_element(By.TAG_NAME, "body")
-        # The offset is from the center, not top-left.
-        x = -body.size["width"] / 2 + self.left
-        y = -body.size["height"] / 2 + self.top
-        chain = ActionChains(driver)
-        chain.move_to_element_with_offset(body, x, y).click().perform()
-
-    @property
-    def left(self):
-        return self._left
-
-    @property
-    def top(self):
-        return self._top
-
-    def __str__(self):
-        return f"CoordClick(coords: ({self.left}, {self.top}))"
-
-    __repr__ = __str__
-
-    def __eq__(self, other):
-        if not isinstance(other, MiniWoBCoordClick):
-            return False
-        return (self.left, self.top) == (other.left, other.top)
-
-    def __hash__(self):
-        return hash((self.__class__.__name__, self.left, self.top))
-
-    def to_dict(self):
-        return {"type": "CoordClick", "left": self._left, "top": self._top}
-
-
-class MiniWoBElementClick(MiniWoBAction):
-    """An action that clicks on a DOM element regardless of its position
-    or visibility.
-
-    This is done via a JavaScript call.
-
-    Args:
-        element: One of the following:
-            - the DOMElement object to click
-            - ref (int) of the DOMElement object to click
-        fail_hard (bool): If True, throw an error when the click cannot
-            be successfully performed
-    """
-
-    def __init__(self, element, fail_hard=False):
-        self._element = element
-        self._ref = element.ref
-        if self._ref is None:
-            raise RuntimeError(f"Cannot click on {element} with _ref = None")
-        self._fail_hard = fail_hard
-
-    def __call__(self, driver):
-        if self.element.tag == "select":
-            # SPECIAL CASE: <select>
-            body = driver.find_element(By.TAG_NAME, "body")
-            chain = ActionChains(driver)
-            chain.move_to_element_with_offset(
-                body, self.element.left + 5, self.element.top + 5
-            ).click().perform()
-            return
-        result = driver.execute_script(f"return core.elementClick({self._ref});")
-        if result is not True:
-            if self._fail_hard:
-                raise RuntimeError(f"{self} failed: {result}")
-            else:
-                logging.warning("%s failed: %s", self, result)
-
-    @property
-    def ref(self):
-        return self._ref
-
-    @property
-    def element(self):
-        return self._element
-
-    def __str__(self):
-        return f"click({self.element})"
-
-    __repr__ = __str__
-
-    def __eq__(self, other):
-        """Compare based on element refs."""
-        if not isinstance(other, MiniWoBElementClick):
-            return False
-        return (self.ref, self._fail_hard) == (other.ref, other._fail_hard)
-
-    def __hash__(self):
-        return hash((self.__class__.__name__, self.ref, self._fail_hard))
-
-    def to_dict(self):
-        return {
-            "type": "ElementClick",
-            "element": self.element.to_dict(),
+class ActionTypes(IntEnum):
+    NONE = 0
+    COORD_CLICK = 1
+    ELEMENT_CLICK = 2
+    TYPE = 3
+    FOCUS_AND_TYPE = 4
+
+
+TYPING_MAX_LENGTH = 32
+ASCII_CHARSET = frozenset(chr(x) for x in range(32, 128))
+MAX_REF = 1000000
+
+
+def get_action_space(screen_width, screen_height):
+    """Return the space of serialized actions."""
+    space = spaces.Dict(
+        {
+            "action_type": spaces.Discrete(len(ActionTypes)),
+            # coords (left, top) is used for COORD_CLICK
+            "coords": spaces.Box(
+                np.array([0.0, 0.0], dtype=np.float32),
+                np.array([screen_width, screen_height], dtype=np.float32),
+            ),
+            # ref (element ref ID) is used for ELEMENT_CLICK and FOCUS_AND_TYPE
+            "ref": spaces.Discrete(MAX_REF, start=1),
+            # text is only used for TYPE and FOCUS_AND_TYPE
+            "text": spaces.Text(TYPING_MAX_LENGTH, charset=ASCII_CHARSET),
         }
+    )
+    return space
+
+
+def create_none_action():
+    """Return a valid action object that does nothing."""
+    return {
+        "action_type": ActionTypes.NONE,
+        "coords": np.zeros(2, dtype=np.float32),
+        "ref": 1,
+        "text": " ",
+    }
+
+
+def create_coord_click_action(left, top):
+    """Return a valid action object with type COORD_CLICK."""
+    action = create_none_action()
+    action.update(
+        {
+            "action_type": ActionTypes.COORD_CLICK,
+            "coords": np.array([left, top], dtype=np.float32),
+        }
+    )
+    return action
 
 
-class MiniWoBType(MiniWoBAction):
-    """An action that sends keystrokes to the focused element.
-
-    This is done via Selenium.
-
-    Args:
-        text (str or list[str]): Things to type.
-            Non-printable characters defined in
-            selenium.webdriver.common.keys.Keys can also be used to send
-            special keys (arrows, backspace, etc.)
-    """
-
-    def __init__(self, text):
-        self._text = text
-
-    def __call__(self, driver):
-        chain = ActionChains(driver)
-        chain.send_keys(self._text)
-        chain.perform()
-
-    @property
-    def text(self):
-        return self._text
-
-    def __str__(self):
-        return f"Type({repr(self._text)})"
-
-    __repr__ = __str__
-
-    def __eq__(self, other):
-        if not isinstance(other, MiniWoBType):
-            return False
-        return self.text == other.text
-
-    def __hash__(self):
-        return hash((self.__class__.__name__, self.text))
-
-    def to_dict(self):
-        return {"type": "Type", "text": self.text}
-
-
-class MiniWoBFocusAndType(MiniWoBAction):
-    """An action that first performs an element click to focus an element,
-    followed by sending keystrokes.
-
-    Implemented by executing a MiniWoBElementClick followed by MiniWoBType.
-
-    Args:
-        element: One of the following:
-            - the DOMElement object to click
-            - ref (int) of the DOMElement object to click
-        text (str or list[str]): Things to type.
-    """
-
-    def __init__(self, element, text):
-        self._click = MiniWoBElementClick(element)
-        self._type = MiniWoBType(text)
-
-    def __call__(self, driver):
-        self._click(driver)
-        self._type(driver)
-
-    @property
-    def ref(self):
-        return self._click.ref
-
-    @property
-    def element(self):
-        return self._click.element
-
-    @property
-    def text(self):
-        return self._type.text
-
-    def __str__(self):
-        return f"type({self._click.element}, {repr(self._type.text)})"
-
-    __repr__ = __str__
-
-    def __eq__(self, other):
-        if not isinstance(other, MiniWoBFocusAndType):
-            return False
-        return (self._click, self._type) == (other._click, other._type)
+def create_element_click_action(ref):
+    """Return a valid action object with type ELEMENT_CLICK."""
+    action = create_none_action()
+    action.update(
+        {
+            "action_type": ActionTypes.ELEMENT_CLICK,
+            "ref": ref,
+        }
+    )
+    return action
 
-    def __hash__(self):
-        return hash((self._click, self._type))
 
-    def to_dict(self):
-        return {
-            "type": "ElementClick",
-            "element": self.element.to_dict(),
-            "text": self.text,
+def create_type_action(text):
+    """Return a valid action object with type TYPE."""
+    action = create_none_action()
+    action.update(
+        {
+            "action_type": ActionTypes.TYPE,
+            "text": text,
+        }
+    )
+    return action
+
+
+def create_focus_and_type_action(ref, text):
+    """Return a valid action object with type FOCUS_AND_TYPE."""
+    action = create_none_action()
+    action.update(
+        {
+            "action_type": ActionTypes.FOCUS_AND_TYPE,
+            "ref": ref,
+            "text": text,
         }
+    )
+    return action
+
+
+def execute_coord_click(left, top, driver):
+    """Click at coordinates (left, top)."""
+    body = driver.find_element(By.TAG_NAME, "body")
+    # The offset is from the center, not top-left.
+    x = -body.size["width"] / 2 + left
+    y = -body.size["height"] / 2 + top
+    chain = ActionChains(driver)
+    chain.move_to_element_with_offset(body, x, y).click().perform()
+
+
+def execute_element_click(ref, driver):
+    """Click on the DOM element specified by a ref ID."""
+    # TODO: Handle <select> correctly.
+    result = driver.execute_script(f"return core.elementClick({ref});")
+    if result is not True:
+        logging.warning("Clicking %s failed: %s", ref, result)
+
+
+def execute_type(text, driver):
+    """Send keystrokes to the focused element."""
+    chain = ActionChains(driver)
+    chain.send_keys(text)
+    chain.perform()
+
+
+def execute_focus_and_type(ref, text, driver):
+    """Click the specified DOM element and then send keystrokes."""
+    execute_element_click(ref, driver)
+    execute_type(text, driver)
+
+
+def execute_action(action, driver):
+    """Execute the action on the ChromeDriver."""
+    action_type = action["action_type"]
+    if action_type == ActionTypes.NONE:
+        pass
+    elif action_type == ActionTypes.COORD_CLICK:
+        left = float(action["coords"][0])
+        top = float(action["coords"][1])
+        execute_coord_click(left, top, driver)
+    elif action_type == ActionTypes.ELEMENT_CLICK:
+        ref = int(action["ref"])
+        execute_element_click(ref, driver)
+    elif action_type == ActionTypes.TYPE:
+        text = action["text"]
+        execute_type(text, driver)
+    elif action_type == ActionTypes.FOCUS_AND_TYPE:
+        ref = int(action["ref"])
+        text = action["text"]
+        execute_focus_and_type(ref, text, driver)
+    else:
+        raise ValueError(f"Unknown action type: {action_type}")
diff --git a/miniwob/state.py b/miniwob/dom.py
similarity index 80%
rename from miniwob/state.py
rename to miniwob/dom.py
index ca0f0950..f70edccb 100644
--- a/miniwob/state.py
+++ b/miniwob/dom.py
@@ -1,142 +1,16 @@
+"""Encapsulation of the DOM element tree."""
 import re
 
-from miniwob.utils import Phrase
-
-
-class MiniWoBState:
-    """MiniWoB state.
-
-    Warning: The return types might be changed in the future!!!
-    """
-
-    # Task dimensions
-    HEIGHT = ROWS = 210
-    WIDTH = COLS = 160
-    PROMPT_HEIGHT = PROMPT_ROWS = 50
-
-    def __init__(self, utterance, fields, dom_info):
-        """Initialize a MiniWoBState.
-
-        Args:
-            utterance (unicode)
-            fields (Fields)
-            dom_info (dict)
-        """
-        ################
-        # Parse utterance
-        assert isinstance(utterance, str)
-        self._phrase = Phrase(utterance)
-        self._fields = fields
-        ################
-        # Store DOM
-        self._dom_elements = []
-        self._root_dom = DOMElement(
-            dom_info, parent=None, dom_elements=self._dom_elements
-        )
-        ################
-        # Screenshot (None by default)
-        self._screenshot = None
-
-    @property
-    def utterance(self):
-        """Task utterance.
-
-        Returns:
-            unicode
-        """
-        return self._phrase.text
-
-    @property
-    def phrase(self):
-        """The Phrase object of the utterance.
-
-        Returns:
-            Phrase
-        """
-        return self._phrase
-
-    @property
-    def tokens(self):
-        """Tokens.
-
-        Returns
-            list[unicode]
-        """
-        return list(self._phrase.tokens)
-
-    def detokenize(self, start, end):
-        """Return the substring of the original string that corresponds
-        to tokens[start:end].
-
-        Args:
-            start (int)
-            end (int)
-        Returns:
-            unicode
-        """
-        return self._phrase.detokenize(start, end)
-
-    @property
-    def fields(self):
-        """Key-value fields extracted from the utterance.
-
-        Returns:
-            Fields
-        """
-        return self._fields
-
-    @property
-    def dom(self):
-        """The root DOM structure.
-
-        Returns:
-            DOMElement
-        """
-        return self._root_dom
-
-    @property
-    def dom_elements(self):
-        """List of all DOM elements, flattened.
-
-        Returns:
-            list[DOMElement]
-        """
-        return self._dom_elements
-
-    def __str__(self):
-        return f"MiniWoBState(utterance: {repr(self.utterance)})"
-
-    __repr__ = __str__
-
-    def set_screenshot(self, pil_image):
-        """Add screenshot to the state.
-
-        Args:
-            pil_image (PIL Image)
-        """
-        self._screenshot = pil_image
-
-    @property
-    def screenshot(self):
-        """Return screenshot, or None if not exist.
-
-        Returns:
-            PIL Image or None
-        """
-        return self._screenshot
-
 
 class DOMElement:
     """Encapsulate the DOM element."""
 
-    def __init__(self, raw_dom, parent=None, dom_elements=None):
+    def __init__(self, raw_dom, parent=None):
         """Create a new DOMElement based on the data from getDOMInfo in JavaScript.
 
         Args:
             raw_dom (dict): A dict with values from getDOMInfo in JavaScript.
             parent (DOMElement|None): the parent DOMElement, or None
-            dom_elements (list|None): If specified, append this DOMElement
-                object to the list
         """
         self._parent = parent
         self._tag = raw_dom["tag"].lower()
@@ -167,17 +41,12 @@ def __init__(self, raw_dom, parent=None, dom_elements=None):
         # Recurse on the children
         self._children = []
         for raw_child in raw_dom["children"]:
-            self._children.append(
-                DOMElement(raw_child, parent=self, dom_elements=dom_elements)
-            )
+            self._children.append(DOMElement(raw_child, parent=self))
         # Fix a bug where sometimes children are created even though all children are <t>
         # (which will incorrectly make this element a non-leaf and thus unclickable)
         if self._children and all(child.tag == "t" for child in self._children):
             self._text = " ".join(child.text for child in self._children)
             self._children = []
-        # Add to the collection
-        if dom_elements is not None:
-            dom_elements.append(self)
 
     def __eq__(self, other):
         if not isinstance(other, DOMElement):
@@ -313,6 +182,14 @@ def children(self):
         """Return the list of children (list[DOMElement])."""
         return self._children
 
+    @property
+    def subtree_elements(self):
+        """Return the list of elements in the subtree, including self (list[DOMElement])."""
+        elements = [self]
+        for child in self.children:
+            elements += child.subtree_elements
+        return elements
+
     @property
     def parent(self):
         """Return the parent (DOMElement)."""
diff --git a/miniwob/environment.py b/miniwob/environment.py
index 17eaca1d..b0a4b7f6 100644
--- a/miniwob/environment.py
+++ b/miniwob/environment.py
@@ -1,46 +1,45 @@
+"""MiniWoB environment."""
 import logging
-import os
-import sys
 
+import gymnasium as gym
+
+from miniwob.action import get_action_space
 from miniwob.instance import MiniWoBInstance
+from miniwob.observation import get_observation_space
 
 
-class MiniWoBEnvironment:
+class MiniWoBEnvironment(gym.Env):
     """MiniWoB environment."""
 
-    def __init__(self, subdomain):
-        """Creates a new MiniWoBEnvironment with no instances.
-        Must call configure() to set up instances.
+    # render_mode = None: Headless Chrome (default)
+    # render_mode = "human": Show the Chrome screen
+    metadata = {"render_modes": ["human"]}
+    reward_range = (-1, 1)
+
+    def __init__(
+        self,
+        subdomain,
+        render_mode=None,
+        base_url=None,
+        threading=True,
+        reward_processor=None,
+        wait_ms=0.0,
+        block_on_reset=True,
+        refresh_freq=0,
+        data_mode="train",
+    ):
+        """Creates a new MiniWoBEnvironment.
 
         Args:
             subdomain (str): MiniWoB task name (e.g., "click-test")
-        """
-        self.subdomain = subdomain
-        self.instances = []
-        self.died = False
-
-    @property
-    def num_instances(self):
-        return len(self.instances)
-
-    def configure(self, num_instances=1, seeds=None, **kwargs):
-        """Creates the required number of MiniWoBInstance.
-
-        Args:
-            num_instances (int): Number of instances
-
-        kwargs are passed into the constructor of MiniWoBInstance:
-            headless (bool): Whether to render GUI
+            render_mode (str): Render mode
             base_url (str): Base URL, which is usually one of the following
                 - http://localhost:8000/     (served by http-serve)
                 - file:///path/to/miniwob-plusplus/html/
-            cache_state (bool): Whether to cache and return the initial
-                state; only make sense if the task interface never changes
+                If None, infers the file:// path from this module's location.
             threading (bool): Whether to run the instances in separate threads
             reward_processor (callable; optional): A function that takes
                 the metadata and return a reward (see miniwob.reward)
-            seeds (list[object]): Random seeds to set for each instance;
-                len(seeds) must be equal to num_instances.
             wait_ms (float): Pause the instance after each action for this
                 amount of time (in milliseconds).
             block_on_reset (bool): On reset, block until the page loads.
@@ -48,145 +47,130 @@ def configure(self, num_instances=1, seeds=None, **kwargs):
                 refresh the page at the beginning of the next episode.
                 Takes time but cleans up any lingering states and memory leaks.
                 *** Must specify `seeds` at each reset call.
-            initial_mode (str): Initial data mode (e.g., "train", "test")
+            data_mode (str): Data mode (e.g., "train", "test"). Used in some
+                subdomains.
         """
-        assert seeds is not None, "seeds must be specified"
-        assert len(seeds) == num_instances, "len(seeds) must be equal to num_instances"
-        self.configure_kwargs = kwargs
-        for instance in self.instances:
-            instance.close()
-        self.instances = []
-        for index in range(num_instances):
-            logging.info("Starting WebDriver Instance %d", index)
-            instance = MiniWoBInstance(index, self.subdomain, seeds[index], **kwargs)
-            instance.start()
-            self.instances.append(instance)
-        for instance in self.instances:
-            instance.wait()
-
-    def reset(self, seeds=None, mode=None, record_screenshots=False):
-        """Forces stop and start all instances.
+        if render_mode and render_mode not in self.metadata["render_modes"]:
+            raise ValueError(f"Invalid render mode: {render_mode}")
+        self.render_mode = render_mode
+        self.instance_kwargs = {
+            "subdomain": subdomain,
+            "headless": (render_mode is None),
+            "base_url": base_url,
+            "threading": threading,
+            "reward_processor": reward_processor,
+            "wait_ms": wait_ms,
+            "block_on_reset": block_on_reset,
+            "refresh_freq": refresh_freq,
+            "data_mode": data_mode,
+        }
+        self.instance = None
+        self._hard_reset_instance()
+        self.action_space = get_action_space(
+            screen_width=self.instance.task_width,
+            screen_height=self.instance.task_height,
+        )
+        self.observation_space = get_observation_space(
+            screen_width=self.instance.task_width,
+            screen_height=self.instance.task_height,
+        )
+
+    def _hard_reset_instance(self):
+        """Closes the current MiniWoBInstance (if exists) and starts a new one."""
+        if self.instance:
+            self.instance.close()
+        logging.info("Starting WebDriver Instance")
+        self.instance = MiniWoBInstance(index=0, **self.instance_kwargs)
+        self.instance.start()
+        self.instance.wait()
+
+    def reset(self, seed=None, options=None):
+        """Reset the instance.
 
         Args:
-            seeds (list[object]): Random seeds to set for each instance;
-                If specified, len(seeds) must be equal to the number of instances.
-                A None entry in the list = do not set a new seed.
-            mode (str): If specified, set the data mode to this value before
-                starting new episodes.
-            record_screenshots (bool): Whether to record screenshots of the states.
+            seed (optional int): Random seed.
+            options (optional dict): An option dict with the following allowed keys:
+                - data_mode (str): set the data mode to this value.
+                - record_screenshots (bool): Whether to record screenshots.
         Returns:
-            states (list[MiniWoBState])
+            observation (MiniWoBState)
+            info (dict)
         """
-        # If an instance died, call configure
-        if any(instance.died for instance in self.instances):
-            logging.warning("An instance died. Reset instance ...")
-            self.configure(len(self.instances), seeds, **self.configure_kwargs)
-        # Parse arguments
-        if seeds is None:
-            seeds = [None] * len(self.instances)
-        else:
-            assert isinstance(seeds, (list, tuple)) and len(seeds) == len(
-                self.instances
-            )
-        if mode is not None:
-            self.set_mode(mode)
-        self.set_record_screenshots(record_screenshots)
-        # The ith entry in `states` will be set by the ith instance
-        states = [None] * len(self.instances)
-        for i, instance in enumerate(self.instances):
-            instance.call(instance.reset, states, seeds[i])
-        for instance in self.instances:
-            instance.wait()
-        self.died = any(instance.died for instance in self.instances)
-        return states
-
-    def step(self, actions):
-        """Applies an action on each instance and returns the results.
+        # The seed in Env is actually not used
+        super().reset(seed=seed)
+        # Hard reset the instances if needed
+        if not self.instance or self.instance.died:
+            logging.warning("Hard-resetting the instance ...")
+            self._hard_reset_instance()
+        # Process the options
+        options = options or {}
+        if "data_mode" in options:
+            self.set_data_mode(options["data_mode"])
+        if "record_screenshots" in options:
+            self.set_record_screenshots(options["record_screenshots"])
+        # We pass lists for the instance to modify in-place.
+        obs = [None]
+        infos = [None]
+        self.instance.call(self.instance.reset, obs, infos, seed)
+        self.instance.wait()
+        return obs[0], infos[0]
+
+    def step(self, action):
+        """Applies an action on the instance and returns the result.
 
         Args:
-            actions (list[MiniWoBAction or None])
+            action (MiniWoBAction)
 
         Returns:
-            tuple (states, rewards, dones, info)
-            states (list[MiniWoBState])
-            rewards (list[float])
-            dones (list[bool])
+            observation (MiniWoBState)
+            reward (float)
+            terminated (bool)
+            truncated (bool)
             info (dict): additional debug information.
-                Global debug information is directly in the root level
-                Local information for instance i is in info['n'][i]
         """
-        assert len(actions) == len(
-            self.instances
-        ), "len(action) is {} but there are {} instances".format(
-            len(actions), len(self.instances)
-        )
-        # Initialize with reasonable values
-        states = [None] * len(self.instances)
-        rewards = [-1.0] * len(self.instances)
-        dones = [True] * len(self.instances)
-        info = {"n": [{} for _ in self.instances]}
-        # Have the instances replace the values
-        for i, instance in enumerate(self.instances):
-            instance.call(instance.step, actions[i], states, rewards, dones, info["n"])
-        for instance in self.instances:
-            instance.wait()
-        self.died = any(instance.died for instance in self.instances)
-        return states, rewards, dones, info
-
-    def set_mode(self, mode):
-        """Set the data mode ("train", "test", etc.) of all instances.
+        # We pass lists for the instance to modify in-place.
+        obs = [None]
+        rewards = [-1.0]
+        dones = [True]
+        truncs = [False]
+        infos = [None]
+        self.instance.call(self.instance.step, action, obs, rewards, dones, infos)
+        self.instance.wait()
+        return obs[0], rewards[0], dones[0], truncs[0], infos[0]
+
+    def render(self):
+        # The currently supported render modes do not require computing the render.
+        return None
+
+    def set_data_mode(self, mode):
+        """Set the data mode ("train", "test", etc.) of the instance.
         Will have effect starting from the next episode.
 
         Args:
             mode (str)
         """
-        for instance in self.instances:
-            instance.mode = mode
+        self.instance.mode = mode
 
     def set_record_screenshots(self, record_screenshots):
-        """Adjust whether the record the screenshots of the states.
+        """Adjust whether the record the screenshots.
 
         Args:
             record_screenshots (bool)
         """
-        for instance in self.instances:
-            instance.record_screenshots = record_screenshots
+        self.instance.record_screenshots = record_screenshots
 
     def visualize_attention(self, attentions):
         """Sends the attention weights to be visualized.
 
         Args:
-            attentions (list[*]): attention weight for each instance.
-                Each list element is one of:
+            attentions: attention weights, which is one of:
                 - None: Do not do anything
                 - np.array or 2d list of shape (num_grid_rows, num_grid_cols)
                 - np.array or 2d list of shape (0, 0): Clear the visualization
         """
-        for i, instance in enumerate(self.instances):
-            instance.call(instance.visualize_attention, attentions[i])
-        for instance in self.instances:
-            instance.wait()
+        self.instance.call(self.instance.visualize_attention, attentions)
+        self.instance.wait()
 
     def close(self):
-        for instance in self.instances:
-            instance.call(instance.close)
-        for instance in self.instances:
-            instance.wait()
-
-
-def test_environment():
-    try:
-        task_name = sys.argv[1]
-    except IndexError:
-        print(f"Usage: python {sys.argv[0]} TASK_NAME")
-        exit(1)
-    env = MiniWoBEnvironment(task_name)
-    base_url = os.environ.get("MINIWOB_BASE_URL")
-    env.configure(num_instances=1, seeds=[0], base_url=base_url)
-    states = env.reset()
-    print(states[0].dom.visualize())
-    env.close()
-
-
-if __name__ == "__main__":
-    test_environment()
+        self.instance.call(self.instance.close)
+        self.instance.wait()
diff --git a/miniwob/fields.py b/miniwob/fields.py
index 4990b141..4237dbc2 100644
--- a/miniwob/fields.py
+++ b/miniwob/fields.py
@@ -1,8 +1,7 @@
+"""Task-specific key-value pairs extracted from the task instructions."""
 import collections
 import json
-import os
 import re
-import sys
 
 # Mapping from task_name to field extractor
 # Each extractor is a function that takes the utterance string and
@@ -1122,29 +1121,3 @@ def extract_flight_subtasks(utterance):
 FIELD_EXTRACTORS["flight.Delta"] = extract_flight_subtasks
 FIELD_EXTRACTORS["flight.JetBlue"] = extract_flight_subtasks
 FIELD_EXTRACTORS["flight.United"] = extract_flight_subtasks
-
-################################################################
-# Script for printing out the utterances for a particular task
-
-
-def extract_utterances():
-    try:
-        task_name = sys.argv[1]
-    except IndexError:
-        print(f"Usage: {sys.argv[0]} task_name", file=sys.stderr)
-        exit(1)
-    FIELD_EXTRACTORS[task_name] = lambda utt: Fields({})
-    from miniwob.environment import MiniWoBEnvironment
-
-    env = MiniWoBEnvironment(task_name)
-    base_url = os.environ.get("MINIWOB_BASE_URL")
-    env.configure(num_instances=4, seeds=list(range(4)), base_url=base_url)
-    for i in range(25):
-        states = env.reset()
-        for state in states:
-            print("UTT:\t{}".format(state.utterance.replace("\n", " ")))
-    env.close()
-
-
-if __name__ == "__main__":
-    extract_utterances()
diff --git a/miniwob/instance.py b/miniwob/instance.py
index ec324ae7..266a2d74 100644
--- a/miniwob/instance.py
+++ b/miniwob/instance.py
@@ -1,3 +1,4 @@
+"""Low-level interface with ChromDriver via Selenium."""
 import json
 import logging
 import pathlib
@@ -14,19 +15,23 @@
 from selenium.webdriver.support import expected_conditions as EC
 from selenium.webdriver.support.ui import WebDriverWait
 
+from miniwob.action import execute_action
+from miniwob.dom import DOMElement
 from miniwob.fields import Fields, get_field_extractor
+from miniwob.observation import (
+    create_empty_observation,
+    create_empty_screenshot,
+    create_observation,
+)
 from miniwob.reward import get_original_reward
-from miniwob.screenshot import get_screenshot
-from miniwob.state import MiniWoBState
+from miniwob.screenshot import get_screenshot, pil_to_numpy_array
 
 HTML_DIR = pathlib.Path(__file__).parent.parent / "html"
 DEFAULT_BASE_URL = f"file://{HTML_DIR}/"
 
 
 class MiniWoBInstance(Thread):
-    """Interface between Python and Chrome driver via Selenium.
-    Manages a single instance.
-    """
+    """Interface between Python and ChromeDriver via Selenium."""
 
     # Added some space for title bar
     WINDOW_WIDTH = 500
@@ -42,28 +47,26 @@ class MiniWoBInstance(Thread):
     def __init__(
         self,
         index,
-        subdomain,
-        seed,
+        subdomain=None,
         headless=False,
         base_url=None,
-        cache_state=False,
         threading=True,
         reward_processor=None,
         wait_ms=0.0,
         block_on_reset=True,
         refresh_freq=0,
-        initial_mode="train",
+        data_mode="train",
     ):
         """Starts a new Selenium WebDriver session.
 
         Args:
             index (int): Instance index
             subdomain (str): MiniWoB task name (e.g., "click-test")
-            seed (object): Random seed
             headless (bool): Whether to render GUI
-            base_url (str): Base URL (default to localhost at port 8000)
-            cache_state (bool): Whether to cache and return the initial
-                state; only make sense if the task interface never changes
+            base_url (str): Base URL, which is usually one of the following
+                - http://localhost:8000/     (served by http-serve)
+                - file:///path/to/miniwob-plusplus/html/
+                If None, infers the file:// path from this module's location.
             threading (bool): Whether to run this instance as a Thread
             reward_processor (callable; optional): A function that takes
                 the metadata and return a reward (see miniwob.reward)
@@ -74,14 +77,13 @@ def __init__(
                 refresh the page at the beginning of the next episode.
                 Takes time but cleans up any lingering states and memory leaks.
                 *** Must specify `seeds` at each reset call.
-            initial_mode (str): Initial data mode (e.g., "train", "test")
+            data_mode (str): Data mode (e.g., "train", "test")
         """
         super().__init__()
         # Overrides Thread.daemon: Kill this thread when the parent is killed
         self.daemon = True
         self.died = False
         self.index = index
-        self.init_seed = repr(seed)
         self.headless = headless
         base_url = base_url or DEFAULT_BASE_URL
         if subdomain.startswith("flight."):
@@ -103,14 +105,13 @@ def __init__(
             self.task_width = self.TASK_WIDTH
             self.task_height = self.TASK_HEIGHT
         self.field_extractor = get_field_extractor(subdomain)
-        self.cache_state = cache_state
         self.threading = threading
         self.reward_processor = reward_processor
         self.wait_ms = wait_ms
         self.block_on_reset = block_on_reset
         self.refresh_freq = refresh_freq
         self.num_episodes = 0
-        self.mode = initial_mode
+        self.mode = data_mode
         self.record_screenshots = False
         if reward_processor is None:
             # Use the original reward
@@ -185,8 +186,6 @@ def create_driver(self):
         except TimeoutException as e:
             logging.error("Page did not load properly. Wrong MINIWOB_BASE_URL?")
             raise e
-        # Seed the seed
-        self.driver.execute_script(f"Math.seedrandom({self.init_seed});")
 
     def close(self):
         """Tear down the WebDriver."""
@@ -199,13 +198,14 @@ def close(self):
             traceback.print_exc()
         self.died = True
 
-    def reset(self, states, seed):
+    def reset(self, obs, infos, seed):
         """Forces stop and start this instance.
-        Also sets states[i] to be the initial state
+        Also sets obs[i] to be the initial observation
         (where i = self.index).
 
         Args:
-            states (list)
+            obs (list)
+            infos (list)
             seed (object): Seed to set for the next episode
         """
         if self.refresh_freq:
@@ -215,34 +215,36 @@ def reset(self, states, seed):
         i = self.index
         self.force_stop()
         self.begin_task(seed=seed)
-        states[i] = self.get_state()
-        if self.cache_state:
-            self.initial_state = states[i]
+        obs[i], extra_metadata = self.get_observation()
+        metadata = self.get_metadata()
+        metadata.update(extra_metadata)
+        infos[i] = metadata
 
-    def step(self, action, states, rewards, dones, info_n):
+    def step(self, action, obs, rewards, dones, infos):
         """Applies an action on this instance.
-        Also sets states[i], rewards[i], dones[i], and info['n'][i]
+        Also sets obs[i], rewards[i], dones[i], and infos[i]
         (where i = self.index).
 
         Args:
             action (MiniWoBAction)
-            states (list)
+            obs (list)
             rewards (list)
             dones (list)
-            info_n (list)
+            infos (list)
         """
         i = self.index
         self.perform(action)
         metadata = self.get_metadata()
         rewards[i] = self.reward_processor(metadata)
         dones[i] = metadata["done"]
-        if not metadata["done"]:
-            if not self.cache_state:
-                states[i] = self.get_state()
-            else:
-                states[i] = self.initial_state
+        if metadata["done"]:
+            obs[i] = self.get_empty_observation()
+            extra_metadata = {}
+        else:
+            obs[i], extra_metadata = self.get_observation()
         metadata["elapsed"] = max(0.0, time.time() - self.start_time)
-        info_n[i] = metadata
+        metadata.update(extra_metadata)
+        infos[i] = metadata
 
     ################################
     # Primitive actions
@@ -297,15 +299,24 @@ def perform(self, action):
                     self.index,
                 )
             else:
-                action(self.driver)
+                execute_action(action, self.driver)
         if self.wait_ms:
             time.sleep(self.wait_ms / 1000.0)
 
-    def get_state(self):
-        """Get the current state.
+    def get_empty_observation(self):
+        """Get an empty observation for a terminated session."""
+        return create_empty_observation(self.task_width, self.task_height)
+
+    def get_observation(self):
+        """Get the current observation.
 
         Returns:
-            MiniWoBState
+            a tuple (observation, extra_metadata).
+            observation: Observation object from the observation space.
+            extra_metadata: A dict containing the following extra information:
+                - root_dom: DOMElement object for the root DOM element.
+                - fields: Fields object storing task-specific key-value pairs
+                    extracted from the instruction.
         """
         # Get the utterance
         response = self.driver.execute_script("return core.getUtterance();")
@@ -317,12 +328,15 @@ def get_state(self):
             fields = self.field_extractor(utterance)
         # Get the DOM
         dom_info = self.driver.execute_script("return core.getDOMInfo();")
-        state = MiniWoBState(utterance, fields, dom_info)
+        root_dom = DOMElement(dom_info)
         # Get screenshot if requested
         if self.record_screenshots:
             img = get_screenshot(self.driver, self.task_width, self.task_height)
-            state.set_screenshot(img)
-        return state
+            img = pil_to_numpy_array(img)
+        else:
+            img = create_empty_screenshot(self.task_width, self.task_height)
+        observation = create_observation(utterance, root_dom, img)
+        return observation, {"root_dom": root_dom, "fields": fields}
 
     def get_metadata(self):
         """Get other metadata.
diff --git a/miniwob/observation.py b/miniwob/observation.py
new file mode 100644
index 00000000..c2b43311
--- /dev/null
+++ b/miniwob/observation.py
@@ -0,0 +1,149 @@
+"""MiniWoB observation space."""
+import numpy as np
+from gymnasium import spaces
+
+UTTERANCE_MAX_LENGTH = 256
+ATTRIBUTE_MAX_LENGTH = 64
+TEXT_MAX_LENGTH = 256
+ASCII_CHARSET = frozenset(chr(x) for x in range(32, 128))
+MIN_REF = -1000000
+MAX_REF = 1000000
+
+
+def get_observation_space(screen_width, screen_height):
+    """Return the space of observations."""
+    utterance_space = spaces.Text(
+        max_length=UTTERANCE_MAX_LENGTH,
+        charset=ASCII_CHARSET,
+    )
+    element_space = spaces.Dict(
+        {
+            # Non-zero integer ID:
+            # `ref` for normal elements start from 1, while `ref` for text
+            # pseudo-elements counts down from -1.
+            "ref": spaces.Discrete(MAX_REF - MIN_REF, start=MIN_REF),
+            # `ref` ID of the parent (0 = no parent, for root element).
+            "parent": spaces.Discrete(MAX_REF),
+            # Position (left, top)
+            "pos": spaces.Box(
+                np.array([float("-inf"), float("-inf")]),
+                np.array([float("inf"), float("inf")]),
+            ),
+            # Size (width, height)
+            "size": spaces.Box(
+                np.array([0.0, 0.0]), np.array([float("inf"), float("inf")])
+            ),
+            # Tag:
+            # For normal elements, this is the uppercased tag name (e.g., "DIV").
+            # For <input> elements, the input type is appended (e.g., "INPUT_text").
+            # Each non-empty text node becomes a pseudo-elements with tag "t".
+            "tag": spaces.Text(max_length=ATTRIBUTE_MAX_LENGTH, charset=ASCII_CHARSET),
+            # Text content of leaf nodes.
+            "text": spaces.Text(
+                min_length=0, max_length=TEXT_MAX_LENGTH, charset=ASCII_CHARSET
+            ),
+            # Value of <input> elements.
+            "value": spaces.Text(
+                min_length=0, max_length=TEXT_MAX_LENGTH, charset=ASCII_CHARSET
+            ),
+            # HTML id attribute
+            "id": spaces.Text(
+                min_length=0, max_length=ATTRIBUTE_MAX_LENGTH, charset=ASCII_CHARSET
+            ),
+            # HTML class attribute (multiple classes are separated by spaces)
+            "classes": spaces.Text(
+                min_length=0, max_length=ATTRIBUTE_MAX_LENGTH, charset=ASCII_CHARSET
+            ),
+            # Colors (RGBA)
+            "bg_color": spaces.Box(
+                np.array([0.0] * 4, dtype=np.float32),
+                np.array([255.0] * 4, dtype=np.float32),
+            ),
+            "fg_color": spaces.Box(
+                np.array([0.0] * 4, dtype=np.float32),
+                np.array([255.0] * 4, dtype=np.float32),
+            ),
+            # Flags:
+            # focused: whether the element is being focused on
+            # tampered: whether the element has been tampered (clicked, focused, typed, etc.)
+            # targeted: whether the element is an event target (for recorded demonstrations)
+            # is_leaf: whether the element is a leaf
+            "flags": spaces.MultiBinary(n=4),
+        }
+    )
+    screenshot_space = spaces.Box(
+        # Each position stores the RGB values. Note the swapped axes (height first).
+        np.zeros((screen_height, screen_width, 3), dtype=np.uint8),
+        np.ones((screen_height, screen_width, 3), dtype=np.uint8) * 255.0,
+        dtype=np.uint8,
+    )
+    observation_space = spaces.Dict(
+        {
+            "utterance": utterance_space,
+            "dom_elements": spaces.Sequence(element_space),
+            "screenshot": screenshot_space,
+        }
+    )
+    return observation_space
+
+
+def serialize_dom_element(element):
+    """Serialize the given DOMElement to fit the element space."""
+    serialized = {
+        "ref": element.ref,
+        "parent": element.parent.ref if element.parent else 0,
+        "pos": np.array([element.left, element.top], dtype=np.float32),
+        "size": np.array([element.width, element.height], dtype=np.float32),
+        "tag": element.tag[:ATTRIBUTE_MAX_LENGTH],
+        "text": (element.text or "")[:TEXT_MAX_LENGTH],
+        "value": str(element.value or "")[:TEXT_MAX_LENGTH],
+        "id": element.id,
+        "classes": element.classes,
+        "bg_color": np.array(element.bg_color, dtype=np.float32),
+        "fg_color": np.array(element.fg_color, dtype=np.float32),
+        "flags": np.array(
+            [
+                element.focused,
+                element.tampered,
+                element.targeted,
+                element.is_leaf,
+            ],
+            dtype=np.int8,
+        ),
+    }
+    return serialized
+
+
+def create_empty_screenshot(screen_width, screen_height):
+    """Returns an all-black screenshot."""
+    return np.zeros((screen_height, screen_width, 3), dtype=np.uint8)
+
+
+def create_empty_observation(screen_width, screen_height):
+    """Returns an empty observation for a terminated session."""
+    observation = {
+        "utterance": "",
+        "dom_elements": [],
+        "screenshot": create_empty_screenshot(screen_width, screen_height),
+    }
+    return observation
+
+
+def create_observation(utterance, root_dom, screenshot):
+    """Returns an observation that fits in the observation space.
+
+    Args:
+        utterance: Instruction text extracted from the task.
+        root_dom: DOMElement object for the root element.
+        screenshot: Screenshot as an RGB array.
+    Returns:
+        the observation object
+    """
+    dom_elements = root_dom.subtree_elements
+    serialized_elements = [serialize_dom_element(element) for element in dom_elements]
+    observation = {
+        "utterance": utterance[:UTTERANCE_MAX_LENGTH],
+        "dom_elements": serialized_elements,
+        "screenshot": screenshot,
+    }
+    return observation
diff --git a/miniwob/registration.py b/miniwob/registration.py
new file mode 100644
index 00000000..cf83c3c8
--- /dev/null
+++ b/miniwob/registration.py
@@ -0,0 +1,132 @@
+"""Gymnasium environment registration."""
+from gymnasium.envs.registration import register
+
+MINIWOB_TASKS = [
+    "bisect-angle",
+    "book-flight",
+    "book-flight-nodelay",
+    "chase-circle",
+    "choose-date",
+    "choose-date-easy",
+    "choose-date-medium",
+    "choose-date-nodelay",
+    "choose-list",
+    "circle-center",
+    "click-button",
+    "click-button-sequence",
+    "click-checkboxes",
+    "click-checkboxes-large",
+    "click-checkboxes-soft",
+    "click-checkboxes-transfer",
+    "click-collapsible",
+    "click-collapsible-2",
+    "click-collapsible-2-nodelay",
+    "click-collapsible-nodelay",
+    "click-color",
+    "click-dialog",
+    "click-dialog-2",
+    "click-link",
+    "click-menu",
+    "click-menu-2",
+    "click-option",
+    "click-pie",
+    "click-pie-nodelay",
+    "click-scroll-list",
+    "click-shades",
+    "click-shape",
+    "click-tab",
+    "click-tab-2",
+    "click-tab-2-easy",
+    "click-tab-2-hard",
+    "click-tab-2-medium",
+    "click-test",
+    "click-test-2",
+    "click-test-transfer",
+    "click-widget",
+    "copy-paste",
+    "copy-paste-2",
+    "count-shape",
+    "count-sides",
+    "drag-box",
+    "drag-cube",
+    "drag-item",
+    "drag-items",
+    "drag-items-grid",
+    "drag-shapes",
+    "drag-sort-numbers",
+    "email-inbox",
+    "email-inbox-delete",
+    "email-inbox-forward",
+    "email-inbox-forward-nl",
+    "email-inbox-forward-nl-turk",
+    "email-inbox-important",
+    "email-inbox-nl-turk",
+    "email-inbox-noscroll",
+    "email-inbox-reply",
+    "email-inbox-star-reply",
+    "enter-date",
+    "enter-password",
+    "enter-text",
+    "enter-text-2",
+    "enter-text-dynamic",
+    "enter-time",
+    "find-midpoint",
+    "find-word",
+    "focus-text",
+    "focus-text-2",
+    "grid-coordinate",
+    "guess-number",
+    "highlight-text",
+    "highlight-text-2",
+    "identify-shape",
+    "login-user",
+    "login-user-popup",
+    "moving-items",
+    "multi-layouts",
+    "multi-orderings",
+    "navigate-tree",
+    "number-checkboxes",
+    "read-table",
+    "read-table-2",
+    "resize-textarea",
+    "right-angle",
+    "scroll-text",
+    "scroll-text-2",
+    "search-engine",
+    "simon-says",
+    "simple-algebra",
+    "simple-arithmetic",
+    "social-media",
+    "social-media-all",
+    "social-media-some",
+    "terminal",
+    "text-editor",
+    "text-transform",
+    "tic-tac-toe",
+    "unicode-test",
+    "use-autocomplete",
+    "use-autocomplete-nodelay",
+    "use-colorwheel",
+    "use-colorwheel-2",
+    "use-slider",
+    "use-slider-2",
+    "use-spinner",
+    "visual-addition",
+]
+FLIGHTWOB_TASKS = ["AA", "Alaska", "Alaska-auto", "Alaska-auto-medium"]
+
+
+def register_miniwob_envs():
+    """Register MiniWoB and FlightWoB environments."""
+    for name in MINIWOB_TASKS:
+        register(
+            id=f"miniwob/{name}-v1",
+            entry_point="miniwob.environment:MiniWoBEnvironment",
+            kwargs={"subdomain": name},
+        )
+    for name in FLIGHTWOB_TASKS:
+        register(
+            id=f"miniwob/flight.{name}-v1",
+            entry_point="miniwob.environment:MiniWoBEnvironment",
+            kwargs={"subdomain": f"flight.{name}"},
+        )
diff --git a/miniwob/screenshot.py b/miniwob/screenshot.py
index fd587fe9..41d5330a 100644
--- a/miniwob/screenshot.py
+++ b/miniwob/screenshot.py
@@ -1,4 +1,4 @@
-# Dealing with screenshots
+"""Screenshot utilities."""
 import json
 import sys
 from io import StringIO
@@ -32,7 +32,7 @@ def pil_to_numpy_array(pil_image):
         numpy array of shape (height, width, 3)
         where 3 is the number of channels (RGB).
     """
-    return np.array(pil_image).astype(np.float32)
+    return np.array(pil_image).astype(np.uint8)
 
 
 def create_gif(path_prefix):
diff --git a/miniwob/utils.py b/miniwob/utils.py
index ff83cd11..63b7c04b 100644
--- a/miniwob/utils.py
+++ b/miniwob/utils.py
@@ -1,3 +1,4 @@
+"""Generic utilities."""
 import re
 
 


Test patch:
diff --git a/miniwob/tests/test_action.py b/miniwob/tests/test_action.py
index 4cc92e0a..467c9020 100644
--- a/miniwob/tests/test_action.py
+++ b/miniwob/tests/test_action.py
@@ -1,10 +1,10 @@
 import pytest
 
 from miniwob.action import (
-    MiniWoBCoordClick,
-    MiniWoBElementClick,
-    MiniWoBFocusAndType,
-    MiniWoBType,
+    create_coord_click_action,
+    create_element_click_action,
+    create_focus_and_type_action,
+    create_type_action,
 )
 from miniwob.environment import MiniWoBEnvironment
 
@@ -21,68 +21,43 @@ class RepeatedTester:
 
     @pytest.fixture
     def env(self):
-        env = MiniWoBEnvironment(self.TASK_NAME)
-        if self.FRAGILE is True:
-            env.configure(headless=True, num_instances=1, seeds=[1], wait_ms=300)
-        elif self.FRAGILE == "instance":
-            env.configure(headless=True, num_instances=1, seeds=[1])
-        elif self.FRAGILE == "delay":
-            env.configure(
-                headless=True, num_instances=3, seeds=list(range(3)), wait_ms=1000
-            )
+        if self.FRAGILE:
+            env = MiniWoBEnvironment(subdomain=self.TASK_NAME, wait_ms=300)
         else:
-            env.configure(headless=True, num_instances=3, seeds=list(range(3)))
+            env = MiniWoBEnvironment(subdomain=self.TASK_NAME)
         yield env
         env.close()
 
     def test_run(self, env):
         for i in range(self.N):
             print(f"Iteration {i + 1} / {self.N}")
-            states = env.reset()
-            for j, state in enumerate(states):
-                print(f"Fields {j}: {state.fields}")
-            for s in range(self.MAX_STEPS):
-                print(f"Step {s + 1} / {self.MAX_STEPS}")
-                actions = [self.get_action(x, s) for x in states]
-                states, rewards, dones, info = env.step(actions)
-                if all(x for x in dones):
+            obs, info = env.reset()
+            reward = -1
+            for step in range(self.MAX_STEPS):
+                action = self.get_action(obs, info, step)
+                obs, reward, terminated, _, _ = env.step(action)
+                assert reward >= 0
+                if terminated:
                     break
-                assert all(x >= 0 for x in rewards)
             else:
                 assert False, f"Number of steps exceeded {self.MAX_STEPS}"
-            assert all(x > 0 for x in rewards)
+            assert reward >= 0
 
-    def get_action(self, state, step):
+    def get_action(self, obs, info, step):
         """Returns a MiniWoBAction that clicks the right thing."""
         raise NotImplementedError
 
-    def create_element_click_action(self, element):
-        action = MiniWoBElementClick(element, fail_hard=True)
-        print(f"Clicking with {action}")
-        return action
-
-    def click_button(self, state, text):
+    def click_button(self, obs, text):
         """Create an action that clicks on the button with the specified text."""
-        for element in state.dom_elements:
-            if element.tag == "button" and element.text == text:
-                return self.create_element_click_action(element)
+        for element in obs["dom_elements"]:
+            if element["tag"] == "button" and element["text"] == text:
+                return create_element_click_action(element["ref"])
         assert False, "Submit button not found"
 
     def create_coord_click_action(self, element):
-        action = MiniWoBCoordClick(
-            element.left + (element.width / 2), element.top + (element.height / 2)
-        )
-        print(f"Clicking with {action}")
-        return action
-
-    def create_type_action(self, text):
-        action = MiniWoBType(text)
-        print(f"Typing with {action}")
-        return action
-
-    def create_focus_and_type_action(self, element, text):
-        action = MiniWoBFocusAndType(element, text)
-        print(f"Focus and Type: {action}")
+        left, top = element["pos"].tolist()
+        width, height = element["size"].tolist()
+        action = create_coord_click_action(left + (width / 2), top + (height / 2))
         return action
 
 
@@ -93,10 +68,10 @@ def create_focus_and_type_action(self, element, text):
 class TestClickTest2(RepeatedTester):
     TASK_NAME = "click-test-2"
 
-    def get_action(self, state, step):
-        for element in state.dom_elements:
-            if element.tag == "button" and element.text == "ONE":
-                return self.create_element_click_action(element)
+    def get_action(self, obs, info, step):
+        for element in obs["dom_elements"]:
+            if element["tag"] == "button" and element["text"] == "ONE":
+                return create_element_click_action(element["ref"])
         # No button is found, which is weird
         assert False, 'Button "ONE" not found'
 
@@ -104,10 +79,10 @@ def get_action(self, state, step):
 class TestClickButton(RepeatedTester):
     TASK_NAME = "click-button"
 
-    def get_action(self, state, step):
-        target = state.fields["target"]
-        for element in state.dom_elements:
-            if element.tag == "button" and element.text == target:
+    def get_action(self, obs, info, step):
+        target = info["fields"]["target"]
+        for element in obs["dom_elements"]:
+            if element["tag"] == "button" and element["text"] == target:
                 return self.create_coord_click_action(element)
         # No button is found, which is weird
         assert False, f'Button "{target}" not found'
@@ -116,9 +91,9 @@ def get_action(self, state, step):
 class TestFocusText(RepeatedTester):
     TASK_NAME = "focus-text"
 
-    def get_action(self, state, step):
-        for element in state.dom_elements:
-            if element.tag == "input_text":
+    def get_action(self, obs, info, step):
+        for element in obs["dom_elements"]:
+            if element["tag"] == "input_text":
                 return self.create_coord_click_action(element)
         # No input is found, which is weird
         assert False, "Input box not found"
@@ -127,42 +102,43 @@ def get_action(self, state, step):
 class TestIdentifyShape(RepeatedTester):
     TASK_NAME = "identify-shape"
 
-    def get_action(self, state, step):
-        shape = self._identify_shape(state)
-        for element in state.dom_elements:
-            if element.tag == "button" and element.text == shape:
-                return self.create_element_click_action(element)
+    def get_action(self, obs, info, step):
+        shape = self._identify_shape(obs)
+        for element in obs["dom_elements"]:
+            if element["tag"] == "button" and element["text"] == shape:
+                return create_element_click_action(element["ref"])
         # No button is found, which is weird
         assert False, f'Button "{shape}" not found'
 
-    def _identify_shape(self, state):
-        for element in state.dom_elements:
-            if element.tag == "svg":
-                child = element.children[0]
-                print(child)
-                if child.tag == "circle":
+    def _identify_shape(self, obs):
+        for element in obs["dom_elements"]:
+            if element["tag"] == "svg":
+                child = [
+                    x for x in obs["dom_elements"] if x["parent"] == element["ref"]
+                ][0]
+                if child["tag"] == "circle":
                     return "Circle"
-                elif child.tag == "text":
-                    if child.text.isdigit():
+                elif child["tag"] == "text":
+                    if child["text"].isdigit():
                         return "Number"
                     else:
                         return "Letter"
-                elif child.tag == "rect":
+                elif child["tag"] == "rect":
                     return "Rectangle"
-                elif child.tag == "polygon":
+                elif child["tag"] == "polygon":
                     return "Triangle"
 
 
 class TestClickDialog2(RepeatedTester):
     TASK_NAME = "click-dialog-2"
 
-    def get_action(self, state, step):
-        target = state.fields["target"]
+    def get_action(self, obs, info, step):
+        target = info["fields"]["target"]
         if target == "x":
-            target = None
-        for element in state.dom_elements:
-            if element.tag == "button" and element.text == target:
-                return self.create_element_click_action(element)
+            target = ""
+        for element in obs["dom_elements"]:
+            if element["tag"] == "button" and element["text"] == target:
+                return create_element_click_action(element["ref"])
         # No button is found, which is weird
         assert False, f'Button "{target}" not found'
 
@@ -175,75 +151,75 @@ class TestEnterText(RepeatedTester):
     TASK_NAME = "enter-text"
     MAX_STEPS = 3
 
-    def get_action(self, state, step):
+    def get_action(self, obs, info, step):
         if step == 0:
             # Click on the textbox
-            for element in state.dom_elements:
-                if element.tag == "input_text":
-                    assert not element.focused
-                    return self.create_element_click_action(element)
+            for element in obs["dom_elements"]:
+                if element["tag"] == "input_text":
+                    assert not element["flags"][0].item()
+                    return create_element_click_action(element["ref"])
             assert False, "Input text not found"
         elif step == 1:
             # Assert that the input is focused
-            for element in state.dom_elements:
-                if element.tag == "input_text":
-                    assert element.focused
+            for element in obs["dom_elements"]:
+                if element["tag"] == "input_text":
+                    assert element["flags"][0].item()
                     break
             else:
                 assert False, "Input text not found"
             # Type the text
-            target = state.fields["target"]
+            target = info["fields"]["target"]
             if len(target) > 2:
                 # Hmm... Let's try the LEFT arrow key
                 target = target[:-2] + target[-1] + "\ue012" + target[-2]
-            return self.create_type_action(target)
+            return create_type_action(target)
         elif step == 2:
             # Click submit
-            return self.click_button(state, "Submit")
+            return self.click_button(obs, "Submit")
 
 
 class TestEnterTextFocusAndType(RepeatedTester):
     TASK_NAME = "enter-text"
     MAX_STEPS = 2
 
-    def get_action(self, state, step):
+    def get_action(self, obs, info, step):
         if step == 0:
             # Type into the textbox
-            target = state.fields["target"]
-            for element in state.dom_elements:
-                if element.tag == "input_text":
-                    return self.create_focus_and_type_action(element, target)
+            target = info["fields"]["target"]
+            for element in obs["dom_elements"]:
+                if element["tag"] == "input_text":
+                    return create_focus_and_type_action(element["ref"], target)
             assert False, "Input text not found"
         elif step == 1:
             # Click submit
-            return self.click_button(state, "Submit")
+            return self.click_button(obs, "Submit")
 
 
 class TestClickCheckboxes(RepeatedTester):
     TASK_NAME = "click-checkboxes"
     MAX_STEPS = 7
 
-    def get_action(self, state, step):
-        if not state:
+    def get_action(self, obs, info, step):
+        if not obs:
             return
-        # print state.dom.visualize()
+        # print obs.dom.visualize()
         things_to_click = [
-            state.fields[key] for key in state.fields.keys if key != "button"
+            info["fields"][key] for key in info["fields"].keys if key != "button"
         ]
-        for element in state.dom_elements:
-            if element.tag == "label":
-                checkbox, text = element.children
-                if checkbox.value:
-                    things_to_click.remove(text.text)
+        for element in obs["dom_elements"]:
+            if element["tag"] == "label":
+                checkbox, text = (
+                    x for x in obs["dom_elements"] if x["parent"] == element["ref"]
+                )
+                if checkbox["value"]:
+                    things_to_click.remove(text["text"])
                     continue
-                elif text.text in things_to_click:
-                    # Click on <input>:
-                    # return self.create_element_click_action(checkbox)
+                elif text["text"] in things_to_click:
                     # Click on <label>:
-                    return self.create_element_click_action(element)
+                    return create_element_click_action(element["ref"])
         # Click the submit button
         assert not things_to_click
-        return self.click_button(state, "Submit")
+        return self.click_button(obs, "Submit")
 
 
 class TestChooseDateEasy(RepeatedTester):
@@ -251,20 +227,20 @@ class TestChooseDateEasy(RepeatedTester):
     MAX_STEPS = 3
     FRAGILE = True
 
-    def get_action(self, state, step):
+    def get_action(self, obs, info, step):
         if step == 0:
-            for element in state.dom_elements:
-                if element.tag == "input_text":
-                    return self.create_element_click_action(element)
+            for element in obs["dom_elements"]:
+                if element["tag"] == "input_text":
+                    return create_element_click_action(element["ref"])
             assert False, "Input text not found"
         elif step == 1:
-            target = state.fields["day"]
-            for element in state.dom_elements:
-                if element.tag == "a" and element.text == target:
-                    return self.create_element_click_action(element)
+            target = info["fields"]["day"]
+            for element in obs["dom_elements"]:
+                if element["tag"] == "a" and element["text"] == target:
+                    return create_element_click_action(element["ref"])
             assert False, f"Day {target} not found"
         elif step == 2:
-            return self.click_button(state, "Submit")
+            return self.click_button(obs, "Submit")
 
 
 class TestUseAutocomplete(RepeatedTester):
@@ -273,7 +249,7 @@ class TestUseAutocomplete(RepeatedTester):
     FRAGILE = True
 
     def check(self, element, fields):
-        t = element.text
+        t = element["text"]
         if t is None:
             return False
         if "end" in fields.keys:
@@ -281,22 +257,22 @@ def check(self, element, fields):
         else:
             return t.startswith(fields["start"])
 
-    def get_action(self, state, step):
+    def get_action(self, obs, info, step):
         if step == 0:
-            for element in state.dom_elements:
-                if element.tag == "input_text":
-                    return self.create_focus_and_type_action(
-                        element, state.fields["start"]
+            for element in obs["dom_elements"]:
+                if element["tag"] == "input_text":
+                    return create_focus_and_type_action(
+                        element["ref"], info["fields"]["start"]
                     )
             assert False, "Input text not found"
         elif step == 1:
-            # print state.dom.visualize()
-            for element in state.dom_elements:
-                if element.tag == "div" and self.check(element, state.fields):
-                    return self.create_element_click_action(element)
+            # print obs.dom.visualize()
+            for element in obs["dom_elements"]:
+                if element["tag"] == "div" and self.check(element, info["fields"]):
+                    return create_element_click_action(element["ref"])
             assert False, "Correct entry not found"
         elif step == 2:
-            return self.click_button(state, "Submit")
+            return self.click_button(obs, "Submit")
 
 
 class TestUseAutocompleteNoDelay(RepeatedTester):
@@ -305,7 +281,7 @@ class TestUseAutocompleteNoDelay(RepeatedTester):
     FRAGILE = "instance"
 
     def check(self, element, fields):
-        t = element.text
+        t = element["text"]
         if t is None:
             return False
         if "end" in fields.keys:
@@ -313,22 +289,22 @@ def check(self, element, fields):
         else:
             return t.startswith(fields["start"])
 
-    def get_action(self, state, step):
+    def get_action(self, obs, info, step):
         if step == 0:
-            for element in state.dom_elements:
-                if element.tag == "input_text":
-                    return self.create_focus_and_type_action(
-                        element, state.fields["start"]
+            for element in obs["dom_elements"]:
+                if element["tag"] == "input_text":
+                    return create_focus_and_type_action(
+                        element["ref"], info["fields"]["start"]
                     )
             assert False, "Input text not found"
         elif step == 1:
-            # print state.dom.visualize()
-            for element in state.dom_elements:
-                if element.tag == "div" and self.check(element, state.fields):
-                    return self.create_element_click_action(element)
+            # print obs.dom.visualize()
+            for element in obs["dom_elements"]:
+                if element["tag"] == "div" and self.check(element, info["fields"]):
+                    return create_element_click_action(element["ref"])
             assert False, "Correct entry not found"
         elif step == 2:
-            return self.click_button(state, "Submit")
+            return self.click_button(obs, "Submit")
 
 
 class TestClickColor(RepeatedTester):
@@ -349,13 +325,13 @@ class TestClickColor(RepeatedTester):
         (255, 192, 203): "pink",
     }
 
-    def get_action(self, state, step):
-        for element in state.dom_elements:
-            if "color" in element.classes:
-                r, g, b, a = element.bg_color
+    def get_action(self, obs, info, step):
+        for element in obs["dom_elements"]:
+            if "color" in element["classes"]:
+                r, g, b, a = element["bg_color"].tolist()
                 name = self.COLORS[int(r * 255), int(g * 255), int(b * 255)]
-                if name == state.fields["target"]:
-                    return self.create_element_click_action(element)
+                if name == info["fields"]["target"]:
+                    return create_element_click_action(element["ref"])
         assert False, "Correct entry not found"
 
 
@@ -363,61 +339,37 @@ class TestEnterTime(RepeatedTester):
     TASK_NAME = "enter-time"
     MAX_STEPS = 2
 
-    def get_action(self, state, step):
+    def get_action(self, obs, info, step):
         if step == 0:
-            target = state.fields["target"]
+            target = info["fields"]["target"]
             if target.startswith("1:"):
                 target = "0" + target  # Typing '14' will change the number to 2
-            for element in state.dom_elements:
-                if element.tag == "input_time":
-                    return self.create_focus_and_type_action(element, target)
+            for element in obs["dom_elements"]:
+                if element["tag"] == "input_time":
+                    return create_focus_and_type_action(element["ref"], target)
             assert False, "Input text not found"
         elif step == 1:
-            return self.click_button(state, "Submit")
-
-
-# IMPOSSIBLE TO DO RIGHT NOW, since the select items have wrong coordinates
-# class TestChooseList(RepeatedTester):
-#    TASK_NAME = 'choose-list'
-#    MAX_STEPS = 3
-#    FRAGILE = 'delay'
-#
-#    def get_action(self, state, step):
-#        if step == 0:
-#            for element in state.dom_elements:
-#                if element.tag == 'select':
-#                    return self.create_element_click_action(element)
-#            assert False, 'Select not found'
-#        elif step == 1:
-#            print state.dom.visualize()
-#            for element in state.dom_elements:
-#                if element.text == state.fields['target']:
-#                    return self.create_element_click_action(element)
-#            assert False, 'Correct entry not found'
-#        elif step == 2:
-#            return self.click_button(state, 'Submit')
+            return self.click_button(obs, "Submit")
 
 
 class TestClickPie(RepeatedTester):
     TASK_NAME = "click-pie-nodelay"
     MAX_STEPS = 2
 
-    def get_action(self, state, step):
-        print(state.dom.visualize())
+    def get_action(self, obs, info, step):
         if step == 0:
             path = None
-            for element in state.dom_elements:
-                if element.tag == "path":
+            for element in obs["dom_elements"]:
+                if element["tag"] == "path":
                     path = element
-                elif element.text == "+":
-                    return self.create_element_click_action(path)
+                elif element["text"] == "+":
+                    return create_element_click_action(path["ref"])
             assert False, "Select not found"
         elif step == 1:
-            print(state.dom.visualize())
             path = None
-            for element in state.dom_elements:
-                if element.tag == "path":
+            for element in obs["dom_elements"]:
+                if element["tag"] == "path":
                     path = element
-                elif element.text == state.fields["target"]:
-                    return self.create_element_click_action(path)
+                elif element["text"] == info["fields"]["target"]:
+                    return create_element_click_action(path["ref"])
             assert False, "Correct entry not found"
diff --git a/miniwob/tests/test_api.py b/miniwob/tests/test_api.py
new file mode 100644
index 00000000..ee1836cf
--- /dev/null
+++ b/miniwob/tests/test_api.py
@@ -0,0 +1,29 @@
+import gymnasium
+from gymnasium import spaces
+from gymnasium.utils.env_checker import check_env
+from gymnasium.wrappers import FlattenObservation
+
+import miniwob  # noqa: F401
+
+
+class TestGymAPI:
+    def test_click_test_env(self):
+        """Check that the click-test environment follows Gym API."""
+        env = gymnasium.make("miniwob/click-test-v1")
+        check_env(env.unwrapped)
+        env.close()
+
+    def test_flattened_observation_space(self):
+        """Verify the flattened observation space."""
+        env = gymnasium.make("miniwob/login-user-v1")
+        check_env(env.unwrapped)
+        assert isinstance(env.observation_space, spaces.Dict)
+        assert set(env.observation_space) == {"utterance", "dom_elements", "screenshot"}
+        # dom_elements is a Sequence space and cannot be flattened.
+        # But each element in the Sequence can be flattened.
+        env = FlattenObservation(env)
+        assert isinstance(env.observation_space, spaces.Dict)
+        assert isinstance(env.observation_space["utterance"], spaces.Box)
+        assert isinstance(env.observation_space["dom_elements"], spaces.Sequence)
+        assert isinstance(env.observation_space["screenshot"], spaces.Box)
+        env.close()
diff --git a/miniwob/tests/test_environment.py b/miniwob/tests/test_environment.py
index eccbeaf7..9cb11ed2 100644
--- a/miniwob/tests/test_environment.py
+++ b/miniwob/tests/test_environment.py
@@ -3,7 +3,7 @@
 import numpy as np
 import pytest
 
-from miniwob.action import MiniWoBCoordClick, MiniWoBElementClick, MiniWoBTerminate
+from miniwob.action import create_coord_click_action, create_element_click_action
 from miniwob.environment import MiniWoBEnvironment
 
 
@@ -13,8 +13,7 @@ class MiniWoBTester:
 
     @pytest.fixture
     def env(self):
-        env = MiniWoBEnvironment(self.TASK_NAME)
-        env.configure(num_instances=3, seeds=[1, 2, "hello"], headless=True)
+        env = MiniWoBEnvironment(subdomain=self.TASK_NAME)
         yield env
         env.close()
 
@@ -25,23 +24,22 @@ class TestMiniWoBEnvironment(MiniWoBTester):
     ################################
     # Helpers
 
-    def get_coord_click(self, state):
+    def get_coord_click(self, obs):
         """Get the action that clicks the button."""
-        for element in state.dom_elements:
-            if element.tag == "button":
-                action = MiniWoBCoordClick(element.left + 5, element.top + 5)
-                print(f"Clicking with {action}")
+        for element in obs["dom_elements"]:
+            if element["tag"] == "button":
+                left, top = element["pos"].tolist()
+                action = create_coord_click_action(left + 5, top + 5)
                 return action
-        raise ValueError(f"Cannot find button: {str(state.dom_elements)}")
+        raise ValueError("Cannot find button.")
 
-    def get_element_click(self, state):
+    def get_element_click(self, obs):
         """Get the action that clicks the button."""
-        for element in state.dom_elements:
-            if element.tag == "button":
-                action = MiniWoBElementClick(element, fail_hard=True)
-                print(f"Clicking with {action}")
+        for element in obs["dom_elements"]:
+            if element["tag"] == "button":
+                action = create_element_click_action(element["ref"])
                 return action
-        raise ValueError(f"Cannot find button: {str(state.dom_elements)}")
+        raise ValueError("Cannot find button.")
 
     # get_click = get_coord_click
     get_click = get_element_click
@@ -50,129 +48,79 @@ def get_element_click(self, state):
     # Tests
 
     def test_do_nothing(self, env):
-        """Test the ability to start up multiple instances."""
-        assert len(env.instances) == 3
+        """Test the ability to start an instance for the click-test task."""
+        obs, info = env.reset()
+        assert obs["utterance"] == "Click the button."
+        assert any(element["tag"] == "button" for element in obs["dom_elements"])
 
     def test_run(self, env):
-        """Test reset and step."""
-        print("=" * 40)
-        states = env.reset()
-        print([x.utterance for x in states])
-        assert all(x.utterance == "Click the button." for x in states)
-        print([x.fields for x in states])
-        assert all(x.fields.keys == ["dummy"] for x in states)
-        print([x.dom for x in states])
-        print(states[0].dom_elements)
-        print(states[0].dom.visualize())
-        ################################
-        print("=" * 40)
-        states, rewards, dones, info = env.step([None, None, None])
-        print([x.utterance for x in states])
-        print([x.dom for x in states])
-        print(rewards)
-        assert rewards == [0, 0, 0]
-        print(dones)
-        assert dones == [False, False, False]
-        print(info)
-        ################################
+        """Test reset() and step()."""
+        obs, info = env.reset()
+        assert obs["utterance"] == "Click the button."
+        # Test empty action
+        obs, reward, terminated, truncated, info = env.step(None)
+        assert obs["utterance"] == "Click the button."
+        assert reward == 0
+        assert terminated is False
+        assert terminated is False
         # Test clicking
-        print("=" * 40)
-        action = self.get_click(states[1])
-        states, rewards, dones, info = env.step([None, action, None])
-        assert states[1] is None
-        states = [states[0], states[2]]
-        print([x.utterance for x in states])
-        print([x.dom for x in states])
-        print(rewards)
-        assert rewards[0] == 0 and rewards[1] > 0 and rewards[2] == 0
-        print(dones)
-        assert dones == [False, True, False]
-        print(info)
-        ################################
+        action = self.get_click(obs)
+        obs, reward, terminated, truncated, info = env.step(action)
+        assert reward > 0
+        assert terminated is True
+        assert truncated is False
+        # Test reset
+        obs, info = env.reset()
+        assert obs["utterance"] == "Click the button."
+        # Test clicking again
+        action = self.get_click(obs)
+        obs, reward, terminated, truncated, info = env.step(action)
+        assert reward > 0
+        assert terminated is True
+        assert truncated is False
+
+    def test_timeout(self, env):
+        """Test environment timeout."""
+        obs, info = env.reset()
+        assert obs["utterance"] == "Click the button."
         # Wait for timeout
-        print("=" * 40)
-        for i in range(1, 13):
-            print(f"Sleeping ... ({i})")
-            time.sleep(1)
-        states, rewards, dones, info = env.step([None, None, None])
-        assert states == [None, None, None]
-        print(rewards)
-        assert rewards[0] < 0 and rewards[1] > 0 and rewards[2] < 0
-        print(dones)
-        assert dones == [True, True, True]
-        print(info)
-        ################################
+        time.sleep(12)
+        obs, reward, terminated, truncated, info = env.step(None)
+        assert reward < 0
+        assert terminated is True
+        assert truncated is False
         # Start again
-        print("=" * 40)
-        states = env.reset()
-        print([x.utterance for x in states])
-        print([x.dom for x in states])
-        ################################
-        print("=" * 40)
-        states, rewards, dones, info = env.step([None, None, None])
-        print([x.utterance for x in states])
-        print([x.dom for x in states])
-        print(rewards)
-        assert rewards == [0, 0, 0]
-        print(dones)
-        assert dones == [False, False, False]
-        print(info)
+        obs, info = env.reset()
+        assert obs["utterance"] == "Click the button."
+        obs, reward, terminated, truncated, info = env.step(None)
+        assert reward == 0
+        assert terminated is False
+        assert truncated is False
 
     def test_speed(self, env):
-        states = env.reset()
+        """Test the processing speed for step()."""
+        obs, info = env.reset()
         start_time = time.time()
         elapsed = []
         N = 50
         for i in range(1, N + 1):
             print("Iteration", i, "/", N)
-            actions = [None] * len(states)
-            states, rewards, dones, info = env.step(actions)
+            obs, reward, terminated, truncated, info = env.step(None)
+            assert terminated is False
             elapsed.append(time.time() - start_time)
             start_time = time.time()
         mean = sum(elapsed) / len(elapsed)
         print("Average time:", mean)
         print("SD:", sum((x - mean) ** 2 for x in elapsed) / len(elapsed))
 
-    def test_reset(self, env):
-        print("=" * 40)
-        states = env.reset()
-        # Test clicking
-        action = self.get_click(states[1])
-        states, rewards, dones, info = env.step([None, action, None])
-        assert dones == [False, True, False]
-        print(info)
-        # Should issue a warning on the second instance
-        states, rewards, dones, info = env.step([None, action, None])
-        assert dones == [False, True, False]
-        print(info)
-        # Hard reset
-        states = env.reset()
-        states, rewards, dones, info = env.step([None, None, None])
-        assert dones == [False, False, False]
-        print(info)
-        # Now click the first and second
-        action_0 = self.get_click(states[0])
-        action_1 = self.get_click(states[1])
-        states, rewards, dones, info = env.step([action_0, action_1, None])
-        assert dones == [True, True, False]
-        print(info)
-
     def test_attention(self, env):
-        print("=" * 40)
+        """Test that visualize_attention() does not crash."""
         env.reset()
         attention = np.random.rand(20, 20) * 0.02
-        env.visualize_attention([attention, None, None])
-        time.sleep(2)
-        env.visualize_attention([[[]], None, None])
-        time.sleep(2)
-
-    def test_suicide(self, env):
-        print("=" * 40)
-        states = env.reset()
-        action = MiniWoBTerminate()
-        states, rewards, dones, info = env.step([None, action, None])
-        assert dones == [False, True, False]
-        assert rewards[1] == -1
+        env.visualize_attention(attention)
+        time.sleep(1)
+        env.visualize_attention(None)
+        time.sleep(1)
 
 
 ################################################
@@ -181,118 +129,92 @@ def test_suicide(self, env):
 class TestMiniWoBSeed(MiniWoBTester):
     TASK_NAME = "click-button"
 
-    def get_button_click(self, state):
+    def get_button_click(self, obs, info):
         """Get the action that clicks the button."""
-        for element in state.dom_elements:
-            if element.tag == "button" and element.text == state.fields["target"]:
-                action = MiniWoBElementClick(element, fail_hard=True)
-                print(f"Clicking with {action}")
+        for element in obs["dom_elements"]:
+            if (
+                element["tag"] == "button"
+                and element["text"] == info["fields"]["target"]
+            ):
+                action = create_element_click_action(element["ref"])
                 return action
-        raise ValueError(f"Cannot find button: {str(state.dom_elements)}")
-
-    def get_input_click(self, state):
-        """Get the action that clicks an input element."""
-        for element in state.dom_elements:
-            if element.tag == "input_text":
-                action = MiniWoBElementClick(element, fail_hard=True)
-                print(f"Clicking with {action}")
-                return action
-        raise ValueError(f"Cannot find input: {str(state.dom_elements)}")
+        raise ValueError("Cannot find button.")
 
     def test_seed(self, env):
-        print("=" * 40)
-        states = env.reset(seeds=["hello", "itsme", "hello"])
-        print(states[0].dom.visualize())
-        # Check that everything is the same for instances 0 and 2
-        utt_hello = states[0].utterance
-        utt_itsme = states[1].utterance
-        assert utt_hello == states[2].utterance != utt_itsme
-        ref_to_text_hello = {x.ref: x.text for x in states[0].dom_elements}
-        ref_to_text_itsme = {x.ref: x.text for x in states[1].dom_elements}
-        assert ref_to_text_hello == {x.ref: x.text for x in states[2].dom_elements}
-        # Test clicking an input (no effect)
-        actions_input = [self.get_input_click(state) for state in states]
-        actions_button = [self.get_button_click(state) for state in states]
-        states, rewards, dones, info = env.step(actions_input)
-        assert dones == [False, False, False]
-        assert ref_to_text_hello == {x.ref: x.text for x in states[0].dom_elements}
-        assert ref_to_text_itsme == {x.ref: x.text for x in states[1].dom_elements}
-        assert ref_to_text_hello == {x.ref: x.text for x in states[2].dom_elements}
-        # Now click the correct button using the actions computed before
-        states, rewards, dones, info = env.step(actions_button)
-        assert dones == [True, True, True]
-        # Now run everything again but with shuffled seeds
-        states = env.reset(seeds=["hello", "hello", "itsme"])
-        assert states[0].utterance == states[1].utterance == utt_hello
-        assert states[2].utterance == utt_itsme
-        assert ref_to_text_hello == {x.ref: x.text for x in states[0].dom_elements}
-        assert ref_to_text_hello == {x.ref: x.text for x in states[1].dom_elements}
-        assert ref_to_text_itsme == {x.ref: x.text for x in states[2].dom_elements}
-        # Test clicking with the old action objects
-        actions_input = [actions_input[2], actions_input[0], actions_input[1]]
-        states, rewards, dones, info = env.step(actions_input)
-        assert dones == [False, False, False]
-        actions_button = [actions_button[2], actions_button[0], actions_button[1]]
-        states, rewards, dones, info = env.step(actions_button)
-        assert dones == [True, True, True]
+        obs_1, info_1 = env.reset(seed=31416)
+        obs_2, info_2 = env.reset(seed=227)
+        obs_3, info_3 = env.reset(seed=227)
+        obs_4, info_4 = env.reset(seed=31416)
+        # Check that everything is the same for the same seed
+        assert obs_1["utterance"] == obs_4["utterance"]
+        assert obs_2["utterance"] == obs_3["utterance"]
+        ref_to_text_1 = {x["ref"]: x["text"] for x in obs_1["dom_elements"]}
+        ref_to_text_2 = {x["ref"]: x["text"] for x in obs_2["dom_elements"]}
+        ref_to_text_3 = {x["ref"]: x["text"] for x in obs_3["dom_elements"]}
+        ref_to_text_4 = {x["ref"]: x["text"] for x in obs_4["dom_elements"]}
+        assert ref_to_text_1 == ref_to_text_4
+        assert ref_to_text_2 == ref_to_text_3
+        assert ref_to_text_1 != ref_to_text_2
+        # Compute the correct action from obs 1
+        # and apply it on obs 4 (same seed)
+        action = self.get_button_click(obs_1, info_1)
+        obs, reward, terminated, truncated, info = env.step(action)
+        assert reward > 0
+        assert terminated is True
 
 
 class TestMiniWoBMode(MiniWoBTester):
     TASK_NAME = "click-test-transfer"
 
-    def get_button_click(self, state, text):
+    def get_button_click(self, obs, text):
         """Get the action that clicks the button."""
-        for element in state.dom_elements:
-            if element.tag == "button" and element.text == text:
-                action = MiniWoBElementClick(element, fail_hard=True)
-                print(f"Clicking with {action}")
+        for element in obs["dom_elements"]:
+            if element["tag"] == "button" and element["text"] == text:
+                action = create_element_click_action(element["ref"])
                 return action
-        raise ValueError(f"Cannot find button: {str(state.dom_elements)}")
+        raise ValueError("Cannot find button.")
 
     def test_mode(self, env):
         """Test if setting the mode works.
         - mode = 'train': click on button ONE
         - mode = 'test':  click on button TWO
         """
-        print("=" * 40)
         # Training time
-        states = env.reset()
-        targets = ["ONE", "TWO", "ONE"]
-        actions = []
-        for state, target in zip(states, targets):
-            assert state.utterance == "Click button ONE."
-            actions.append(self.get_button_click(state, target))
-        states, rewards, dones, info = env.step(actions)
-        assert dones == [True, True, True]
-        assert rewards[0] > 0 and rewards[1] < 0 and rewards[2] > 0
+        obs, info = env.reset()
+        assert obs["utterance"] == "Click button ONE."
+        action = self.get_button_click(obs, "ONE")
+        obs, reward, terminated, truncated, info = env.step(action)
+        assert reward > 0
+        obs, info = env.reset()
+        assert obs["utterance"] == "Click button ONE."
+        action = self.get_button_click(obs, "TWO")
+        obs, reward, terminated, truncated, info = env.step(action)
+        assert reward < 0
         # Test time
-        env.set_mode("test")
-        states = env.reset()
-        actions = []
-        for state, target in zip(states, targets):
-            assert state.utterance == "Click button TWO."
-            actions.append(self.get_button_click(state, target))
-        states, rewards, dones, info = env.step(actions)
-        assert dones == [True, True, True]
-        assert rewards[0] < 0 and rewards[1] > 0 and rewards[2] < 0
+        env.set_data_mode("test")
+        obs, info = env.reset()
+        assert obs["utterance"] == "Click button TWO."
+        action = self.get_button_click(obs, "ONE")
+        obs, reward, terminated, truncated, info = env.step(action)
+        assert reward < 0
         # Test time again; mode should be persistent
-        states = env.reset()
-        actions = []
-        for state, target in zip(states, targets):
-            assert state.utterance == "Click button TWO."
-            actions.append(self.get_button_click(state, target))
-        states, rewards, dones, info = env.step(actions)
-        assert dones == [True, True, True]
-        assert rewards[0] < 0 and rewards[1] > 0 and rewards[2] < 0
-        # Training time again: set mode with reset(mode=...)
-        states = env.reset(mode="train")
-        actions = []
-        for state, target in zip(states, targets):
-            assert state.utterance == "Click button ONE."
-            actions.append(self.get_button_click(state, target))
-        states, rewards, dones, info = env.step(actions)
-        assert dones == [True, True, True]
-        assert rewards[0] > 0 and rewards[1] < 0 and rewards[2] > 0
+        obs, info = env.reset()
+        assert obs["utterance"] == "Click button TWO."
+        action = self.get_button_click(obs, "TWO")
+        obs, reward, terminated, truncated, info = env.step(action)
+        assert reward > 0
+        # Training time again: set mode with reset()
+        obs, info = env.reset(options={"data_mode": "train"})
+        assert obs["utterance"] == "Click button ONE."
+        action = self.get_button_click(obs, "ONE")
+        obs, reward, terminated, truncated, info = env.step(action)
+        assert reward > 0
+        # Training time again; mode should be persistent
+        obs, info = env.reset()
+        assert obs["utterance"] == "Click button ONE."
+        action = self.get_button_click(obs, "TWO")
+        obs, reward, terminated, truncated, info = env.step(action)
+        assert reward < 0
 
 
 ################################################
@@ -302,29 +224,19 @@ class TestMiniWoBFields(MiniWoBTester):
     TASK_NAME = "email-inbox-forward-nl"
 
     def test_fields(self, env):
-        print("=" * 40)
         # Training time
-        states = env.reset()
-        for state in states:
-            print(state.utterance)
-            print(state.fields)
-            assert "by" in state.fields.keys
-            assert "to" in state.fields.keys
-            assert state.fields["by"] in state.utterance
-            assert state.fields["to"] in state.utterance
+        obs, info = env.reset()
+        assert "by" in info["fields"].keys
+        assert "to" in info["fields"].keys
+        assert info["fields"]["by"] in obs["utterance"]
+        assert info["fields"]["to"] in obs["utterance"]
         # Test time
-        states = env.reset(mode="test")
-        for state in states:
-            print(state.utterance)
-            print(state.fields)
-            assert state.fields.keys == ["dummy"]
-            assert state.utterance
-        # Training time
-        states = env.reset(mode="train")
-        for state in states:
-            print(state.utterance)
-            print(state.fields)
-            assert "by" in state.fields.keys
-            assert "to" in state.fields.keys
-            assert state.fields["by"] in state.utterance
-            assert state.fields["to"] in state.utterance
+        obs, info = env.reset(options={"data_mode": "test"})
+        assert info["fields"].keys == ["dummy"]
+        assert obs["utterance"]
+        # Training time again
+        obs, info = env.reset(options={"data_mode": "train"})
+        assert "by" in info["fields"].keys
+        assert "to" in info["fields"].keys
+        assert info["fields"]["by"] in obs["utterance"]
+        assert info["fields"]["to"] in obs["utterance"]