Add detailed documentation for transition tables

Explains array shapes, indexing, and how to use the tables
in the Q-iteration algorithm with concrete examples.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Author: eugenevinitsky

homeworks/homework_1/problem_1/problem_1.py

@@ -17,17 +17,63 @@
 WARNING: Do NOT modify or regenerate the transition tables (.npy files).
 The grading server uses the same tables - changing them will cause your submission to fail.
 
+TRANSITION TABLES:
+=================
+
+The state space is discretized into a 200x200 grid:
+    - s0 (first index): Position index, ranging from 0 to 199
+        - Index 0 corresponds to position -1.2 (leftmost)
+        - Index 199 corresponds to position 0.6 (rightmost)
+        - The goal region (position >= 0.5) is roughly indices 189-199
+    - s1 (second index): Velocity index, ranging from 0 to 199
+        - Index 0 corresponds to velocity -0.07 (moving left fastest)
+        - Index 199 corresponds to velocity 0.07 (moving right fastest)
+        - Index ~100 corresponds to velocity ~0 (stationary)
+
+The transition tables are numpy arrays:
+
     P (transition_next_states.npy): shape (200, 200, 3, 2)
-        P[s0, s1, a] = [s0', s1']
-        The next state indices when taking action a from state (s0, s1)
+        P[s0, s1, a] = [s0', s1']  (a numpy array of 2 integers)
+
+        Given current state indices (s0, s1) and action a, returns the
+        next state indices [s0', s1'].
+
+        Example:
+            next_state = P[100, 100, 2]  # State (100,100), action 2 (push right)
+            s0_next, s1_next = next_state[0], next_state[1]
+            # Now you can look up Q[s0_next, s1_next, :] to get Q-values at next state
 
     R (transition_rewards.npy): shape (200, 200, 3)
-        R[s0, s1, a] = reward
-        The reward is -1 for each step until the goal is reached.
+        R[s0, s1, a] = reward (a single float)
+
+        The immediate reward for taking action a in state (s0, s1).
+        In MountainCar, this is always -1 for every step (encourages reaching
+        the goal quickly).
+
+        Example:
+            reward = R[100, 100, 2]  # Always -1.0
 
     D (transition_dones.npy): shape (200, 200, 3)
-        D[s0, s1, a] = done (True/False)
-        True when the car reaches the goal (position >= 0.5)
+        D[s0, s1, a] = done (a boolean: True or False)
+
+        Whether the episode terminates after taking action a in state (s0, s1).
+        True only when the car reaches the goal (position >= 0.5).
+
+        Example:
+            done = D[195, 150, 2]  # True if this transition reaches the goal
+            # If done is True, there is no future reward (episode ends)
+
+USING THE TABLES IN Q-ITERATION:
+===============================
+
+For each state-action pair (s0, s1, a):
+    1. Look up the next state: s0', s1' = P[s0, s1, a]
+    2. Look up the reward: r = R[s0, s1, a]
+    3. Look up if terminal: d = D[s0, s1, a]
+    4. Apply Bellman update:
+       Q_new[s0, s1, a] = r + gamma * (1 - d) * max_a' Q[s0', s1', a']
+
+       Note: (1 - d) ensures we don't add future value for terminal states
 
 Your task:
     Implement Q-iteration using the Bellman optimality equation:
@@ -38,7 +84,7 @@
 
 Submission:
     - checkpoint.pt: Your Q-table saved with torch.save (shape: 200x200x3)
-    - policy.py is provided - do not modify it
+    - policy.py: Implement the forward() method to select actions using your Q-table
 
 Expected performance:
     A well-implemented solution should consistently reach the goal.