qlearning.py

"""
qlearning.py — Tabular Q-Learning for the Garbage Collecting Robot.

Training runs directly against GarbageRobotEnv (no HTTP server needed).
The Q-table is persisted to disk as JSON and loaded by inference.py at startup.

State representation:
    (robot_x, robot_y, sorted_garbage_tuple)
    e.g. (2, 3, ((1,1),(4,4)))   — compact, hashable, fully describes the relevant world

Actions:
    0=UP  1=DOWN  2=LEFT  3=RIGHT  4=COLLECT

Usage:
    # Train all tasks and save
    python3 qlearning.py --train --episodes 8000

    # Evaluate silently (uses saved Q-table)
    python3 qlearning.py --eval

Fix applied:
    - load() previously had two separate key-reconstruction passes, where the
      first pass result (variable `k`) was computed but then immediately discarded.
      The second pass also misidentified the garbage sub-list when it had exactly
      2 integer elements (treating [gx, gy] pairs as flat coords instead of a
      tuple-of-tuples). Replaced both passes with a single, unambiguous decode:
        parsed = [rx, ry, [[gx1,gy1],[gx2,gy2],...]]
      where the third element is always the nested garbage list.
"""

import os
import json
import random
import argparse
from collections import defaultdict
from environment import GarbageRobotEnv
from scenarios import SCENARIOS

# ── Constants ──────────────────────────────────────────────────────────────

ACTIONS      = ["UP", "DOWN", "LEFT", "RIGHT", "COLLECT"]
ACTION_IDX   = {a: i for i, a in enumerate(ACTIONS)}
Q_TABLE_PATH = os.environ.get("Q_TABLE_PATH", "qtable.json")

# ── Hyperparameters ─────────────────────────────────────────────────────────

ALPHA         = 0.15
GAMMA         = 0.97
EPSILON_START = 1.0
EPSILON_END   = 0.05
EPSILON_DECAY = 0.9995


# ── State Encoding ──────────────────────────────────────────────────────────

def encode_state(obs: dict) -> tuple:
    """
    Convert a raw observation dict into a hashable tuple suitable as a Q-table key.

    Key structure: (robot_x, robot_y, ((gx1,gy1),(gx2,gy2),...))
    Garbage positions are sorted so order doesn't create phantom new states.
    """
    rx, ry = obs["robot_position"]
    garbage = tuple(sorted((int(g[0]), int(g[1])) for g in obs["garbage_positions"]))
    return (int(rx), int(ry), garbage)


# ── Q-Table ─────────────────────────────────────────────────────────────────

class QTable:
    """
    Dictionary-backed Q-table with defaultdict initialisation.
    Values default to a small optimistic initial value to encourage exploration.
    """

    def __init__(self, optimistic_init: float = 0.5):
        self.optimistic_init = optimistic_init
        self._q: dict = {}

    def _ensure(self, state: tuple):
        if state not in self._q:
            self._q[state] = [self.optimistic_init] * len(ACTIONS)

    def get(self, state: tuple, action_idx: int) -> float:
        self._ensure(state)
        return self._q[state][action_idx]

    def update(self, state: tuple, action_idx: int, value: float):
        self._ensure(state)
        self._q[state][action_idx] = value

    def best_action(self, state: tuple) -> int:
        """Return the index of the greedy best action."""
        self._ensure(state)
        return int(max(range(len(ACTIONS)), key=lambda i: self._q[state][i]))

    def best_q(self, state: tuple) -> float:
        self._ensure(state)
        return max(self._q[state])

    # ── Persistence ─────────────────────────────────────────────────────────

    def save(self, path: str = Q_TABLE_PATH):
        """
        Serialise Q-table to JSON.

        Key format saved to disk:
            [rx, ry, [[gx1,gy1], [gx2,gy2], ...]]
        This is unambiguous: element 0 and 1 are ints, element 2 is always a
        list-of-lists, even when there is only one garbage piece.
        """
        serialisable = {}
        for (rx, ry, garbage), v in self._q.items():
            key = json.dumps([rx, ry, [list(g) for g in garbage]])
            serialisable[key] = v
        with open(path, "w") as f:
            json.dump(serialisable, f)
        print(f"[Q-Table] Saved {len(self._q):,} states → {path}")

    def load(self, path: str = Q_TABLE_PATH) -> bool:
        """
        Load Q-table from JSON.

        FIX: The previous implementation had two redundant key-reconstruction
        loops. The first built variable `k` which was immediately discarded;
        the second pass misclassified [gx, gy] pairs (lists of 2 ints) as flat
        coordinates rather than garbage-position tuples, corrupting multi-garbage
        states.

        New single-pass decode relies on the unambiguous 3-element structure:
            parsed[0] = rx  (int)
            parsed[1] = ry  (int)
            parsed[2] = [[gx1,gy1], ...]  (always a list-of-lists)
        """
        if not os.path.exists(path):
            return False
        with open(path, "r") as f:
            raw = json.load(f)
        self._q = {}
        for k_str, v in raw.items():
            parsed = json.loads(k_str)
            # Robustly handle both new format [rx, ry, [[gx,gy],...]]
            # and old format [rx, ry, [gx, gy]] (single garbage, flat list).
            rx, ry = int(parsed[0]), int(parsed[1])
            raw_garbage = parsed[2]
            if raw_garbage and isinstance(raw_garbage[0], list):
                # New / multi-garbage format: [[gx1,gy1],[gx2,gy2],...]
                garbage = tuple(tuple(p) for p in raw_garbage)
            elif raw_garbage and isinstance(raw_garbage[0], int):
                # Old single-garbage flat format: [gx, gy]
                garbage = (tuple(raw_garbage),)
            else:
                garbage = ()
            self._q[(rx, ry, garbage)] = v
        print(f"[Q-Table] Loaded {len(self._q):,} states ← {path}")
        return True

    def __len__(self):
        return len(self._q)


# ── Observation Helper ───────────────────────────────────────────────────────

def _obs_from_env(env) -> dict:
    """Build an obs dict directly from GarbageRobotEnv fields."""
    obs_obj = env.get_observation()
    return {
        "robot_position":     obs_obj.robot_position,
        "garbage_positions":  list(obs_obj.garbage_positions),
        "obstacle_positions": list(obs_obj.obstacle_positions),
        "grid_size":          obs_obj.grid_size,
        "battery_level":      obs_obj.battery_level,
        "inventory_count":    obs_obj.inventory_count,
        "message":            obs_obj.message,
        "robot_mode":         obs_obj.robot_mode,
        "home_position":      obs_obj.home_position,
        "unload_station":     obs_obj.unload_station,
        "current_storage_load": obs_obj.current_storage_load,
        "storage_capacity":   obs_obj.storage_capacity,
        "distance_from_home": obs_obj.distance_from_home,
    }


# ── Training ─────────────────────────────────────────────────────────────────

def train(
    task_ids=None,
    episodes: int = 8000,
    qtable: QTable = None,
    verbose: bool = True,
) -> QTable:
    """
    Run Q-learning over the given task_ids for `episodes` total episodes.
    Tasks are sampled uniformly so the agent generalises across difficulties.
    """
    if task_ids is None:
        task_ids = list(SCENARIOS.keys())

    if qtable is None:
        qtable = QTable()

    env     = GarbageRobotEnv()
    epsilon = EPSILON_START

    best_scores: dict = {t: 0.0 for t in task_ids}

    for ep in range(1, episodes + 1):
        task_id = random.choice(task_ids)
        env.reset(task_id)
        obs       = _obs_from_env(env)
        state     = encode_state(obs)

        total_reward = 0.0
        done         = False

        while not done:
            if random.random() < epsilon:
                action_idx = random.randrange(len(ACTIONS))
            else:
                action_idx = qtable.best_action(state)

            action     = ACTIONS[action_idx]
            result     = env.step(action)
            next_obs   = result["observation"]
            reward     = result["reward"]
            done       = result["done"]

            next_state = encode_state(next_obs)

            # Bellman update
            old_q     = qtable.get(state, action_idx)
            td_target = reward + (0.0 if done else GAMMA * qtable.best_q(next_state))
            new_q     = old_q + ALPHA * (td_target - old_q)
            qtable.update(state, action_idx, new_q)

            state = next_state
            obs   = next_obs
            total_reward += reward

        score = env.grade(task_id)
        if score > best_scores[task_id]:
            best_scores[task_id] = score

        epsilon = max(EPSILON_END, epsilon * EPSILON_DECAY)

        if verbose and ep % 500 == 0:
            avg_best = sum(best_scores.values()) / len(best_scores)
            print(
                f"  Ep {ep:5d}/{episodes}  ε={epsilon:.4f}  "
                f"states={len(qtable):,}  "
                f"best_scores={best_scores}  avg={avg_best:.2f}"
            )

    return qtable


# ── Inference Helper (used by inference.py) ──────────────────────────────────

class QLearningAgent:
    """
    Thin wrapper around a loaded Q-table for use by inference.py.
    Falls through (returns None) when the state has never been seen during training.
    """

    def __init__(self, path: str = Q_TABLE_PATH):
        self.qtable = QTable()
        self.loaded = self.qtable.load(path)

    def get_action(self, obs: dict) -> str | None:
        if not self.loaded:
            return None
        state = encode_state(obs)
        if state not in self.qtable._q:
            return None
        return ACTIONS[self.qtable.best_action(state)]


# ── Evaluation ───────────────────────────────────────────────────────────────

def evaluate(qtable: QTable, task_ids=None, runs: int = 5) -> dict:
    """Run `runs` greedy episodes per task and return average scores."""
    if task_ids is None:
        task_ids = list(SCENARIOS.keys())

    env     = GarbageRobotEnv()
    results = {}

    for task_id in task_ids:
        scores = []
        for _ in range(runs):
            env.reset(task_id)
            obs  = _obs_from_env(env)
            done = False
            while not done:
                state      = encode_state(obs)
                action_idx = qtable.best_action(state)
                result     = env.step(ACTIONS[action_idx])
                obs        = result["observation"]
                done       = result["done"]
            scores.append(env.grade(task_id))
        avg = sum(scores) / len(scores)
        results[task_id] = round(avg, 3)
        print(f"  {task_id:12s}  avg score = {avg:.3f}  ({scores})")

    return results


# ── CLI Entry Point ───────────────────────────────────────────────────────────

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Q-Learning for Garbage Robot")
    parser.add_argument("--train",    action="store_true", help="Run training")
    parser.add_argument("--eval",     action="store_true", help="Run evaluation only")
    parser.add_argument("--episodes", type=int, default=8000)
    parser.add_argument("--tasks",    nargs="+", default=None)
    parser.add_argument("--output",   default=Q_TABLE_PATH)
    args = parser.parse_args()

    if args.train:
        print("=" * 55)
        print("  Q-Learning Training — Garbage Collecting Robot")
        print("=" * 55)
        task_ids = args.tasks or list(SCENARIOS.keys())
        print(f"  Tasks    : {task_ids}")
        print(f"  Episodes : {args.episodes}")
        print(f"  α={ALPHA}  γ={GAMMA}  ε {EPSILON_START}→{EPSILON_END}  decay={EPSILON_DECAY}")
        print()

        qt = train(task_ids=task_ids, episodes=args.episodes, verbose=True)
        qt.save(args.output)

        print("\n  — Evaluation on greedy policy —")
        evaluate(qt, task_ids)

    elif args.eval:
        print("=" * 55)
        print("  Q-Learning Evaluation")
        print("=" * 55)
        qt = QTable()
        if not qt.load(args.output):
            print(f"[ERROR] No Q-table found at {args.output}. Run with --train first.")
        else:
            evaluate(qt)
    else:
        parser.print_help()