Sliding Puzzles Gym: A Scalable Benchmark for State Representation in Visual Reinforcement Learning

We thank the reviewer for their positive assessment of our experimental design and literature review, and for the opportunity to clarify the core motivation behind SPGym regarding the evaluation of visual representation capabilities.

Why SPGym Tests Visual Representation Capabilities

The central idea behind SPGym is to create a controlled environment where the primary challenge being scaled is the agent's ability to process and understand diverse visual inputs, even though the agent is trained end-to-end on a standard RL task. We achieve this through specific design choices:

1. POMDP Formulation: As detailed in Section 3, SPGym is formulated as a Partially Observable Markov Decision Process (POMDP) defined by $(\mathcal{S}, \mathcal{X}, \mathcal{A}, \mathcal{P}, \mathcal{R}, \mathcal{S}_0)$. Critically, the agent never observes the true underlying puzzle state $s \in \mathcal{S}$. Instead, it only receives visual observations $x \in \mathcal{X}$. The agent must infer the relevant state information solely from these high-dimensional visual inputs.

2. Isolating the Visual Challenge: The key design element is that across all experiments and all compared agents, the core components of the MDP remain fixed: the state space $\mathcal{S}$ (tile permutations), the action space $\mathcal{A}$ (up, down, left, right), the deterministic transition dynamics $\mathcal{P}$ (how tiles move), the reward function $\mathcal{R}$ (based on Manhattan distance), and the initial state distribution $\mathcal{S}_0$. The only thing that changes between experimental conditions (e.g., different pool sizes) or agent comparisons (e.g., different representation learning modules) is the emission function that maps the underlying state $s$ to the visual observation $x$. We vary this by overlaying the puzzle state with different pools of images.

3. Performance Reflects Representation Quality: Because all core task elements (dynamics, rewards, etc.) are constant, any difference in performance (e.g., sample efficiency) between agents or across different visual diversity levels must be attributed to how effectively the agent's visual encoder maps the observation $x$ to a useful internal representation. An agent cannot solve the puzzle without implicitly or explicitly understanding the tile configuration from the image. Better visual representations enable the policy to make better decisions, leading to faster learning and higher success rates. While policy learning is intrinsically linked, the bottleneck being systematically stressed and evaluated is the visual representation learning component.

Empirical Evidence

To provide further empirical support for this link between representation quality and task performance, we conducted linear probe evaluations (inspired by feedback from reviewer o6ZL) on the frozen encoders learned by PPO and SAC agents. We trained linear classifiers to predict tile positions from the learned features. Our results show a statistically significant negative correlation (Pearson r=-0.81, p=1.1e-13) between the probe's test accuracy and the number of environment steps the RL agent needed to reach 80% success. This demonstrates that encoders capable of producing more informative features (higher probe accuracy) enable faster learning on the downstream RL task. The probe accuracies for tested agents follows:

Conclusion

In essence, while the agent learns end-to-end, SPGym isolates the difficulty scaling to the visual domain. An agent succeeding in SPGym, especially across varying image pools, demonstrates robustness in its visual processing specific to the task structure. We believe this setup effectively probes an agent's ability to form useful representations from pixels under varying visual conditions, which is a critical aspect of visual RL.

We will revise the Introduction and Methodology sections to incorporate this detailed explanation and motivation, ensuring it is clear how our design choices allow for the evaluation of visual representation learning capabilities in a controlled manner. We hope this addresses the reviewer's concern and clarifies the value proposition of SPGym. We appreciate the reviewer's willingness to reconsider their evaluation based on this clarification.

We thank the reviewer for the positive feedback on SPGym's potential and the detailed, constructive comments regarding the evaluation settings and experimental design. We address the concerns below:

1. In-Distribution Evaluation

We understand the concern that Table 2 results might seem like overfitting. However, our primary goal with the in-distribution (ID) setting is not to evaluate generalization across images but to measure the sample efficiency with which different RL agents learn useful visual representations under controlled visual diversity. We chose pool size 5 because it offers a balance: it introduces enough visual diversity to discriminate between different representation learning approaches (as seen by the varying steps-to-convergence) while remaining learnable for most agents within our computational budget.

As detailed in our response to Reviewer R6Sn, SPGym is formulated as a POMDP where only the visual observation function changes between settings; the underlying task dynamics remain fixed. Therefore, differences in sample efficiency (Table 2) directly reflect how effectively each agent's representation learning component handles the visual aspect of the task.

2. Out-of-Distribution Evaluation

We acknowledge the reviewer's point that zero OOD success might seem like a limitation. However, we view this consistent failure not as a flaw of the benchmark, but as a crucial diagnostic finding about the limitations of current end-to-end visual RL methods when faced with generalizing visual features learned purely through a task-specific RL objective. Standard RL objectives, like the one used here, do not explicitly optimize for generalization to unseen visual appearances of the same underlying state. SPGym effectively highlights this gap.

Our choice to present the main OOD results (Table 3) using agents trained on pool size 5 stems from this being the setting where most methods achieved high ID success. Evaluating OOD generalization is most meaningful when agents have demonstrated competence on their training distribution. While some base agents achieved ID success on larger pools, they still failed completely OOD, reinforcing the finding that simply increasing training diversity within this range wasn't sufficient for generalization with current methods. Showing zero OOD success across methods aims to motivate research into algorithms with better visual generalization properties.

3. Training Set Size

We appreciate the suggestion to explore significantly larger training pools. In preliminary experiments, we attempted training with pools of thousands of images. However, we found that with such high visual diversity, the RL training became unstable. Because each observation was essentially unique (rarely seen more than once), the gradients from the RL objective were insufficient to train a useful visual encoder from scratch. Without a stable encoder, the policy failed to learn.

Therefore, our approach was to identify the approximate limit of visual diversity that current standard algorithms (PPO, SAC, DreamerV3) could handle within a reasonable budget (10M steps), which led to testing pools up to 20, 50, and 100, respectively. This reveals the scaling limitations of these methods rather than aiming for OOD generalization via massive datasets, which might require different training paradigms (e.g., pretraining on external data, different objectives).

4. Tiered Generalization Evaluation

We thank the reviewer for this excellent suggestion. We agree that evaluating generalization across different levels of visual similarity would provide deeper insights. As a preliminary step in this direction, we evaluated the trained PPO and SAC agents on augmented versions of their training images. We observed a strong correlation (Pearson r=-0.81, p=2.5e-12) between success rates on these augmented images and the agents' sample efficiency, suggesting a link between robustness to simple transformations and learning speed. We plan to include these results and discuss the tiered approach as a key direction for future work in the revised manuscript:

We will revise the paper to clarify the evaluation rationale, incorporate the results on augmented images, and explicitly frame the tiered generalization as important future work, addressing the points raised. Thank you again for the constructive feedback.

We sincerely thank the reviewer for their thorough and constructive feedback. We appreciate the recognition of SPGym's potential and the detailed suggestions, which will significantly improve the paper.

1. Answers to Direct Questions

1. Missing Tile: The missing tile's starting position is randomized in each episode as part of the initial shuffling.
2. Training Termination ("Early Stopping"): Our primary motivation for terminating runs based on sustained success rate was computational efficiency, given the number of experiments. A secondary reason was to evaluate OOD generalization before potential extreme overfitting. We agree "early stopping" is imprecise terminology, as agents can still improve solution efficiency, and will rephrase this in the revision.
3. PPO 4x4 Convergence: PPO does not converge within 10M steps on the 4x4 grid. Extended runs (100M steps) show it requires approximately 24.5M ± 7.6M steps (avg. 5 seeds) to reach 80% success, indicating feasibility but requiring much more data. We will add this finding.

2. Addressing Key Concerns

1. Disentanglement of Representation and Policy Learning: We acknowledge the reviewer's point on end-to-end learning. However, SPGym's design isolates the visual representation challenge as the key variable by keeping all underlying MDP components fixed and only varying the visual observation function. Performance differences thus reflect agents' visual processing capabilities. For a detailed explanation of the POMDP formulation and supporting empirical evidence from linear probing, please see our response to Reviewer R6Sn. We will add these clarifications and results to the revision.

2. Optimal Solution Claim: You are correct. Given our termination criterion and that only PPO was tested for continued training, the claim that best-performing agents fall short might be too strong. We will revise this to reflect that agents under our protocol didn't reach the optimum, acknowledging the potential for improvement with longer training.

3. Reward Function: The Manhattan distance reward is standard in puzzle literature (Korf et al., 1985; Burns et al., 2012; Lee et al., 2022; Moon et al., 2024) and encourages minimizing steps by accumulating negative rewards until the goal (+1) is reached. While local optima are possible, this aligns the RL objective with finding efficient solutions. A perfect dense reward is non-trivial, and sparse rewards are infeasible here.

4. Evaluation Criteria Clarity (80% Threshold): We apologize for the omission. We calculate this by finding the first environment step where the average success rate across all parallel environments reaches 80% for each run, then averaging these step counts across seeds. We will clarify this.

5. Data Augmentation & Hyperparameter Selection: We acknowledge the limitation of optimizing augmentations only for RAD and applying that strategy universally. This was a trade-off for controlled comparison vs. method-specific tuning. Similarly, hyperparameters were tuned for training efficiency. Evaluating based on generalization is important future work. We will clarify these limitations.

6. Fairness of Pretraining Comparison (PPO): We agree it's not a direct 'from scratch' comparison. The intent was to assess the impact of leveraging pretrained features (ID/OOD) vs. learning purely from the RL signal in SPGym. We will clarify this motivation.

7. Generalization Evaluation (Table 3 & Pool Size Choice): The consistent OOD failure across methods and pool sizes is a key diagnostic finding about current end-to-end visual RL limitations. Our rationale for focusing Table 3 on pool size 5, preliminary findings on larger pools, and the interpretation of the OOD results are discussed in detail in our response to Reviewer ukCk (due to space limitations).

8. Procedural Generation: We agree the current results don't show a demonstrated advantage over ImageNet. We will moderate the claim, framing it as a feature with potential benefits (diversity, control, storage) for future investigation.

3. Other Points

• Missing References: Thank you. We will add CARLA and COOM to the related work discussion.
• Code: Code is in the supplementary material, omitted from the paper for anonymity. We will add links in the camera-ready version.
• Minor Points: We will refine terminology (model-based definition, visual diversity) and acknowledge unevaluated features (other modalities) and future analysis directions (image difficulty).

We hope these responses and planned revisions address the reviewer's concerns. We value the feedback and are committed to improving the paper. We appreciate the reviewer's openness to reconsidering their evaluation.