Spatial-Gym: A Gymnasium Environment for Spatial Reasoning Benchmarking

Abstract

Spatial-Gym is a Gymnasium-compatible environment designed for evaluating spatial reasoning capabilities of Large Language Models (LLMs) and other AI agents. Built upon the spatial puzzle dataset introduced in SPaRC (Kaesberg et al.), this environment provides a standardized interface for benchmarking agent performance on grid-based spatial reasoning tasks. The environment supports multiple observation formats, customizable rendering modes for human and LLM interaction, and comprehensive evaluation metrics for systematic analysis of spatial reasoning abilities.

Key Features

• Standardized RL Interface: Full Gymnasium API compliance for seamless integration with existing RL frameworks
• Dual Observation Modes: Structured tensor representation or JSON-based symbolic encoding
• Multi-Modal Rendering: Human-readable visualizations and LLM-optimized text representations
• Flexible Dataset Support: Compatible with HuggingFace datasets following the SPaRC format
• Comprehensive Metrics: Episode-level tracking of success rate, path efficiency, and reasoning patterns
• Backtracking Support: Optional state reversibility for exploring different solution strategies

Installation

From PyPI

pip install Spatial-Gym

From Source

git clone https://github.com/lkaesberg/Spatial-Gym.git
cd Spatial-Gym
pip install -e .

Quick Start

import gymnasium as gym
import Spatial_Gym

# Initialize environment with default configuration
env = gym.make(
    "Spatial-Gym",
    df_name='lkaesberg/SPaRC',
    df_split='all',
    df_set='test',
    render_mode='human',
    observation='new',
    traceback=True,
    max_steps=1000
)

# Standard RL loop
observation, info = env.reset()
terminated = False

while not terminated:
    action = env.action_space.sample()  # Replace with your agent
    observation, reward, terminated, truncated, info = env.step(action)
    env.render()

env.close()

Environment Configuration

Environment Specification

Action Space

Discrete(4): Four directional moves in the grid environment.

Observation Space

Tensor Format (observation='new')

A dictionary containing:

• base (Dict[str, np.ndarray]): One-hot encoded spatial features
  • visited: Binary grid marking visited cells
  • gaps: Binary grid indicating traversable/non-traversable cells
  • agent_location: One-hot encoding of agent position
  • target_location: One-hot encoding of goal position
  • Additional puzzle-specific properties (e.g., stars, triangles)
• color (np.ndarray): Integer grid (1-8) representing color properties
• additional_info (np.ndarray): Puzzle-specific metadata (polyshape IDs, counts)

JSON Format (observation='SPaRC')

String-encoded JSON representing the grid state with symbolic notation, following the original SPaRC specification.

Reward Structure

• +1.0: Successfully solving the puzzle
• -1.0: Invalid termination or failure state
• +0.01: Incremental reward for remaining on valid solution path (encourages exploration while maintaining progress)

Episode Termination

• Success: Agent reaches target location satisfying all puzzle constraints
• Failure: Agent enters invalid state or violates puzzle rules
• Truncation: Maximum step limit reached

API Reference

Core Methods

env.reset(options: Optional[Dict] = None) -> Tuple[Observation, Dict]

Initializes or resets the environment to a new puzzle state.

• Parameters:
  • options: Optional dictionary with 'puzzle_id' key to load specific puzzle
• Returns: Initial observation and info dictionary

env.step(action: int) -> Tuple[Observation, float, bool, bool, Dict]

Executes one environment step given an action.

• Parameters:
  • action: Integer in range [0, 3] representing directional move
• Returns: Observation, reward, terminated flag, truncated flag, info dictionary

env.render() -> Optional[np.ndarray]

Generates visual or textual representation of current state based on render_mode.

env.close()

Releases environment resources and closes rendering windows.

Experimental Setup

Dataset

The environment uses puzzles from the SPaRC dataset, which contains spatial reasoning challenges of varying complexity. Each puzzle is defined by:

• Grid dimensions (variable size)
• Initial agent position
• Target position
• Spatial constraints (gaps, regions, colored elements)
• Solution paths of varying lengths

Evaluation Metrics

The info dictionary returned by step() and reset() contains:

• Success Rate: Binary indicator of puzzle completion
• Path Length: Number of steps taken
• Optimality: Ratio of actual path length to shortest possible path
• Invalid Actions: Count of rule violations
• Puzzle Metadata: Difficulty rating, constraint types, grid size

Use Cases

Benchmarking LLM Spatial Reasoning

import gymnasium as gym
import Spatial_Gym
from your_llm_wrapper import LLMAgent

env = gym.make("Spatial-Gym", render_mode='llm', observation='SPaRC')
agent = LLMAgent(model="gpt-4")

observation, info = env.reset()
for _ in range(100):  # Evaluate on 100 puzzles
    done = False
    while not done:
        action = agent.predict(env.render())  # LLM sees text representation
        observation, reward, terminated, truncated, info = env.step(action)
        done = terminated or truncated
    
    # Log metrics
    print(f"Puzzle {info['puzzle_id']}: Success={info['success']}, Steps={info['steps']}")
    observation, info = env.reset()

Reinforcement Learning Training

from stable_baselines3 import PPO

env = gym.make("Spatial-Gym", observation='new', max_steps=500)
model = PPO("MultiInputPolicy", env, verbose=1)
model.learn(total_timesteps=1_000_000)
model.save("spatial_reasoning_agent")

Prime Intellect verifiers environment

environments/spatial_gym/ wraps Spatial-Gym as a verifiers environment for evaluation and RL training. It covers all three settings of the Spatial-Gym paper (step-by-step, step-by-step with backtracking, and one-shot). The reward is 1.0 for a correct path and 0.0 otherwise. See environments/spatial_gym/README.md.

uv pip install -e environments/spatial_gym
vf-eval spatial-gym-env -m <model> --client.base-url <url> --client.api-key-var <ENV_VAR> -n 20 --no-push

Repository Structure

Spatial-Gym/
├── Spatial_Gym/           # Core environment implementation
│   ├── __init__.py        # Package initialization
│   ├── Spatial_Gym.py     # Main environment class
│   ├── register_env.py    # Gymnasium registration
│   └── render/            # Rendering modules
│       ├── human_renderer.py
│       └── llm_renderer.py
├── llm_testing/           # LLM evaluation utilities
│   ├── llm_host.py        # LLM interaction wrapper
│   └── parse_logs.py      # Result analysis tools
├── Final_Product.py       # Interactive demo script
├── human_play.py          # Human player interface
├── pyproject.toml         # Package configuration
└── README.md              # This file

Testing

Spatial-Gym includes a comprehensive test suite to ensure environment stability and correctness.

Running Tests

# Install with test dependencies
pip install -e ".[test]"

# Run all tests
pytest tests/ -v

# Run specific test categories
pytest tests/test_environment.py -v       # Environment API tests
pytest tests/test_random_agent.py -v      # Random agent tests
pytest tests/test_predefined_paths.py -v  # Path validation tests

# Run with coverage
pytest tests/ --cov=Spatial_Gym --cov-report=html

Test Coverage

The test suite includes 43+ tests covering:

• ✅ Environment initialization and configuration
• ✅ Gymnasium API compliance
• ✅ Random agent behavior (stress tests)
• ✅ Predefined valid and invalid paths
• ✅ Multi-episode stability
• ✅ Different observation formats
• ✅ Rendering modes

Continuous Integration

Tests automatically run on:

• Every push and pull request
• Multiple OS (Ubuntu, macOS)
• Python versions 3.9, 3.10, 3.11

See tests/README.md for detailed testing documentation.

Citation

If you use Spatial-Gym, please cite:

@misc{kaesberg2026spatialgym,
  title  = {Mind the Gap Between Spatial Reasoning and Acting! Step-by-Step Evaluation of Agents With Spatial-Gym},
  author = {Kaesberg, Lars Benedikt and Yang, Tianyu and Bauer, Niklas and Ruas, Terry and Wahle, Jan Philip and Gipp, Bela},
  year   = {2026},
  eprint = {2604.09338},
  archivePrefix = {arXiv}
}

Spatial-Gym is built on the SPaRC puzzles. Please also cite:

@inproceedings{kaesberg-etal-2025-sparc,
  title     = "{SP}a{RC}: A Spatial Pathfinding Reasoning Challenge",
  author    = "Kaesberg, Lars Benedikt and Wahle, Jan Philip and Ruas, Terry and Gipp, Bela",
  booktitle = "Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing",
  month     = nov,
  year      = "2025",
  address   = "Suzhou, China",
  publisher = "Association for Computational Linguistics",
  url       = "https://aclanthology.org/2025.emnlp-main.526/",
  doi       = "10.18653/v1/2025.emnlp-main.526",
  pages     = "10359--10390"
}

Related work from our group

Contributing

We welcome contributions! Please follow these guidelines:

1. Fork the repository
2. Create a feature branch (git checkout -b feature/new-feature)
3. Commit changes with descriptive messages
4. Add tests for new functionality
5. Submit a pull request

For bug reports and feature requests, please use the GitHub issue tracker.

License

This project is licensed under the MIT License - see the LICENCE file for details.

Contact

For questions, suggestions, or collaboration inquiries, please contact:

• GitHub Issues: https://github.com/lkaesberg/Spatial-Gym/issues