Lab 01 — MDP + Q-Learning, Task Planner & VLA Action Decoder

Phase: 15 — Embodied AI & Robotics: RL, Planning & VLA Difficulty: ⭐⭐⭐☆☆ (the algorithms are small; the convergence intuition is ⭐⭐⭐⭐⭐) Time: 4–6 hours

An embodied agent is not a chatbot: it acts in a world, the world pushes back, and there is no "regenerate." This lab builds the control core of one from scratch in stdlib Python — a deterministic MDP (GridWorld), the planning baseline you use when you have the model (value_iteration), the learning algorithm for when you don't (q_learning with seeded epsilon-greedy), a symbolic task planner (STRIPS/BFS over operators with preconditions and effects), and a VLA action decoder that models RT-2/OpenVLA's "actions as tokens": a tokenized (instruction + perception) input is decoded by an injected policy into discrete action tokens, then de-tokenized into continuous actions. The two soul tests prove the learner is right: the Q-policy reaches the goal in the optimal number of steps, and every action it takes matches value iteration's optimum.

What you build

• GridWorld — a deterministic MDP. States are cells; actions are up/down/left/right; step(state, action) -> (next_state, reward, done) is a pure transition (walls and edges keep you in place, the goal gives +goal_reward and terminates, every other step costs step_penalty).
• value_iteration(env, gamma, theta) — the planning baseline (dynamic programming). Iterate the Bellman optimality backup V(s) = max_a [r + gamma·V(s')] to convergence; return (V, policy).
• q_learning(env, episodes, alpha, gamma, epsilon, rng) — tabular Q-learning. The update is Q(s,a) += alpha·(r + gamma·max_a' Q(s',a') − Q(s,a)); behaviour is seeded epsilon-greedy.
• greedy_policy(Q) and rollout(env, policy, max_steps) — turn a value table into behaviour and run it.
• Operator + plan(start, goal, operators) — a STRIPS/BFS symbolic planner: shortest sequence of operator names (each with preconditions/add/del effects) whose execution satisfies the goal, or None if unreachable. execute_plan runs a plan and rejects an invalid one.
• ActionTokenizer + vla_decode(...) — discretize a continuous action into bins ↔ tokens, then decode a tokenized input into a sequence of action tokens via an injected deterministic policy and de-tokenize back to actions (RT-2/OpenVLA "actions as tokens").

Key concepts

Files

Run

pip install -r requirements.txt
pytest test_lab.py -v                       # against your lab.py (red until implemented)
LAB_MODULE=solution pytest test_lab.py -v   # against the reference (must be green)
python solution.py                          # the worked example

Success criteria

• All tests pass against your implementation (and LAB_MODULE=solution).
• You can explain the Bellman optimality backup and why iterating it converges to V*.
• You can explain why test_q_learning_converges_to_optimal_open_grid is the RL soul test: a table of zeros, fed only (s, a, r, s') samples, becomes a policy that reaches the goal in the fewest possible steps.
• You can explain why test_q_learning_matches_value_iteration_policy checks action value, not action identity — because optimal policies are not unique when actions tie on value.
• You can explain why ε-greedy must be seeded for the determinism test to pass, and why Q-learning still learns the optimum despite acting non-greedily (off-policy).
• You can trace plan(...) returning the shortest valid operator sequence and None when the goal is unreachable, and vla_decode(...) producing the scripted action sequence deterministically.

How this maps to the real stack

Limits of the miniature (be honest in the interview): the world is deterministic and fully observed — real robotics is stochastic and partially observable (POMDP), where you need belief states, function approximation, and on-policy methods like PPO; the action space is one discrete scalar, not a 7-DoF continuous arm; the planner searches a tiny symbolic state space (real PDDL is PSPACE-hard and uses heuristics); and the VLA "policy" is a scripted stub, not a billion-parameter VLM — the point is the interface (tokenize → decode → de-tokenize), not the weights.

Extensions (build these on real hardware / libraries)

• Add stochastic transitions (slip probability) so step returns a distribution; switch the Bellman backup to the expectation and watch the optimal policy become risk-averse.
• Replace the Q-table with a tiny linear function approximator over (row, col) features and re-derive the update — the bridge to DQN.
• Add SARSA (on-policy: target uses the actual next action, not the max) and compare the policies near the cliff edge; this is the classic SARSA-vs-Q-learning lesson.
• Wrap the env in Gymnasium and train a real DQN/PPO from stable-baselines3; confirm it recovers the same optimal path.
• Extend the planner to PDDL with parameterized operators and run Fast Downward; then have an LLM propose the operators and use execute_plan as the affordance checker (SayCan-style).
• Fine-tune a small OpenVLA checkpoint on a binned action space and compare its decode loop to vla_decode.

Interview / resume

• Talking points: "What's the difference between planning and learning, and when do you reach for each?" "Why is Q-learning off-policy and DQN possible?" "How do RT-2/OpenVLA put robot actions in an LLM's vocabulary?" "Why is latency a safety property in embodied AI?" "How do you ground an LLM planner so it doesn't propose impossible actions?"
• Resume bullet: Implemented the control core of an embodied agent from scratch — a GridWorld MDP, value iteration, tabular Q-learning with seeded ε-greedy (converging to the value-iteration optimum), a STRIPS/BFS task planner, and an RT-2/OpenVLA-style VLA action decoder that treats discretized actions as LLM tokens — verified by a deterministic test suite including RL convergence and planner-equivalence soul tests.