A neural policy solving structured STRIPS planning tasks from first principles.
Combining Potential-Based Reward Shaping (PBRS), Environmental Action Masking, and Batched REINFORCE with entropy regularization to solve 5-block re-stacking problems, paired with an interactive Next.js trajectory inspector featuring an articulated overhead robotic gantry.

---

📑 Table of Contents

• 🎯 Aim of the Project
• 📈 Experimental Evolution & Findings (v1 vs. v2 vs. v3)
  • Side-by-Side Diagnostic Dashboards
  • Unified Comparative Analysis
  • Progression Matrix
• 🧩 Problem Formulation & Environment Design
  • STRIPS First-Order Predicates
  • Action Space & State Transitions
  • Policy Network & Action Masking
• ⚖️ Reward Engineering & Potential Shaping
• 🕹️ Interactive Web Visualizer
  • Visualizer Highlights
  • Keyboard Shortcuts
• 💻 Technical Tooling
• 🚀 Quickstart Guide
  • 1. Running Policy Training (Notebook)
  • 2. Launching Trajectory Visualizer
• 📂 Repository Architecture
• 📚 References

---

🎯 Aim of the Project

The Blocks World is a foundational benchmark in AI planning and symbolic reasoning. In this domain, an agent must manipulate a set of labelled blocks located either on a table or stacked upon one another, transitioning from an arbitrary disordered configuration to a designated goal stack (e.g. $[E \rightarrow D \rightarrow C \rightarrow B \rightarrow A]$) using a single robotic manipulator.

While classical planners like $A^*$, Fast-Forward, or SAT-solvers solve this via explicit search over symbolic state graphs, training a deep reinforcement learning policy to discover valid plans introduces severe challenges:

1. Extreme Reward Sparsity: If only the final completed goal stack receives positive feedback, the probability of randomly stumbling into the exact sequence of 10–15 coordinated actions is negligible ($\sim 0%$).
2. Credit Assignment & Local Minima: Deconstructing an existing tower to reach an obscured block temporarily breaks partial subgoals, penalizing naive distance heuristics.
3. Invalid Transitions & Action Explosions: Out of all possible permutations of actions, only a subset are physically admissible at any state.

The primary aim of this project is to:

• Formulate a clean, vector-based STRIPS-style environment in pure PyTorch/Python.
• Implement and compare Reward Modeling paradigms: from naive sparse metrics to Potential-Based Reward Shaping (PBRS) that preserves policy invariance while steering search.
• Overcome policy entropy collapse through Batched Policy Gradients (REINFORCE) with entropy loss and gradient norm clipping.
• Provide a responsive, dark-mode Next.js + SVG visualization platform to scrub, inspect, diff, and diagnose recorded trajectory rollouts.

---

📈 Experimental Evolution & Findings (v1 vs. v2 vs. v3)

Side-by-Side Diagnostic Dashboards

The training dynamics across all three iterations are summarized below. (Click on any dashboard to view in full resolution):

Phase 1 (v1) — Baseline	Phase 2 (v2) — Subgoal PBRS	Phase 3 (v3) — Batched PG + Entropy
Solve Rate: 10% – 15% Mean Horizon: 50.0 steps (Max) Policy Entropy: Premature Collapse Primary Bottleneck: Local Predicate Loops	Solve Rate: ~60% Mean Horizon: 25 – 35 steps Policy Entropy: Moderate Stability Primary Advance: Sequential Subgoals	Solve Rate: 80% – 95% Mean Horizon: 10 – 15 steps (Optimal) Policy Entropy: Sustained Regularization Primary Advance: Batched PG + Entropy Loss

Unified Comparative Analysis

Analyzing the three iterations reveals how algorithmic adjustments resolved distinct failure modes:

1. The Solve Rate Breakthrough ($10% \rightarrow 60% \rightarrow 95%$)

• v1 Failure: In Phase 1, using naive step penalties with hardcoded predicate bonuses, the policy floundered between $0% - 8%$ during training and peaked at only $10-15%$ in greedy evaluation. The agent repeatedly fell into local minima (e.g. resting a block on the table and picking it back up).
• v2 Recovery: Introducing bottom-up Potential-Based Reward Shaping (PBRS) that only rewarded stacking if the foundation block onTable(E) was firmly grounded allowed the agent to discover multi-step subgoals, driving success up to $\sim 60%$ by Epoch 500.
• v3 Mastery: Combining PBRS with Batched Trajectory Rollouts eliminated high-variance gradient noise, driving success to an asymptotic $80% - 95%$ solve rate across evaluation folds.

2. Horizon Optimization ($50 \rightarrow 30 \rightarrow 12 \text{ steps}$)

• In v1, the trajectory length stayed pegged at the ceiling of $50.0$ steps (the timeout horizon), indicating the agent was timing out without ever satisfying the terminal goal condition.
• In v2, as the agent learned to respect bottom-up tower assembly, the mean steps to solve dropped to $25 - 35$ steps.
• In v3, trajectory lengths plummeted rapidly to $10 - 15$ steps, which matches the theoretical optimal plan length for randomly partitioned 5-block configurations.

3. Entropy Dynamics & Exploration Collapse

• Premature Collapse in v1: In Phase 1, policy entropy experienced a catastrophic collapse from $1.3$ down to $<0.15$ within the first 500 epochs. Because the agent lacked structured credit assignment, it converged prematurely onto sub-optimal deterministic loops before ever seeing positive goal trajectories.
• Controlled Decay in v2: Phase 2 maintained entropy above $0.2$ for longer, allowing sustained exploration of valid configurations.
• Entropy Regularization in v3: Adding an explicit entropy loss penalty ($\beta \mathcal{H}(\pi)$) and gradient norm clipping ($|\mathbf{g}|_2 \le 1.0$) prevented policy overconfidence. The agent maintained sufficient stochastic exploration early on, only consolidating into deterministic choices once valid planning paths were reinforced.

Progression Matrix

Diagnostic Metric	Iteration 1 (v1 Baseline)	Iteration 2 (v2 Subgoal PBRS)	Iteration 3 (v3 Batched + Entropy)
Peak Solve Rate	10% – 15%	~60%	80% – 95%
Mean Horizon	50.0 steps (Timeout)	25 – 35 steps	10 – 15 steps (Optimal)
Entropy Profile	Premature Collapse ($<0.15$)	Moderate Decay	Controlled Regularization
Gradient Updates	Single-trajectory REINFORCE	Single-trajectory REINFORCE	Batched Rollouts ($B=16$)
Gradient Clipping	None	None	$|\mathbf{g}|_2 \le 1.0$
Reward Formulation	Naive step penalty + goal bonus	Bottom-up Potential Function $\Phi(s)$	PBRS + Step Cost + Goal Bonus
Core Resolution	❌ Trapped in cyclic loops	⚠️ High gradient variance	✅ Stable, near-optimal solver

---

🧩 Problem Formulation & Environment Design

For a set of 5 distinct blocks $B = {A, B, C, D, E}$ and a single manipulator arm, the system state is represented via a grounded first-order predicate vector.

STRIPS First-Order Predicates

$$\text{Total Predicates} = 20 + 5 + 5 + 1 + 5 = \mathbf{36}$$

Predicate	Arity / Permutations	Dimension	Semantic Meaning
on(x, y)	$5 \times 4$ permutations ($x \neq y$)	20	Block x is stacked directly on top of block y
onTable(x)	$5$ blocks ($x \in B$)	5	Block x rests directly on the workspace surface
clear(x)	$5$ blocks ($x \in B$)	5	Block x has no other block resting on top of it
armempty	Boolean state	1	The robotic manipulator is unoccupied
holding(x)	$5$ blocks ($x \in B$)	5	The robotic manipulator is gripping block x

Action Space & State Transitions

$$\text{Total Grounded Actions} = 20 + 20 + 5 + 5 = \mathbf{50}$$

Each discrete action adheres to formal STRIPS preconditions, add-effects, and delete-effects:

1. unstack(x, y):
   • Preconditions: armempty ∧ clear(x) ∧ on(x, y)
   • Add-Effects:   holding(x), clear(y)
   • Del-Effects:   armempty, on(x, y), clear(x)

2. stack(x, y):
   • Preconditions: holding(x) ∧ clear(y)
   • Add-Effects:   on(x, y), armempty, clear(x)
   • Del-Effects:   holding(x), clear(y)

3. pickup(x):
   • Preconditions: armempty ∧ clear(x) ∧ onTable(x)
   • Add-Effects:   holding(x)
   • Del-Effects:   armempty, onTable(x), clear(x)

4. putdown(x):
   • Preconditions: holding(x)
   • Add-Effects:   onTable(x), armempty, clear(x)
   • Del-Effects:   holding(x)

Policy Network & Action Masking

The policy model is a multi-layer perceptron (MLP) taking a 72-dimensional vector (concatenation of current state $\mathbf{s}_t \in {0, 1}^{36}$ and goal state $\mathbf{g} \in {0, 1}^{36}$):

$$\mathbf{x} = [\mathbf{s}_t ,|, \mathbf{g}] \in \mathbb{R}^{72}$$

$$\mathbf{h}_1 = \text{ReLU}(\mathbf{W}_1 \mathbf{x} + \mathbf{b}_1), \quad \mathbf{h}_2 = \text{ReLU}(\mathbf{W}_2 \mathbf{h}_1 + \mathbf{b}_2), \quad \mathbf{h}_3 = \text{ReLU}(\mathbf{W}_3 \mathbf{h}_2 + \mathbf{b}_3)$$

$$\mathbf{z} = \mathbf{W}_4 \mathbf{h}_3 + \mathbf{b}_4 \in \mathbb{R}^{50}$$

Environmental Action Masking

To prevent the agent from attempting physically invalid operations (e.g. unstacking a block buried beneath others or gripping while already holding), an eligibility mask $\mathbf{m}_t \in {0, 1}^{50}$ is computed dynamically from the environment preconditions. Invalid action logits are clamped to $-\infty$ prior to categorical sampling:

$$\tilde{z}_i = \begin{cases} z_i & \text{if } m_{t, i} = 1 \ -10^9 & \text{if } m_{t, i} = 0 \end{cases}, \qquad \pi(a_i \mid \mathbf{s}_t) = \frac{\exp(\tilde{z}_i)}{\sum_j \exp(\tilde{z}_j)}$$

---

⚖️ Reward Engineering & Potential Shaping

Early iterations demonstrated that unshaped step penalties with terminal bonuses fail completely due to the credit assignment bottleneck. The agent adopts cyclical behaviors (pickup $\leftrightarrow$ putdown) that minimize immediate penalties without risking exploration.

To solve this, we formulated Potential-Based Reward Shaping (PBRS) based on grounded pile evaluation:

$$R(s_t, a_t, s_{t+1}) = \gamma \Phi(s_{t+1}) - \Phi(s_t) - c_{\text{step}} + \mathbb{I}_{\text{done}} \cdot R_{\text{goal}}$$

where $\Phi(s)$ is a state potential score measuring structural progress toward the goal stack $[E, D, C, B, A]$:

1. Manipulator Arm Condition: $+2$ if the arm state aligns with goal arm configuration (armempty or holding(x)).
2. Bottom-Up Foundation Check:
   • Check if the foundation base onTable(E) exists. If missing, foundation is invalid (subsequent stacked blocks grant no points).
   • If base is valid, sequentially verify each ordered link: on(D, E), on(C, D), on(B, C), on(A, B). Each consecutive intact link adds $+2$.
   • If the stack is unbroken to the top, check if clear(A) holds ($+2$).
3. Predicate Overlap Matching: $+0.1$ for each matching active predicate with the target goal state.
4. Step Penalty & Goal Bonus: $c_{\text{step}} = 0.05$ per action taken; $R_{\text{goal}} = 50.0$ upon reaching terminal state.

---

🕹️ Interactive Web Visualizer

The frontend is a bespoke web interface located in frontend/ built with Next.js 16, React 19, and Tailwind CSS v4 to analyze recorded trajectory rollouts.

Visualizer Highlights

• 🏗️ Overhead Robotic Gantry & Claw Gripper:
  • A responsive overhead trolley rail dynamically travels horizontally to the targeted block column.
  • Dual articulated mechanical claws spread open when armempty and firmly grip when holding(X).
• 🏛️ Column Stability Engine:
  • Resolves multi-block towers from dynamic on(X, Y) and onTable(X) relationships.
  • Column stability algorithms ensure block stacks retain consistent horizontal coordinates across timesteps, eliminating visual hopping.
• 🔍 Predicate Transition Diffing:
  • Live inspector highlights newly added STRIPS predicates in green (+) and deleted predicates in red (-) at each timestep.
  • Interactive block hover illuminates corresponding active predicates across the canvas.
• ⏱️ Timeline Scrubber & Playback Controls:
  • Full-range interactive slider, playback rates ($0.25\times$ to $4\times$), auto-loop toggle, and jumpable step markers.
• 💾 In-Memory Streaming & Offline File Support:
  • Parses diagnostics files with $10,000+$ trajectories using browser IndexedDB and streaming disk cache.

Keyboard Shortcuts

Shortcut	Description
←	Step backward one timestep
→	Step forward one timestep
Space	Toggle auto-playback Play / Pause
Home	Jump to initial step ($t=0$)
End	Jump to terminal step ($t=T-1$)
L	Toggle continuous loop mode
?	Open keyboard shortcuts reference modal
Esc	Dismiss active modal or overlay

---

💻 Technical Tooling

---

🚀 Quickstart Guide

1. Running Policy Training (Notebook)

Clone the repository and open the self-contained training notebook:

git clone https://github.com/slyofzero/BlocksWorldSolver.git
cd BlocksWorldSolver

Open BlocksWorld.ipynb in Jupyter Lab, VS Code, or Google Colab:

• Run cells 1 through 10 to initialize the STRIPS environment, action pre-conditions, and RewardModel.
• Execute train_policy() in cell 13 or run save_diagnostics(5) in cell 16 to generate multi-run diagnostics.
• Trajectories and metric histories are saved automatically to training_runs/train_*/diagnostics.json.

2. Launching Trajectory Visualizer

To explore trajectories with the dark-mode Next.js visualizer:

cd frontend

# Install dependencies using Bun (or npm)
bun install

# Start development server
bun dev

Visit http://localhost:3000 in your browser. The visualizer will automatically detect and load the pre-recorded runs from training_runs/ and frontend/data/runs/demo_run.json.

---

📂 Repository Architecture

BlocksWorldSolver/
├── assets/
│   ├── demo.gif                 # Interactive visualizer screen recording
│   ├── v1_dashboard.png         # Iteration 1 training curves & failure diagnosis
│   ├── v2_dashboard.png         # Iteration 2 subgoal reward shaping metrics
│   └── v3_dashboard.png         # Iteration 3 batched REINFORCE + entropy metrics
│
├── frontend/                    # Next.js 16 / React 19 Trajectory Visualizer
│   ├── app/
│   │   ├── components/
│   │   │   ├── BlocksWorldCanvas.tsx      # SVG canvas with overhead gantry & claw
│   │   │   ├── InspectorPanel.tsx         # Real-time STRIPS predicate diffs
│   │   │   ├── Controls.tsx               # Playback speeds, loop & scrub buttons
│   │   │   └── EpisodeBrowserModal.tsx    # Filter & select episodes by reward
│   │   ├── lib/
│   │   │   ├── layoutEngine.ts            # Column stability and stack layout algorithm
│   │   │   ├── predicateParser.ts         # STRIPS logic & syntax evaluator
│   │   │   └── runService.ts              # Trajectory streaming and caching
│   │   └── page.tsx                       # Main dashboard layout
│   └── package.json
│
├── training_runs/
│   └── train_0/
│       ├── diagnostics.json     # Full episodic metrics (returns, loss, entropy)
│       └── diagnostics.png      # 8-metric matplotlib summary chart
│
├── BlocksWorld.ipynb            # End-to-end environment, policy, and training code
├── trajectories.json            # Sample evaluation trajectory rollouts
└── README.md                    # Project documentation

---

📚 References

• Sutton & Barto — Reinforcement Learning: An Introduction (Policy Gradient Theorem & REINFORCE).
• Ng, Harada, & Russell (1999) — Policy Invariance Under Reward Transformations: Theory and Application to Reward Shaping (Potential-Based Reward Shaping).
• Fikes & Nilsson (1971) — STRIPS: A New Approach to the Application of Theorem Proving to Problem Solving.
• Mnih et al. (2016) — Asynchronous Methods for Deep Reinforcement Learning (Entropy Regularization).

---

"Curious about how machines learn structure, not just patterns."