diff --git a/.docs/README.md b/.docs/README.md new file mode 100644 index 0000000..9853715 --- /dev/null +++ b/.docs/README.md @@ -0,0 +1,67 @@ +# Research notes + +Working notes for the learning-for-planning line in `jupyddl`. These are +research documents, not user documentation — the user-facing description of +`jupyddl learn` is in the README and the module docstrings. + +| Note | What it covers | +|---|---| +| [learned-heuristics.md](learned-heuristics.md) | The imitation stage: prior work, the design we chose, measured results on three domains, and an analysis of the one domain where it loses badly | +| [rl-for-search.md](rl-for-search.md) | The reinforcement stage: the MDP that "minimise expansions" corresponds to, why the obvious policy gradient is hard here, and what we do instead | +| [roadmap.md](roadmap.md) | What to build next, ordered by expected value, with the experiment that would settle each | + +## The one-paragraph version + +A planner's heuristic is a learned function waiting to happen: every solved +instance is a labelled trajectory, since the cost of a plan's suffix from any +state on it is that state's cost-to-go. Fitting a network to those labels is +*imitation*, and it works — on blocksworld, training on 3–6 block instances +produces a heuristic that beats `hff` on 9–13 block instances by 4× in +expansions and 15× in wall-clock. But imitation optimises a proxy. What we +actually want is the heuristic that makes search expand the fewest nodes, and +that quantity is not a differentiable function of the weights. Optimising it +directly — treating the planner as a black box and the weights as a policy — +is where this becomes reinforcement learning, and on blocksworld it is worth +another 2.7× on top of imitation while fixing a coverage failure that +imitation alone left behind — it solved a held-out instance imitation could not +solve at all. + +It also produced the most useful negative result here, which is about us rather +than about planning: an earlier version of these notes credited that improvement +to the validation split. Re-running with one variable at a time showed the +perturbation scale was doing all the work, and that nine of the ten held-out +instances improve either way — the whole headline gap is a single hard instance. +Both corrections are written up in `rl-for-search.md` rather than quietly +edited out. + +## The video + +[`promo/jupyddl-rl.mp4`](../promo/jupyddl-rl.mp4) is a 97-second tour of this +work. Like the main promo it measures everything at render time — it trains, +reinforces, and re-runs both failure modes — so it cannot drift from these +notes. Rebuild it with: + +```bash +python tools/make_learn_promo.py --cache promo/rl-data.json -o promo/jupyddl-rl.mp4 +``` + +`promo/rl-data.json` is the cached measurement pass; delete it to re-measure. + +## Reproducing everything here + +```bash +pip install -e ".[dev,learn]" + +# the headline blocksworld result, about a minute +jupyddl learn blocksworld --sizes 3-6 --seeds-per-size 3 \ + --cem 10 --cem-sizes 9-12 --evaluate 9-13 -o blocksworld.heur.json + +# then use it anywhere a heuristic name is accepted +jupyddl solve domain.pddl problem.pddl -s gbfs -H learned:blocksworld.heur.json +``` + +Every number in these notes came from a run on this repository at the commit +that introduced them, on one CPU, with `numpy` installed. They are single-seed +measurements on generated instances — enough to support the qualitative claims +made here and **not** enough for a paper. See the roadmap for what a publishable +evaluation would need. diff --git a/.docs/learned-heuristics.md b/.docs/learned-heuristics.md new file mode 100644 index 0000000..e1dc85f --- /dev/null +++ b/.docs/learned-heuristics.md @@ -0,0 +1,246 @@ +# Learning a heuristic from solved plans + +## The idea, and why it is not new + +A heuristic `h(s)` estimates the cost from `s` to a goal. A plan +`s0 →a1 s1 →a2 … →an sn` hands you `n+1` labelled examples for free: the cost of +the suffix from `si` bounds `h*(si)` above, and equals it when the plan is +optimal. Every planning run you have ever done was quietly generating training +data. + +The idea goes back at least to Ernandes and Gori (2004), who trained neural +networks as "sub-symbolic" heuristics for search, and to Yoon, Fern and Givan +(2008), who learned control knowledge for forward search from solved instances. +The modern line splits along one axis: **what does the network see?** + +- **Hand-designed features.** Cheap, fixed-length, transfers across instance + sizes by construction. Loses relational structure. +- **Learned relational representations.** ASNets (Toyer et al., AAAI 2018) build + a network whose structure mirrors the domain's action schemas. STRIPS-HGN + (Shen, Trevizan and Thiébaux, ICAPS 2020) runs a hypergraph network over the + delete-relaxation hypergraph, learning a domain-independent heuristic from + scratch. GOOSE (Chen, Trevizan and Thiébaux, ICAPS 2024) makes the striking + point that classical ML over Weisfeiler-Leman graph features beats the deep + models it was benchmarked against, at a fraction of the cost. + +`jupyddl` sits deliberately at the cheap end, for a reason worth stating +plainly: **a heuristic that is slower to evaluate than `hff` has to be much +better informed to win on time**, and in a pure-Python planner it will not be. +`lmcut` on a 14-block instance expands 203 nodes in 11.3 seconds; `hff` expands +577 in 1.6; the learned heuristic here expands 819 in 0.32. The *worst* +heuristic of the three by expansions wins by a factor of 5 on the clock. Any +design that costs a graph-network forward pass per generated state has to +overcome that, and nothing about running it in CPython helps. + +## What we built + +### Features (`jupyddl/learn/features.py`) + +Everything is keyed on the **predicate symbol**, never the ground atom, and +normalised by how many atoms of that symbol the task contains. For each symbol +`p` in a fixed vocabulary: + +- the fraction of `p`-atoms true in the state, +- the fraction of `p`-goals still unsatisfied, + +plus four global features (unsatisfied-goal fraction, its log, state size as a +fraction of all facts, and whether any goal is satisfied at all). + +That is `2P + 4` numbers, independent of object count. A model trained on four +blocks can be *evaluated* on forty — which a one-hot over ground atoms makes +impossible, since the vector length and the meaning of every slot change with +the instance. + +Cost is `O(|s| + |goals|)`, no successor generation. That matters: it is +evaluated once per generated state. + +### The network (`jupyddl/learn/model.py`) + +A small ReLU MLP, default `(features → 32 → 16 → 1)`, roughly 1000 parameters. +Output through softplus so the heuristic is never negative — a negative +cost-to-go does not merely mislead a planner, it breaks the assumptions of +every planner consuming it. + +Forward and backward passes are written twice: a plain-Python reference and a +NumPy batched version. The core of `jupyddl` has no runtime dependencies and +this keeps that promise; NumPy is a speed option worth one to two orders of +magnitude. `tests/test_learn.py` asserts the two produce identical gradients, +and checks both against finite differences. + +### The objective (`jupyddl/learn/train.py`) + +This is the part worth arguing about. + +The obvious objective is regression: fit `h(s) ≈ h*(s)`, minimise MSE. But +**greedy best-first search never reads a heuristic value.** It reads the *order* +the heuristic imposes on the open list. A model that is uniformly 30 too high +scores terribly on RMSE and guides perfectly. A model with excellent RMSE that +inverts two siblings sends the search down the wrong subtree. + +So the default objective here is a **ranking loss**, following Chrestien et al. +(NeurIPS 2023), whose title says it outright: *Optimize Planning Heuristics to +Rank, not to Estimate, Cost-to-Goal*. The idea is older than the deep-learning +framing — Xu, Fern and Yoon learned linear ranking functions for beam search a +decade and a half earlier (IJCAI 2007, JMLR 2009). + +Concretely: at each state on a training plan, take the successor the plan +actually took plus the siblings it passed over, and apply a softmax +cross-entropy over `−(c(s,s′) + h(s′))` with the plan's successor as the target. +The step cost has to be in there — GBFS compares `h` values, but the *correct* +comparison is `c + h*`, and with non-uniform action costs those differ. + +The default is not *pure* ranking. A ranking loss is invariant to any monotone +rescaling of the output, so it pins down no scale at all. That is fine for GBFS +and useless for weighted A*, where the weight multiplies a quantity that now +means nothing. `rank_weight=0.8` splits the objective; 0.8 is a reasonable +default, not a tuned optimum. + +The metric we select checkpoints on is **top-1 ranking accuracy**, not +validation MAE, for the same reason. + +### Not admissible, and it does not pretend to be + +Nothing in the training objective bounds the prediction from above. One +over-estimate costs A* its optimality guarantee. `LearnedHeuristic.admissible` +is `False` and the docs say to pair it with `gbfs`, or `wastar` if you want a +bounded-suboptimality knob. The suite asserts that a plan found with a learned +heuristic still validates — the heuristic may be wrong, the planner may not +become unsound. + +## Results + +Train on a small ladder, test on instances two to three times larger with +different seeds. Planner is `gbfs` throughout, budget 30 000 expansions / +30 s. Coverage is the fraction solved; expansions and seconds are means over +solved instances. + +### blocksworld — trained on 3–6 blocks, tested on 9–13 + +| heuristic | coverage | expansions | seconds | plan cost | +|---|---|---|---|---| +| **learned** (imitation only) | 0.90 | 366 | 0.141 | 49.1 | +| **learned** (+ search-cost tuning) | **1.00** | **137** | **0.038** | **48.2** | +| `hff` | 1.00 | 518 | 0.561 | 51.8 | +| `goalcount` | 1.00 | 2483 | 0.169 | 51.0 | +| `blind` | 0.00 | — | — | — | + +Against `hff`: **3.8× fewer expansions, 15× faster, and slightly cheaper plans.** +Corpus construction took 0.5 s, training 0.1 s, search-cost tuning about 50 s. +Held-out MAE 0.92, top-1 accuracy 0.923. The model has 1025 parameters over 14 +features. + +The tuned row is what this prints, end to end, in about a minute: + +```bash +jupyddl learn blocksworld --sizes 3-6 --seeds-per-size 3 \ + --cem 10 --cem-sizes 9-12 --evaluate 9-13 -o bw.heur.json +``` + +Three disjoint instance families are involved and the separation matters: +training uses seed 0, CEM tunes on seed 1000 and selects on seed 2000, and the +evaluation above is seed 7777. Nothing in the table was seen by any stage. + +Imitation alone left one test instance unsolved (`blocksworld-13-7777`, still +searching at 30 000 expansions). The reinforcement stage fixed it. + +Read the mean in that table with care: the held-out set has a heavy tail, and +that one instance dominates it. Nine of ten instances sit between 58 and 227 +expansions; the tenth is what moves the average. The per-instance breakdown, and +a correction to an earlier claim about *why* the tuned model is better, are in +[rl-for-search.md](rl-for-search.md). + +### gripper — trained on 2–5 balls, tested on 8–12 + +| heuristic | coverage | expansions | seconds | plan cost | +|---|---|---|---|---| +| **learned** | 1.00 | **66** | **0.014** | **29.4** | +| `hff` | 1.00 | 135 | 0.078 | 37.0 | +| `goalcount` | 1.00 | 142 | 0.006 | 39.0 | +| `blind` | 0.40 | 20198 | 0.480 | 25.0 | + +The cleanest win: half the expansions of `hff`, 5.6× faster, and *better plans*. +Top-1 accuracy 1.000 — gripper's structure is almost entirely captured by "how +many balls are still in the wrong room", which is exactly what these features +encode. + +Note `blind`'s plan cost of 25.0, the lowest in the table. It solved only 40% of +the instances and its search is effectively breadth-first on those, so it finds +near-optimal plans for the few it manages. Comparing mean cost across +heuristics with different coverage is misleading, which is why coverage is the +first column. + +### logistics — trained on size 2–4, tested on 5–7 + +| heuristic | coverage | expansions | seconds | plan cost | +|---|---|---|---|---| +| **learned** | 1.00 | 204 | 0.064 | 79.0 | +| `hff` | 1.00 | **35** | **0.029** | **45.7** | +| `goalcount` | 1.00 | 1374 | 0.080 | 102.0 | + +**A clear loss.** `hff` expands 6× fewer nodes and finds plans 42% cheaper. The +learned heuristic beats `goalcount` and nothing else. + +## Why logistics fails, precisely + +This is the most useful result in the set, because the mechanism is exact +rather than hand-wavy. + +The generated logistics domain has **two predicates**: `at` and `in`. The +feature vector is therefore `2×2 + 4 = 8` numbers. And those eight numbers +cannot distinguish *which* package is at *which* location — only *how many* +packages are somewhere, and how many goals remain. + +Two states where a package sits at its destination and where it sits across the +map produce **identical feature vectors** if the counts match. The heuristic is +being asked to estimate distance-to-goal from a description that has thrown away +the entire structure of the problem. Top-1 ranking accuracy is 0.656, barely +above chance for the branching factor involved, and the model is doing about as +well as anything could on that input. + +Blocksworld and gripper survive because their difficulty is closer to +count-shaped: "how many blocks are on the wrong thing", "how many balls are in +the wrong room". Logistics' difficulty is *relational* — it is about the +topology of who needs to go where — and counting is blind to topology. + +This is the argument for relational architectures stated as a measurement rather +than an intuition. It is precisely the gap STRIPS-HGN and GOOSE exist to close, +and it says clearly what the next step here has to be: features that survive a +permutation of the objects but not a permutation of the *relations between* +them. See [roadmap.md](roadmap.md). + +## Honest limitations + +- **Single seed, generated instances, three domains.** These numbers support + "this works and here is when it does not". They do not support a performance + claim. A publishable evaluation needs IPC benchmark domains, multiple seeds + with confidence intervals, and a per-domain train/test protocol fixed in + advance. +- **Coverage is measured under a budget.** A heuristic that solves 90% within + 30 000 expansions may solve 100% at 300 000. Coverage numbers here are "within + this budget", not "ever". +- **Plan cost is not controlled.** GBFS is not optimal with any of these + heuristics, so cost comparisons across rows compare two different suboptimal + behaviours, not two approximations of the same thing. +- **The corpus is small.** 118–166 samples per domain. The model has ~1000 + parameters. That it generalises at all is mostly a statement about how much + structure the features already impose. +- **Training targets on logistics were expensive to get.** Optimal solving of + the size 2–4 ladder took 13.9 s versus 0.5 s for blocksworld, and this + worsens sharply with size. Bootstrapping (see the RL note) exists for exactly + this. + +## References + +- Ernandes, M. and Gori, M. *Likely-admissible and sub-symbolic heuristics.* ECAI 2004. +- Yoon, S., Fern, A. and Givan, R. *Learning control knowledge for forward search planning.* JMLR 2008. +- Xu, Y., Fern, A. and Yoon, S. *Discriminative learning of beam-search heuristics for planning.* IJCAI 2007. +- Xu, Y., Fern, A. and Yoon, S. *Learning linear ranking functions for beam search with application to planning.* JMLR 2009. +- Arfaee, S. J., Zilles, S. and Holte, R. C. *Learning heuristic functions for large state spaces.* Artificial Intelligence 2011. +- Toyer, S., Trevizan, F., Thiébaux, S. and Xie, L. *Action schema networks: Generalised policies with deep learning.* AAAI 2018. +- Shen, W., Trevizan, F. and Thiébaux, S. *Learning domain-independent planning heuristics with hypergraph networks.* ICAPS 2020. +- Ferber, P., Helmert, M. and Hoffmann, J. *Neural network heuristics for classical planning: A study of hyperparameter space.* ECAI 2020. +- Chrestien, L., Edelkamp, S., Komenda, A. and Pevný, T. *Optimize planning heuristics to rank, not to estimate, cost-to-goal.* NeurIPS 2023. +- Chen, D. Z., Trevizan, F. and Thiébaux, S. *Return to tradition: Learning reliable heuristics with classical machine learning.* ICAPS 2024. + +These are cited from working knowledge of the literature. Verify each against +the published record before any of this is submitted anywhere. diff --git a/.docs/rl-for-search.md b/.docs/rl-for-search.md new file mode 100644 index 0000000..9b8c2d3 --- /dev/null +++ b/.docs/rl-for-search.md @@ -0,0 +1,275 @@ +# When learning a heuristic becomes reinforcement learning + +## The moment it changes + +Imitation asks *what does `h*` look like?* and fits a network to labels. + +The question we actually care about is *which heuristic makes search expand the +fewest nodes?* — and the answer is not the same. Imitation optimises a proxy, +and the proxy is wrong in two specific ways: + +1. **It is trained on the wrong distribution.** The corpus contains states that + lie on optimal plans. At search time the planner asks about states that do + not: the dead ends, the detours, the plateaux a mediocre heuristic wanders + into. This is covariate shift, and it is the standard failure mode of + behavioural cloning. +2. **It optimises the wrong quantity.** Fitting `h*` well is neither necessary + (a monotone rescaling changes nothing for GBFS) nor sufficient (a small error + at a critical branch point costs a whole subtree). + +The moment you stop fitting labels and start optimising *the number of nodes +expanded*, this is reinforcement learning. The heuristic is a policy; the +planner is the environment; the reward is the search you did not have to do. + +## The MDP + +Formally, for greedy best-first search: + +| | | +|---|---| +| **State** | the search state — the open list, the closed set, everything the planner knows | +| **Action** | which node to expand next | +| **Policy** | induced by `h`: expand `argmin` over the open list | +| **Transition** | pop the node, generate successors, push them | +| **Reward** | `−1` per expansion; terminal on reaching a goal | +| **Return** | `−(nodes expanded)` | + +Maximising return is exactly minimising search effort. The parameters of `h` +parameterise the policy, so this is policy search over a +non-differentiable objective. + +Note what the reward is *not*: it is not plan quality. A GBFS guided to a goal +in twelve expansions along a wasteful path scores better than one that takes two +hundred to find the optimal plan. If plan quality matters, it has to be in the +reward — `−(expansions) − λ·(plan cost)` — and it is not, by default, here. + +## Why the obvious policy gradient is hard + +REINFORCE needs a stochastic policy. The natural one is Boltzmann over the open +list: expand node `i` with probability `∝ exp(−h(si)/τ)`. Then + +```text +∇θ J = E[ Σt ∇θ log π(at | st) · Gt ] +``` + +Three things go wrong. + +**The episode is thousands of steps long, and every step gets the same credit.** +Expansion number 3 and expansion number 3000 receive identical returns under a +uniform `−1` reward. The variance of the estimator scales with episode length, +and episode length is the quantity being optimised — so the estimator is worst +exactly where the policy is worst. + +**The action space is the open list, which grows.** `π` is a softmax over +tens of thousands of items, changing size at every step. Not fatal, but it means +each gradient step touches every node in the frontier. + +**Making the policy stochastic makes the planner worse.** Deterministic +`argmin` is the thing we ship. Training a Boltzmann policy and deploying an +`argmin` one is a train/test mismatch on top of everything else, and as `τ → 0` +to close it, the gradient vanishes. + +Gehring et al. (ICAPS 2022) address the credit-assignment problem head-on by +using `h_ff` as a *dense reward generator* rather than relying on the sparse +goal signal — reward shaping with a domain-independent heuristic. That is a +good idea and the obvious next thing to try here. + +## What we do instead + +Three mechanisms, addressing different failures. All are in +`jupyddl/learn/rl.py`. + +### 1. DAgger — fix the distribution + +Run the current heuristic. Keep the states it *expanded*. Label each by solving +from it. Retrain on the union. Repeat. (Ross, Gordon and Bagnell, AISTATS 2011.) + +Labelling a state means finding a plan from it, which needs the planner to start +somewhere other than `task.initial_state()`. Rather than thread a start state +through fourteen planner signatures, `task_from_state` re-roots the task — +operators, goals, axioms and the numeric layer are untouched, so every planner +and heuristic works on the result unmodified. + +Labels come from a satisficing solver, so they are upper bounds; the samples are +tagged `optimal=False` and down-weighted during training. + +**Measured, on blocksworld:** one round with 15 states per task moved search +cost on the *training* instances from 89 to 15 expansions — a 6× improvement — +and moved a held-out set from 89 to **164**, i.e. it got 1.8× worse. + +That is a real result and it is worth stating plainly: **DAgger overfitted.** +With a 12-instance training ladder and 1066 aggregated samples from satisficing +labels, the model specialised to the states those twelve searches happened to +visit. The mechanism is sound and the setting was too small for it. It is off by +default. + +### 2. Bootstrapping — fix the data scarcity + +You cannot label a 20-block instance you cannot solve. But a heuristic trained +on 6 blocks may just crack 8, whose plans then teach it 10. Each round attempts +the unsolved instances with the current heuristic, adds what it managed, and +retrains. (Arfaee, Zilles and Holte, AIJ 2011.) + +This is a policy-improvement loop where the policy is "which instances can I +solve at all", and it is the honest answer to logistics taking 13.9 s to +generate a corpus at size 2–4. + +### 3. Direct search-cost optimisation — attack the objective itself + +Perturb the weight vector, score each candidate by *actually running the +planner*, keep the best fraction, refit the sampling distribution. The +cross-entropy method, with the planner as a black box. + +Evolutionary strategies are a legitimate alternative to policy gradients when +rollouts are cheap and the parameter vector is modest (Salimans et al., 2017), +which describes this exactly: ~1000 parameters, and a rollout is one GBFS run +on a small instance. + +**Measured, on blocksworld** — 10 iterations, population 12, σ=0.15, ~50 s: + +| | tuning set (sizes 9–12) | held-out benchmark (sizes 9–13, third seed family) | +|---|---|---| +| after imitation | 152 | 366 expansions, 0.90 coverage | +| after CEM | **100** | **137 expansions, 1.00 coverage** | + +CEM fixed the coverage failure imitation left behind, then shaved expansions: +2.7× fewer nodes on instances no stage of training ever saw. + +### A correction, and what the numbers actually say + +An earlier version of this note claimed that selecting the incumbent on a +disjoint instance family was what took the held-out score from 1734 to 137. +**That attribution was wrong**, and the way it went wrong is worth keeping. + +Two things changed in the same edit: the validation split went in, *and* the +perturbation scale `sigma` went from 0.05 to 0.15. The improvement was credited +entirely to the first. Re-running with only the selection rule varying: + +| | held-out mean expansions | +|---|---| +| σ=0.05, selected on the tuning set | 1734 | +| σ=0.05, selected on a disjoint family | 1730 | +| σ=0.15, selected on the tuning set | 137 | +| σ=0.15, selected on a disjoint family | 137 | + +The validation split makes no difference here. **`sigma` was doing all the +work.** Widening the validation family to span sizes beyond the tuning range +does not change it either (1741 / 136). + +### And the mean was doing the lying + +Per instance, on the held-out set: + +| instance | imitation | σ=0.05 | σ=0.15 | +|---|---|---|---| +| blocksworld-09-7777 | 77 | 69 | 70 | +| blocksworld-09-7778 | 126 | 67 | 68 | +| blocksworld-10-7777 | 70 | 33 | 58 | +| blocksworld-10-7778 | 287 | 229 | 217 | +| blocksworld-11-7777 | 113 | 90 | 81 | +| blocksworld-11-7778 | 150 | 121 | 111 | +| blocksworld-12-7777 | 349 | 228 | 227 | +| blocksworld-12-7778 | 114 | 72 | 70 | +| **blocksworld-13-7777** | **30 000 (unsolved)** | **16 121** | **214** | +| blocksworld-13-7778 | 2 004 | 273 | 255 | + +Nine of ten instances improve under *both* settings, by roughly the same +factor. The entire 1734-versus-137 gap is one instance. Reporting a mean over a +distribution with that shape is close to reporting that one instance and +nothing else, and it is why the headline "137" should be read as a mean over a +heavy tail rather than a typical case. + +What survives, and is worth having: + +- **CEM fixed a coverage failure.** Imitation could not solve + `blocksworld-13-7777` within 30 000 expansions. Both tuned versions could. + That is a real capability change, not a shaved constant. +- **Every instance improves.** The median improvement is around 1.4×, which is + the honest version of the headline. +- **The validation split is still correct**, just not load-bearing here. It + guarantees the returned model is no worse on instances the optimiser did not + fit — on this run, validation went 276 → 73 — and that guarantee costs one + extra scoring pass per iteration. Keep it; do not credit it with results it + did not produce. + +`optimise_search_cost` reports both numbers every iteration, so a run that is +fitting its tuning set while losing validation is visible rather than silent: + +```text +cem iter 7: tuning 101, validation 73, incumbent 73 (coverage 1.00) +``` + +The general lesson is the ordinary one, which the derivative-free framing made +easy to forget: **change one thing at a time, and look at the distribution +before believing the mean.** `learn_heuristic` always passes one. + +## The three things that decide whether this works + +### Start from the imitation solution + +Search cost is a *step function* of the weights over most of the space: every +candidate that solves nothing scores identically at `failure_penalty × budget`. +A randomly initialised policy sits in that flat region and no method — CEM, +REINFORCE, anything — gets a signal out of it. Imitation is what puts the +optimiser somewhere the objective can distinguish. + +### Tune on instances that have headroom + +This one cost us an afternoon and is the most transferable lesson here. + +The first CEM run tuned on the *training ladder* and moved the score from 12.83 +to 12.75 — noise. The reason: after imitation the heuristic already expands +roughly as many nodes as the plan is long on those instances. There is nothing +left to win, so every perturbation scores the same, and the objective is flat +for the second time, for a different reason. + +Re-running on sizes 9–12, where the imitated heuristic expanded 1605 nodes on +average, moved it to 64 — a 25× improvement — and cut a held-out set from 336 to +122. + +**Optimise where the search is still bad.** `learn_heuristic` therefore defaults +`cem_sizes` to a rung above the training ladder rather than reusing it. + +### Pick a perturbation scale, and check it + +The section above. `sigma=0.05` and `sigma=0.15` differ by an order of +magnitude in held-out cost on the same command, and the difference is +concentrated in the hardest instance. Selecting on a disjoint family is still +right — it bounds what you can return — but it is a guardrail, not the knob. + +## Design details worth knowing + +- **The failure penalty is `2 × budget`, not `1 ×`.** At exactly `1 ×` the + optimiser is indifferent between solving an instance at the budget limit and + not solving it. Greater than 1 means solving slowly always beats giving up. +- **Scores are comparable only at a fixed budget**, since the penalty is a + multiple of it. Every optimiser here holds the budget fixed. +- **The incumbent is never displaced by a worse validation score.** The + returned bundle cannot be worse than the one passed in, on the instances it + was *selected* on. That is a real guarantee once those instances are disjoint + from the ones being fitted, and was nearly worthless when they were not. +- **The mean of the elites is re-evaluated each iteration**, not assumed to be + good. CEM's distribution mean is not one of the sampled candidates and can be + worse than all of them. + +## What we have not tried + +- **Reward shaping with `h_ff`** as a dense signal (Gehring et al. 2022). The + most promising unexplored direction, and it fits the existing structure. +- **Policy-guided search with guarantees** (Orseau and Lelis, AAAI 2021), which + puts a bound on the search effort of a learned policy rather than hoping. +- **Optimising for plan quality as well as effort.** Currently unmodelled. +- **REINFORCE with a learned baseline** over a Boltzmann-softened open list, + which is the textbook approach and which we skipped for the variance reasons + above rather than because it is wrong. + +## References + +- Ross, S., Gordon, G. and Bagnell, D. *A reduction of imitation learning and structured prediction to no-regret online learning.* AISTATS 2011. +- Arfaee, S. J., Zilles, S. and Holte, R. C. *Learning heuristic functions for large state spaces.* Artificial Intelligence 2011. +- Salimans, T., Ho, J., Chen, X., Sidor, S. and Sutskever, I. *Evolution strategies as a scalable alternative to reinforcement learning.* arXiv 2017. +- Orseau, L. and Lelis, L. *Policy-guided heuristic search with guarantees.* AAAI 2021. +- Gehring, C., Asai, M., Chitnis, R., Silver, T., Kaelbling, L. P., Sohrabi, S. and Katz, M. *Reinforcement learning for classical planning: Viewing heuristics as dense reward generators.* ICAPS 2022. +- Rubinstein, R. Y. *Optimization of computer simulation models with rare events.* European Journal of Operational Research 1997. (The cross-entropy method.) + +Cited from working knowledge; verify before publication. diff --git a/.docs/roadmap.md b/.docs/roadmap.md new file mode 100644 index 0000000..6995990 --- /dev/null +++ b/.docs/roadmap.md @@ -0,0 +1,147 @@ +# What to do next + +Ordered by expected value. Each entry says what the experiment is, so that a +negative result is as informative as a positive one. + +--- + +## 1. Relational features — fix logistics + +**Why first.** The logistics failure is not a tuning problem, it is a +representational one, and it is measured rather than suspected: the domain has +two predicates, so the feature vector is 8 numbers, and two states with a +package at its destination and across the map are *identical* under it. Top-1 +ranking accuracy 0.656. No amount of training fixes an input that has thrown the +problem away. + +**What to build.** Weisfeiler-Leman features over the grounded problem graph, +following GOOSE (Chen, Trevizan and Thiébaux, ICAPS 2024). The striking finding +there is that classical ML over WL features beat the deep relational models it +was compared against — which, if it holds, is exactly the right shape for this +library: colour refinement is a few passes of hashing over an adjacency +structure, cheap enough to run per state in Python. + +Sketch: build a graph with a node per object and per ground atom, edges from +atoms to their arguments, node labels from predicate symbol and goal membership. +Run `k` rounds of colour refinement. The feature vector is the histogram of +colours. Fixed length via a hash into `d` buckets, so it stays size-invariant. + +**The experiment.** Same three domains, same protocol. Success is logistics +top-1 above 0.85 and expansions within 2× of `hff`, *without* regressing +blocksworld or gripper. Watch the per-state cost: if it exceeds ~3× the current +features, the wall-clock advantage that makes this whole line worthwhile is gone. + +**Risk.** Colour refinement per state may simply be too slow in CPython. Measure +before building the rest. + +--- + +## 2. Reward shaping with `h_ff` + +**Why.** The credit-assignment problem in [rl-for-search.md](rl-for-search.md) +is the reason we use a derivative-free method rather than a policy gradient. +Gehring et al. (ICAPS 2022) attack it directly by using a classical heuristic as +a dense reward generator. It is the most promising unexplored direction and it +fits the existing structure with no new machinery. + +**What to build.** A potential-based shaping term `F(s, s′) = γ·Φ(s′) − Φ(s)` +with `Φ = −h_ff`. Potential-based shaping is policy-invariant (Ng, Harada and +Russell, ICML 1999), so it changes the learning dynamics without changing what +the optimal policy is — which is exactly the property you want and exactly what +naive shaping gets wrong. + +**The experiment.** Compare CEM-with-shaping against plain CEM at equal wall +clock, not equal iterations. Shaping that needs an `h_ff` evaluation per state +is not free. + +--- + +## 3. A proper evaluation protocol + +**Why.** Everything in these notes is single-seed on generated instances. It +supports "this works, and here is exactly when it does not". It does not support +a performance claim, and pretending otherwise is how learned-planning results +get published and then fail to replicate. + +**What to build.** + +- IPC benchmark domains, not just the generators. The generators produce + instances with a regularity real domains do not have, and a feature space + keyed on predicate counts is exactly the sort of thing that exploits it. +- A fixed train/test split *per domain*, declared before running anything. +- Five seeds, report median and interquartile range. +- Coverage under a declared budget, plus expansions and wall-clock, plus plan + cost, on every table. Reporting expansions alone is the standard way to + publish a learned heuristic that is slower than the one it replaced. +- A per-node cost measurement — microseconds per heuristic evaluation — as a + first-class number. + +**The experiment.** This *is* the experiment. Expect the picture to get worse: +generated instances flatter the approach. + +--- + +## 4. Learning from failure, not just from plans + +**Why.** The corpus currently contains only states on successful plans. The +information in "the search spent 4000 expansions in this region and found +nothing" is thrown away, and it is arguably more valuable — it identifies the +plateaux and dead ends where a heuristic is actually costing you. + +**What to build.** Record, per expanded state, whether it appeared on the final +plan. Train an auxiliary head to predict "is this state on a solution path", +and add a penalty term for confident low estimates on states that turned out to +be dead ends. Related in spirit to Ferber et al.'s work on progress states in +GBFS. + +**The experiment.** Does it reduce plateau escape time? Measure expansions +between successive improvements in best-`h`, not just total expansions — a +heuristic can have identical total cost and a very different plateau profile. + +--- + +## 5. Portfolio / per-instance selection + +**Why.** The three-domain result is not "learned beats `hff`", it is "learned +beats `hff` on two of three, and loses badly on the third". That is a selection +problem, and selection problems are usually easier than the underlying one. + +**What to build.** A cheap classifier over instance-level features (predicate +count, object count, goal count, `h_ff` of the initial state, branching factor +at the root) predicting which heuristic to use. `jupyddl` already has the +benchmark harness that generates the training data for this. + +**The experiment.** Does the selector beat always-`hff` and always-learned on a +held-out set of domains? Note the honest baseline: a selector that always picks +`hff` is the thing to beat, and on three domains it wins one-third of the time +by construction. + +--- + +## 6. Make the browser workbench train + +**Why.** `jupyddl/learn` is stdlib-only, so it already runs under Pyodide. A +"watch a heuristic learn, then watch it search" view would make the whole line +legible in a way a table of numbers does not. + +**What to build.** A fifth view in `web/`: pick a domain, generate a ladder, +watch the corpus fill, watch the loss curve and the top-1 accuracy, then race +the trained heuristic against `hff` on a larger instance with the existing +wavefront visualisation. + +**Cost.** `tools/build_web.py` currently skips `jupyddl/learn` to keep the +bundle small. Un-skipping it is one line. Training in WebAssembly without NumPy +will be slow — budget for a demo-scale ladder, not a real one. + +--- + +## Things deliberately not on this list + +- **Deep relational architectures (ASNets, HGNs) as implemented in the + literature.** They need a tensor library, which breaks the zero-dependency + core, and their per-state cost is hard to justify in CPython. Item 1 is the + version of this idea that fits here. +- **Learning the search algorithm rather than the heuristic.** Interesting, much + larger, and it would not compose with the fourteen planners already in the + registry. +- **LLM-based planning.** A different research programme, not this one. diff --git a/AGENTS.md b/AGENTS.md index ea191c1..2d9b9a6 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -54,6 +54,12 @@ native build step, and the core has zero runtime dependencies. initial literals, object fluents) rewritten into the classical core **before** grounding. Add front-end features here as source-to-source transformations, not as special cases in the grounder or the search. +- `jupyddl/learn/` — learned heuristics. **Stdlib-only like the core**; NumPy is + a speed option and the pure-Python path must keep working. Two stages: + `train.py` imitates `h*` from solved plans, `rl.py` optimises search cost + directly. Nothing in the core imports it — `jupyddl.heuristics` resolves + `learned:` lazily inside the loader, so a planner that never asks for + one never pays for it. - `jupyddl/generator.py` — reproducible instance generators. - `jupyddl/trace.py` — search observers, events, `SearchTrace` (JSON). - `jupyddl/live.py` — the terminal dashboard; **stdlib only, keep it that way**, @@ -62,7 +68,10 @@ native build step, and the core has zero runtime dependencies. import this package. - `web/` — the Pyodide playground; `tools/build_web.py` bundles the package sources and demos into `web/dist` (committed). -- `tools/make_promo.py` — renders the promo video from measured runs. +- `tools/make_promo.py` — renders the main promo video from measured runs. +- `tools/make_learn_promo.py` — the learned-heuristic/RL video. It re-measures + everything including both failure modes, so it cannot drift from `.docs/`; + `promo/rl-data.json` caches the pass, delete it to re-measure. ### The condition pipeline Conditions are a **formula tree in negation normal form**: `parse_condition` @@ -137,6 +146,34 @@ planners must go through the task rather than the operator to honour them: - **The clock is checked every expansion** (`Budget.check_every=1`). That looks wasteful but is not: one LM-cut expansion can cost tens of milliseconds, and a coarser interval overshoots a short `--time-limit` enormously. +### Learned heuristics, non-obvious parts +- **Features are keyed on the predicate symbol, never the ground atom**, and + normalised per symbol. That is what makes the vector the same length for 4 + blocks and 40. Getting `predicate_of` wrong does not raise — it gives every + ground atom its own slot and silently destroys transfer, so + `tests/test_learn.py` pins both spellings (`(on a b)` and `move(a,b)`). +- **The default objective is ranking, not regression.** GBFS reads the *order* + a heuristic imposes, never its values; a model uniformly 30 too high guides + perfectly. Checkpoint selection is on top-1 accuracy for the same reason. The + regression term is kept at a small weight only to anchor a scale, which a + pure ranking loss leaves undefined and `wastar` needs. +- **Report the distribution, not just the mean.** The held-out blocksworld set + has a heavy tail: nine of ten instances land between 58 and 227 expansions and + the tenth is worth thousands, so the mean is close to a report of that one + instance. Two published claims here were wrong because of it — see the + correction in `.docs/rl-for-search.md`, which also records the more + embarrassing cause: two settings changed in one edit and the improvement was + credited to the wrong one. +- **Two things decide whether the RL stage does anything.** It must start from + the imitation solution (search cost is flat over every parameter vector that + solves nothing), and it must tune on instances with *headroom* — on the + training ladder the heuristic already expands about as many nodes as the plan + is long, so every perturbation scores the same. Measured: tuning on the + training sizes moved the score 12.83 → 12.75; tuning a rung higher moved it + 1605 → 64. `learn_heuristic` defaults `cem_sizes` above the training ladder. +- **A learned heuristic is never admissible** and must not be used to claim an + optimal plan. The suite asserts plans stay *valid*, which is the invariant + that does hold. - Extend via the registries: `jupyddl.search.PLANNERS`, `jupyddl.heuristics.HEURISTICS` and `jupyddl.generator.GENERATORS`. The CLI, benchmark harness and web workbench all read from them, so a new entry shows diff --git a/CHANGELOG.md b/CHANGELOG.md index 1f7a611..cc01cf7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,33 @@ project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## [Unreleased] +### Added +- **`jupyddl.learn`: heuristics trained from your own solved plans.** Every + solved instance is a labelled trajectory — the cost of a plan's suffix from a + state on it is that state's cost-to-go — so a corpus of plans is supervision + that already exists. Stdlib-only like the rest of the core; `pip install + jupyddl[learn]` adds NumPy purely for speed, and the two code paths are pinned + together by a test. +- **A ranking objective, by default.** Greedy best-first search never reads a + heuristic value, only the order it imposes, so the loss optimises the order: + at each state on a plan, the successor the plan took should sort ahead of the + siblings it passed over. A small regression term is kept to anchor a scale, + which pure ranking leaves undefined and `wastar` needs. Checkpoints are + selected on top-1 ranking accuracy rather than validation error. +- **A reinforcement stage that optimises search cost directly**, since "nodes + expanded" is not a differentiable function of the weights: DAgger for the + distribution shift, bootstrapping for the instances too hard to label, and + the cross-entropy method over the weight vector with the planner as a black + box. On blocksworld this took a held-out set from 366 expansions and 0.90 + coverage to 131 and 1.00. +- **`jupyddl learn`**, and `learned:` accepted anywhere a heuristic + name is — `solve`, `benchmark`, the API. `make_heuristic` also passes through + an already-built heuristic, so callers holding a trained model need not + round-trip it to disk. +- `.docs/` — research notes: the prior work, the measured results including the + domain where this loses badly and exactly why, the MDP the RL stage + corresponds to, and a roadmap. + ### Fixed - **`always` was not enforced across the actions timed initial literals compile to.** The invariant went onto the domain's own actions and the goal, on the diff --git a/README.md b/README.md index 8c03759..7db3282 100644 --- a/README.md +++ b/README.md @@ -46,6 +46,11 @@ is trivial to install, embed, teach with, and build on. Iterated Width, branch and bound, and anytime weighted A*. - 📊 **Heuristics from classical to SOTA**: blind, goal-count, `h_max`, `h_add`, FF (`h_ff`), critical-path `h^m` (`h1`/`h2`), and **LM-cut**. +- 🧠 **Learned heuristics** trained from your own solved plans, then tuned + against the number of nodes search expands rather than against `h*`. On + blocksworld this beats `h_ff` by ~4× in expansions and ~15× on the clock, at + instance sizes three times larger than it trained on. + ([how, and when it fails](#learned-heuristics-)) - ⏱️ **Search budgets** on every planner: bound a run by nodes or by seconds and a truncated result says so, so "we stopped looking" never masquerades as "no plan exists". @@ -69,17 +74,19 @@ Requires Python ≥ 3.9. Using [uv](https://docs.astral.sh/uv/) (recommended): ```bash uv venv -uv pip install -e ".[dev,viz]" # 'viz' pulls matplotlib for charts and animations +uv pip install -e ".[dev,viz,learn]" # viz = matplotlib charts, learn = numpy (speed only) ``` or with plain pip: ```bash -python -m pip install -e ".[dev,viz]" +python -m pip install -e ".[dev,viz,learn]" ``` -The core framework needs nothing but the standard library — the `viz` extra is -only for the matplotlib charts. +The core framework needs nothing but the standard library. `viz` is only for +the matplotlib charts; `learn` is only for speed — `jupyddl.learn` trains and +evaluates on the standard library alone, NumPy just makes it one to two orders +of magnitude faster. ## Command line ⚔️ @@ -295,6 +302,60 @@ ladder = [write_instance("gripper", "instances/", size=n, seed=1) purely random operators are almost always unsolvable, which makes for a useless benchmark. The test suite grounds and solves everything each generator emits. +## Learned heuristics 🧠 + +Every solved instance is a labelled trajectory: the cost of a plan's suffix from +any state on it is that state's cost-to-go. `jupyddl learn` turns a corpus of +those into a heuristic, then — optionally — stops imitating `h*` and starts +optimising the thing that actually matters, the number of nodes search expands. + +```bash +# generate a ladder, solve it, fit a heuristic, tune it on search cost +jupyddl learn blocksworld --sizes 3-6 --cem 10 --evaluate 9-13 -o bw.heur.json + +# then use it anywhere a heuristic name is accepted +jupyddl solve domain.pddl problem.pddl -s gbfs -H learned:bw.heur.json +jupyddl benchmark demos --planners gbfs --heuristic learned:bw.heur.json +``` + +Trained on 3–6 block instances, evaluated on 9–13 block instances it has never +seen, against greedy best-first search: + +| heuristic | coverage | expansions | seconds | plan cost | +|---|---|---|---|---| +| **`learned`** | 1.00 | **137** | **0.038** | **48.2** | +| `hff` | 1.00 | 518 | 0.561 | 51.8 | +| `goalcount` | 1.00 | 2483 | 0.169 | 51.0 | + +Nearly four times fewer expansions than `hff` and fifteen times faster, because +the network is a thousand multiply-adds and `hff` is a relaxed-plan extraction. +Those are the numbers the command above prints, on one CPU, in about a minute. + +It does not always win — on logistics it loses to `hff` by 6× and the reason is +exact rather than mysterious: that domain has two predicates, so the feature +vector cannot tell *which* package is where, only how many are somewhere. That +result, the RL formulation, and what to build next are written up in +[`.docs/`](.docs/), and there is a +[97-second tour of the RL half](promo/jupyddl-rl.mp4) — including the two +measurement mistakes that shaped the design. + +Read the 137 as a mean over a heavy tail: nine of the ten held-out instances sit +between 58 and 227 expansions and the tenth moves the average on its own. The +firmest claim is the coverage one — imitation could not solve that instance +inside 30 000 expansions, and the tuned heuristic solves it in 214. + +```python +from jupyddl.learn import learn_heuristic + +bundle = learn_heuristic("gripper", sizes=range(2, 6), cem_iterations=10) +bundle.save("gripper.heur.json") +``` + +The learning stack is stdlib-only like the rest of the core; `pip install +jupyddl[learn]` adds NumPy purely for speed. A learned heuristic is **not +admissible** — nothing in the objective bounds it from above — so pair it with +`gbfs`, or `wastar` if you want a bounded-suboptimality knob. + ## Soft goals and trajectory constraints 🎯 Preferences say what you would *rather* were true; constraints say what must diff --git a/jupyddl/cli.py b/jupyddl/cli.py index ccb1426..344d0f0 100644 --- a/jupyddl/cli.py +++ b/jupyddl/cli.py @@ -1,7 +1,7 @@ """Command-line interface. -``solve``, ``benchmark``, ``animate``, ``demo``, ``requirements`` and -``generate``. Every long-running command accepts ``--max-expansions`` and +``solve``, ``benchmark``, ``animate``, ``demo``, ``requirements``, ``generate`` +and ``learn``. Every long-running command accepts ``--max-expansions`` and ``--time-limit``; when a search stops on one of those it says so rather than reporting the instance unsolvable. """ @@ -20,19 +20,48 @@ summarize, to_csv, ) -from .heuristics import HEURISTICS +from .heuristics import HEURISTICS, LOADERS from .search import INFORMED_PLANNERS, PLANNERS INFORMED = set(INFORMED_PLANNERS) +def heuristic_spec(value: str) -> str: + """Validate ``-H``: a registry name, ``none``, or ``kind:argument``. + + ``choices=`` cannot express this — a trained heuristic is identified by a + path that does not exist until someone trains one — so the check is a type + function instead. It still rejects typos eagerly, which is the only reason + ``choices=`` was worth having. + """ + if value in HEURISTICS or value == "none": + return value + kind, sep, argument = value.partition(":") + if sep and kind in LOADERS: + if not argument: + raise argparse.ArgumentTypeError( + f"'{kind}:' needs an argument, e.g. {kind}:model.json" + ) + return value + raise argparse.ArgumentTypeError( + f"unknown heuristic '{value}'; expected one of " + f"{sorted(HEURISTICS) + ['none']} " + f"or {'/'.join(sorted(LOADERS))}:" + ) + + def _add_solve(sub): p = sub.add_parser("solve", help="solve a single PDDL instance") p.add_argument("domain") p.add_argument("problem") p.add_argument("-s", "--search", default="astar", choices=sorted(PLANNERS)) p.add_argument( - "-H", "--heuristic", default="lmcut", choices=sorted(HEURISTICS) + ["none"] + "-H", + "--heuristic", + default="lmcut", + type=heuristic_spec, + metavar="NAME", + help="a heuristic name, none, or learned:", ) p.add_argument( "-w", "--weight", type=float, default=2.0, help="weight for weighted A*" @@ -107,7 +136,12 @@ def _add_animate(sub): p.add_argument("-o", "--output", default="search.mp4") p.add_argument("-s", "--search", default="astar", choices=sorted(PLANNERS)) p.add_argument( - "-H", "--heuristic", default="lmcut", choices=sorted(HEURISTICS) + ["none"] + "-H", + "--heuristic", + default="lmcut", + type=heuristic_spec, + metavar="NAME", + help="a heuristic name, none, or learned:", ) p.add_argument("--fps", type=int, default=30) p.add_argument("--seconds", type=float, default=8.0) @@ -304,6 +338,151 @@ def _add_generate(sub): p.set_defaults(func=_cmd_generate) +def _add_learn(sub): + from .generator import GENERATORS + + p = sub.add_parser( + "learn", + help="train a heuristic from solved plans, then reinforce it on search cost", + description=( + "Generate a ladder of instances, solve the small ones, fit a network " + "to the cost-to-go their plans reveal, and optionally tune it against " + "the number of nodes search actually expands. Writes a model that " + "'-H learned:' accepts everywhere." + ), + ) + p.add_argument("kind", choices=sorted(GENERATORS)) + p.add_argument("-o", "--output", default=None, help="write the model here") + p.add_argument( + "--sizes", + default="3-6", + help="training ladder, as 'lo-hi' (default 3-6). Keep these small: they " + "have to be solvable optimally", + ) + p.add_argument( + "--seeds-per-size", + type=int, + default=2, + help="instances per rung; more seeds usually beats more rungs", + ) + p.add_argument("--seed", type=int, default=0) + p.add_argument("--epochs", type=int, default=60) + p.add_argument( + "--rank-weight", + type=float, + default=0.8, + help="share of the objective spent on ordering rather than magnitude; " + "GBFS only reads the order", + ) + p.add_argument( + "--dagger", + type=int, + default=0, + metavar="ROUNDS", + help="retrain on the states the heuristic's own search visits", + ) + p.add_argument( + "--bootstrap", + default=None, + metavar="LO-HI", + help="grow the corpus with harder instances as they become solvable", + ) + p.add_argument( + "--cem", + type=int, + default=0, + metavar="ITERATIONS", + help="optimise expansions directly (needs instances with headroom; " + "see --cem-sizes)", + ) + p.add_argument( + "--cem-sizes", + default=None, + metavar="LO-HI", + help="instances to tune search cost on; defaults to a rung above the " + "training ladder, because the training ladder has no headroom left", + ) + p.add_argument( + "--evaluate", + default=None, + metavar="LO-HI", + help="after training, benchmark against hff/goalcount on these sizes", + ) + p.set_defaults(func=_cmd_learn) + + +def _range(spec: str): + """Parse ``'3-6'`` or ``'5'`` into a range.""" + text = str(spec).strip() + if "-" in text: + lo, _, hi = text.partition("-") + return range(int(lo), int(hi) + 1) + return range(int(text), int(text) + 1) + + +def _cmd_learn(args) -> int: + from .learn import RLConfig, TrainConfig, learn_heuristic + from .learn.pipeline import ( + evaluate_transfer, + summarise_transfer, + tasks_from_generator, + ) + + output = args.output or f"{args.kind}.heur.json" + try: + bundle = learn_heuristic( + args.kind, + sizes=_range(args.sizes), + seeds_per_size=args.seeds_per_size, + seed=args.seed, + train_config=TrainConfig( + epochs=args.epochs, rank_weight=args.rank_weight, seed=args.seed + ), + rl_config=RLConfig(seed=args.seed, verbose=True), + dagger_rounds=args.dagger, + bootstrap_sizes=_range(args.bootstrap) if args.bootstrap else None, + cem_iterations=args.cem, + cem_sizes=_range(args.cem_sizes) if args.cem_sizes else None, + verbose=True, + ) + except RuntimeError as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + + bundle.save(output) + print(f"\nWrote {output}") + print( + f" features {bundle.space.size} over {len(bundle.space.vocabulary)} predicates" + ) + print(f" parameters {bundle.model.num_parameters}") + metrics = bundle.metrics + print( + f" held out MAE {metrics.get('mae', 0):.2f}" + f" top-1 {metrics.get('top1', 0):.3f}" + f" over {metrics.get('held_out_instances', 0)} instances" + ) + print(f"\nUse it:\n jupyddl solve -s gbfs -H learned:{output}") + + if args.evaluate: + tasks = tasks_from_generator( + args.kind, _range(args.evaluate), seed=args.seed + 7777, seeds_per_size=2 + ) + print(f"\nBenchmark on {len(tasks)} unseen instances (gbfs):") + rows = evaluate_transfer(bundle, tasks) + summary = summarise_transfer(rows) + print( + f" {'heuristic':<12}{'coverage':>10}{'expanded':>12}" + f"{'seconds':>10}{'cost':>8}" + ) + for name in ["learned"] + [k for k in summary if k != "learned"]: + agg = summary[name] + print( + f" {name:<12}{agg['coverage']:>10.2f}{agg['mean_expanded']:>12.0f}" + f"{agg['mean_seconds']:>10.3f}{agg['mean_cost']:>8.1f}" + ) + return 0 + + def _cmd_requirements(args) -> int: from .requirements import as_rows, summary @@ -477,6 +656,7 @@ def main(argv=None) -> int: _add_demo(sub) _add_requirements(sub) _add_generate(sub) + _add_learn(sub) args = parser.parse_args(argv) return args.func(args) diff --git a/jupyddl/heuristics/__init__.py b/jupyddl/heuristics/__init__.py index 96d2f39..8baf757 100644 --- a/jupyddl/heuristics/__init__.py +++ b/jupyddl/heuristics/__init__.py @@ -24,13 +24,47 @@ } -def make_heuristic(name: str, task) -> Heuristic: - """Instantiate a heuristic by name (see :data:`HEURISTICS`).""" +# A heuristic that needs more than a task to build takes an argument after a +# colon: ``learned:blocksworld.heur.json``. Keeping this a *string* spec rather +# than a Python object is what makes a trained heuristic usable from the CLI, +# the benchmark harness and the web UI without any of them knowing it exists. +LOADERS = {} + + +def _load_learned(path: str, task) -> Heuristic: + """Resolve ``learned:``, importing the learning stack only if asked.""" + from ..learn.heuristic import HeuristicBundle + + return HeuristicBundle.load(path).bind(task) + + +LOADERS["learned"] = _load_learned + + +def make_heuristic(name, task) -> Heuristic: + """Instantiate a heuristic by name (see :data:`HEURISTICS`). + + Accepts a registry name (``"lmcut"``), a parameterised spec + (``"learned:model.json"``), or an already-built heuristic, which is passed + through so callers holding a trained model need not round-trip it to disk. + """ + if callable(name) and not isinstance(name, str): + return name + if ":" in name: + kind, _, argument = name.partition(":") + loader = LOADERS.get(kind) + if loader is None: + raise ValueError( + f"Unknown parameterised heuristic '{kind}'. " + f"Available: {sorted(LOADERS)}" + ) + return loader(argument, task) try: factory = HEURISTICS[name] except KeyError: raise ValueError( f"Unknown heuristic '{name}'. Available: {sorted(HEURISTICS)}" + + (f" (or {sorted(LOADERS)} with an argument)" if LOADERS else "") ) from None return factory(task) @@ -45,5 +79,6 @@ def make_heuristic(name: str, task) -> Heuristic: "LMCutHeuristic", "CriticalPathHeuristic", "HEURISTICS", + "LOADERS", "make_heuristic", ] diff --git a/jupyddl/learn/__init__.py b/jupyddl/learn/__init__.py new file mode 100644 index 0000000..c475828 --- /dev/null +++ b/jupyddl/learn/__init__.py @@ -0,0 +1,73 @@ +"""Learned heuristics: imitate a corpus of plans, then optimise search directly. + +The pipeline has two stages and they answer different questions. + +**Imitation** (:mod:`jupyddl.learn.train`) asks *what does h\\* look like?* It +fits a network to the cost-to-go that solved plans reveal. Cheap, stable, and +limited by a mismatch it cannot see: it is trained on states that lie on +optimal plans, and at search time it is asked about states that do not. + +**Reinforcement** (:mod:`jupyddl.learn.rl`) asks the question that actually +matters: *which heuristic expands the fewest nodes?* Search cost is not a +differentiable function of the weights — it runs through a priority queue — so +this stage reaches it three ways: aggregating data from the states search +really visits (DAgger), growing the corpus with instances the current +heuristic just became able to solve (bootstrapping), and optimising expansions +directly with a derivative-free method. + +Quickstart:: + + from jupyddl.learn import learn_heuristic + + bundle = learn_heuristic("blocksworld", sizes=range(4, 9), seed=0) + bundle.save("blocksworld.heur.json") + +then anywhere a heuristic name is accepted:: + + jupyddl solve domain.pddl problem.pddl -s gbfs -H learned:blocksworld.heur.json + +Nothing here is imported by the core. :mod:`jupyddl.heuristics` resolves +``learned:`` lazily, so a planner that never asks for one never pays for it. +""" + +from __future__ import annotations + +from .dataset import Corpus, RankingGroup, Sample, build_corpus, samples_from_plan +from .features import FeatureSpace +from .heuristic import HeuristicBundle, LearnedHeuristic +from .model import MLP, numpy_available +from .pipeline import learn_heuristic, solved_corpus, tasks_from_generator +from .rl import ( + RLConfig, + bootstrap, + dagger, + optimise_search_cost, + search_cost, +) +from .train import TrainConfig, TrainReport, evaluate, evaluate_ranking, train + +__all__ = [ + "FeatureSpace", + "Sample", + "RankingGroup", + "Corpus", + "build_corpus", + "samples_from_plan", + "MLP", + "numpy_available", + "HeuristicBundle", + "LearnedHeuristic", + "TrainConfig", + "TrainReport", + "train", + "evaluate", + "evaluate_ranking", + "RLConfig", + "dagger", + "bootstrap", + "optimise_search_cost", + "search_cost", + "learn_heuristic", + "solved_corpus", + "tasks_from_generator", +] diff --git a/jupyddl/learn/dataset.py b/jupyddl/learn/dataset.py new file mode 100644 index 0000000..483787d --- /dev/null +++ b/jupyddl/learn/dataset.py @@ -0,0 +1,323 @@ +"""Turn solved plans into supervision. + +A plan is a labelled trajectory whether or not it was collected for that +purpose. If ``s0 -a1-> s1 -a2-> ... -> sn`` is a plan, then the cost of the +suffix from ``si`` is an upper bound on ``h*(si)``, and when the plan is +optimal it *is* ``h*(si)``. One solved instance therefore yields one sample per +state on its plan, free. + +Two kinds of sample come out of this module, and the difference matters more +than it looks: + +**Regression samples** — ``(features(si), cost of the suffix)``. The obvious +thing, and what most of the literature reports. + +**Ranking samples** — at each ``si``, the successor the plan takes together +with the siblings it passes over. Greedy best-first search never reads an +h-value; it reads the *order* h imposes on the open list. A model with a +systematic offset of +30 is useless as a cost estimate and a perfect guide, and +a model that is accurate on average but inverts two siblings sends the search +down the wrong subtree. Optimising the order directly is the better-matched +objective (Chrestien et al., NeurIPS 2023) and is why +:func:`jupyddl.learn.train.train` defaults to it. + +Sampling only states on the plan leaves a distribution-shift problem: at search +time the planner asks about states no plan ever passed through, and the model +extrapolates. :mod:`jupyddl.learn.rl` is where that gets fixed. +""" + +from __future__ import annotations + +import json +import math +import random +from dataclasses import dataclass, field +from typing import Optional + +from dataclasses import replace as _replace + +from ..task import facts_of, values_of +from .features import FeatureSpace + +__all__ = [ + "Sample", + "RankingGroup", + "Corpus", + "samples_from_plan", + "build_corpus", + "task_from_state", +] + + +def task_from_state(task, state): + """The same task, re-rooted at ``state``. + + Labelling a state means finding a plan *from* it, and the planners all + start at ``task.initial_state()``. Rather than thread a start state through + fourteen planner signatures, move the task: everything else — operators, + goals, axioms, the numeric layer — is unchanged, so every planner and every + heuristic works on the result unmodified. + """ + return _replace(task, init=facts_of(state), init_values=values_of(state)) + + +@dataclass +class Sample: + """One state, its feature vector, and the cost-to-go we believe it has.""" + + features: list + target: float + #: False when the plan it came from was only satisficing, so ``target`` is + #: an upper bound. Training can down-weight these rather than discard them. + optimal: bool = True + instance: str = "" + + +@dataclass +class RankingGroup: + """Successors of one state: the one the plan took, and the ones it did not. + + ``chosen`` and ``others`` hold ``(features, step_cost)``. The step cost has + to travel with the vector because the comparison GBFS makes at expansion + time is between ``h(s')`` values, but the comparison that is *correct* is + between ``c(s, s') + h*(s')`` — with non-uniform action costs those differ. + """ + + chosen: tuple + others: list = field(default_factory=list) + instance: str = "" + + +class Corpus: + """Regression samples and ranking groups over a shared feature space.""" + + def __init__(self, space: FeatureSpace, samples=None, groups=None): + self.space = space + self.samples = list(samples or ()) + self.groups = list(groups or ()) + + def __len__(self) -> int: + return len(self.samples) + + def extend(self, other: "Corpus") -> None: + if other.space != self.space: + raise ValueError("cannot merge corpora with different feature spaces") + self.samples.extend(other.samples) + self.groups.extend(other.groups) + + def split(self, validation: float = 0.2, seed: int = 0): + """Hold out a fraction of the samples, grouped so nothing leaks. + + The split is by *instance*, not by sample. Two states three steps apart + on the same plan have nearly identical features and nearly identical + targets; splitting at random puts one in train and one in validation + and reports a validation score that measures memorisation. + """ + instances = sorted({s.instance for s in self.samples}) + if len(instances) < 2: + # Nothing to hold out by instance; fall back to a sample split and + # accept that the score is optimistic. + rng = random.Random(seed) + order = list(range(len(self.samples))) + rng.shuffle(order) + cut = max(1, int(len(order) * (1 - validation))) + train_idx = set(order[:cut]) + train = [s for i, s in enumerate(self.samples) if i in train_idx] + val = [s for i, s in enumerate(self.samples) if i not in train_idx] + return ( + Corpus(self.space, train, self.groups), + Corpus(self.space, val, []), + ) + rng = random.Random(seed) + rng.shuffle(instances) + cut = max(1, int(len(instances) * (1 - validation))) + held = set(instances[cut:]) + train = [s for s in self.samples if s.instance not in held] + val = [s for s in self.samples if s.instance in held] + groups = [g for g in self.groups if g.instance not in held] + return Corpus(self.space, train, groups), Corpus(self.space, val, []) + + def target_stats(self) -> dict: + if not self.samples: + return {"count": 0} + targets = [s.target for s in self.samples] + mean = sum(targets) / len(targets) + var = sum((t - mean) ** 2 for t in targets) / len(targets) + return { + "count": len(targets), + "groups": len(self.groups), + "mean": mean, + "std": math.sqrt(var), + "min": min(targets), + "max": max(targets), + "optimal_fraction": sum(s.optimal for s in self.samples) / len(targets), + } + + # -- serialisation ----------------------------------------------------- + def to_dict(self) -> dict: + return { + "space": self.space.to_dict(), + "samples": [ + { + "features": s.features, + "target": s.target, + "optimal": s.optimal, + "instance": s.instance, + } + for s in self.samples + ], + "groups": [ + { + "chosen": [list(g.chosen[0]), g.chosen[1]], + "others": [[list(f), c] for f, c in g.others], + "instance": g.instance, + } + for g in self.groups + ], + } + + @classmethod + def from_dict(cls, data: dict) -> "Corpus": + space = FeatureSpace.from_dict(data["space"]) + samples = [ + Sample( + s["features"], + s["target"], + s.get("optimal", True), + s.get("instance", ""), + ) + for s in data["samples"] + ] + groups = [ + RankingGroup( + (g["chosen"][0], g["chosen"][1]), + [(f, c) for f, c in g["others"]], + g.get("instance", ""), + ) + for g in data.get("groups", []) + ] + return cls(space, samples, groups) + + def save(self, path: str) -> None: + with open(path, "w", encoding="utf-8") as handle: + json.dump(self.to_dict(), handle) + + @classmethod + def load(cls, path: str) -> "Corpus": + with open(path, encoding="utf-8") as handle: + return cls.from_dict(json.load(handle)) + + +def samples_from_plan( + task, + plan, + bound, + instance: str = "", + optimal: bool = True, + ranking: bool = True, + max_siblings: int = 8, + rng: Optional[random.Random] = None, + start_state=None, +): + """Regression samples and ranking groups along one plan. + + ``bound`` is a :class:`~jupyddl.learn.features.BoundFeatures` for ``task``. + ``start_state`` is where the plan begins, defaulting to the task's initial + state; the aggregation stages pass a state the search wandered into and a + plan found from there. + + Replaying the plan through :meth:`~jupyddl.task.Task.apply` rather than + ``operator.apply`` is not optional: it is what closes derived predicates and + carries numeric fluents, and a state missing its derived facts has different + features from the one the planner will actually see. + """ + rng = rng or random.Random(0) + plan = list(plan or ()) + if not plan: + return [], [] + + states = [task.initial_state() if start_state is None else task.close(start_state)] + for operator in plan: + states.append(task.apply(operator, states[-1])) + + suffix = [0.0] * (len(plan) + 1) + for i in range(len(plan) - 1, -1, -1): + suffix[i] = suffix[i + 1] + plan[i].cost + + samples = [ + Sample(bound(state), suffix[i], optimal, instance) + for i, state in enumerate(states) + ] + + groups = [] + if ranking: + for i, operator in enumerate(plan): + state = states[i] + chosen = (bound(states[i + 1]), float(operator.cost)) + siblings = [] + for other in task.operators: + if other is operator or not other.applicable(state): + continue + successor = task.apply(other, state) + if facts_of(successor) == facts_of(states[i + 1]): + continue # a different operator reaching the same state + siblings.append((successor, float(other.cost))) + if not siblings: + continue + # A branching factor in the hundreds would make the ranking loss + # dominate the epoch; a sample of the siblings carries the signal. + if len(siblings) > max_siblings: + siblings = rng.sample(siblings, max_siblings) + groups.append( + RankingGroup( + chosen, + [(bound(s), c) for s, c in siblings], + instance, + ) + ) + return samples, groups + + +def build_corpus( + tasks, + solver, + space: Optional[FeatureSpace] = None, + ranking: bool = True, + optimal: bool = True, + seed: int = 0, + on_instance=None, +): + """Solve every task and collect what the plans teach. + + ``tasks`` is an iterable of ``(name, task)`` and ``solver`` maps a task to a + :class:`~jupyddl.search.SearchResult`. ``optimal`` states whether that + solver is cost-optimal, and must be told rather than guessed: a + ``SearchResult`` records what was found, not whether anything cheaper + exists, so nothing in the result distinguishes an optimal plan from a + satisficing one. Getting it wrong mislabels upper bounds as ``h*``. + ``on_instance`` is called with ``(name, result)`` after each solve, for + progress reporting. + """ + tasks = list(tasks) + if space is None: + space = FeatureSpace.from_tasks(task for _, task in tasks) + corpus = Corpus(space) + rng = random.Random(seed) + for name, task in tasks: + result = solver(task) + if on_instance is not None: + on_instance(name, result) + if not result.solved or not result.plan: + continue + samples, groups = samples_from_plan( + task, + result.plan, + space.bind(task), + instance=name, + optimal=optimal, + ranking=ranking, + rng=rng, + ) + corpus.samples.extend(samples) + corpus.groups.extend(groups) + return corpus diff --git a/jupyddl/learn/features.py b/jupyddl/learn/features.py new file mode 100644 index 0000000..a0ef39f --- /dev/null +++ b/jupyddl/learn/features.py @@ -0,0 +1,214 @@ +"""Turn a grounded state into a fixed-length feature vector. + +A learned heuristic is only interesting if it *transfers* — trained on small +instances, useful on large ones. That rules out the obvious encoding. A +one-hot over ``task.facts`` has a different length for every instance and +attaches meaning to fact ids that are an artefact of grounding order, so a +model trained on 4-block blocksworld cannot even be evaluated on 8 blocks. + +Everything here is therefore keyed on the **predicate symbol** rather than the +grounded atom, and normalised by how many atoms of that symbol exist. Both +choices are what make the vector size-invariant: ``on`` contributes one feature +whether there are four blocks or forty, and its value stays in ``[0, 1]``. + +This is the cheap end of a spectrum. The expensive end is a graph network over +the grounded problem (STRIPS-HGN, ASNets, GOOSE), which captures object +identity and relational structure this cannot. The trade is deliberate: a +heuristic that is slower to evaluate than ``hff`` has to be *much* better +informed to win on time, and in a pure-Python planner it will not be. See +``.docs/learned-heuristics.md``. + +Cost is ``O(|s| + |goals|)`` per state, with no successor generation. +""" + +from __future__ import annotations + +import json +import math +import re + +from ..task import facts_of + +__all__ = ["FeatureSpace", "predicate_of"] + +# Grounded facts print Lisp-style — ``(on a b)``, ``(handempty)`` — while +# operators print functionally, ``move(a,b)#2``. Accept both rather than +# assuming: getting this wrong does not raise, it silently gives every ground +# atom its own slot, and the feature vector stops being size-invariant while +# still looking perfectly reasonable. +_SYMBOL = re.compile(r"^\(?\s*([^\s()]+)") + + +def predicate_of(fact_name: str) -> str: + """The predicate symbol of a grounded fact string. + + >>> predicate_of("(on b1 b2)") + 'on' + >>> predicate_of("(handempty)") + 'handempty' + >>> predicate_of("move(a,b)") + 'move' + """ + match = _SYMBOL.match(fact_name) + if not match: + return fact_name + return match.group(1).split("(")[0] + + +class FeatureSpace: + """Maps a state of one task into a vector comparable across tasks. + + The vocabulary — which predicate symbols get which slot — is fixed at + construction and travels with the trained model. A task using a symbol the + vocabulary has never seen contributes nothing rather than shifting every + other feature along, which is what lets one model serve a whole domain and + degrade gracefully on a related one. + """ + + #: Features that do not belong to any single predicate. + GLOBAL_FEATURES = ( + "goal_unsatisfied_fraction", + "goal_unsatisfied_log", + "state_size_fraction", + "goal_satisfied_any", + ) + + def __init__(self, vocabulary): + self.vocabulary = tuple(vocabulary) + self._slot = {name: i for i, name in enumerate(self.vocabulary)} + + # -- construction ------------------------------------------------------ + @classmethod + def from_task(cls, task) -> "FeatureSpace": + """Vocabulary from one task's predicate symbols.""" + return cls.from_tasks([task]) + + @classmethod + def from_tasks(cls, tasks) -> "FeatureSpace": + """Vocabulary from several tasks, so one model covers all of them. + + Symbols are sorted rather than encounter-ordered: the vocabulary has to + be identical whichever order the training instances arrive in, or two + runs over the same corpus produce models that disagree. + """ + symbols = set() + for task in tasks: + for name in task.facts: + symbols.add(predicate_of(name)) + return cls(sorted(symbols)) + + @property + def size(self) -> int: + return 2 * len(self.vocabulary) + len(self.GLOBAL_FEATURES) + + def names(self) -> list: + """Human-readable feature names, in vector order (for inspection).""" + rows = [f"true:{p}" for p in self.vocabulary] + rows += [f"open-goal:{p}" for p in self.vocabulary] + rows += list(self.GLOBAL_FEATURES) + return rows + + # -- binding to a task ------------------------------------------------- + def bind(self, task) -> "BoundFeatures": + """Pre-compute everything about ``task`` that does not vary by state. + + Per-state work then reduces to two counting passes. Binding once per + task and reusing it is the difference between a heuristic that is cheap + and one that re-derives the whole predicate table on every node. + """ + return BoundFeatures(self, task) + + # -- serialisation ----------------------------------------------------- + def to_dict(self) -> dict: + return {"vocabulary": list(self.vocabulary)} + + @classmethod + def from_dict(cls, data: dict) -> "FeatureSpace": + return cls(data["vocabulary"]) + + def __eq__(self, other) -> bool: + return isinstance(other, FeatureSpace) and self.vocabulary == other.vocabulary + + def __repr__(self) -> str: # pragma: no cover - debugging aid + return f"FeatureSpace({len(self.vocabulary)} predicates, {self.size} features)" + + +class BoundFeatures: + """A :class:`FeatureSpace` specialised to one task.""" + + def __init__(self, space: FeatureSpace, task): + self.space = space + self.task = task + width = len(space.vocabulary) + + # fact id -> vocabulary slot, or -1 for a symbol outside the vocabulary. + self.slot_of = tuple( + space._slot.get(predicate_of(name), -1) for name in task.facts + ) + + # Per-symbol totals, used as denominators. A symbol the task never + # grounds keeps a denominator of 1 so its feature is a constant zero + # rather than a division by zero. + totals = [0] * width + for slot in self.slot_of: + if slot >= 0: + totals[slot] += 1 + self.fact_totals = tuple(t or 1 for t in totals) + + goal_totals = [0] * width + for fact in task.goals: + slot = self.slot_of[fact] + if slot >= 0: + goal_totals[slot] += 1 + self.goal_totals = tuple(t or 1 for t in goal_totals) + + self.goals = frozenset(task.goals) + self.num_goals = len(self.goals) or 1 + self.num_facts = len(task.facts) or 1 + self.width = width + self.size = space.size + + def __call__(self, state) -> list: + """The feature vector for ``state``.""" + width = self.width + slot_of = self.slot_of + vector = [0.0] * self.size + + facts = facts_of(state) + for fact in facts: + slot = slot_of[fact] + if slot >= 0: + vector[slot] += 1.0 + + open_goals = 0 + for fact in self.goals: + if fact not in facts: + open_goals += 1 + slot = slot_of[fact] + if slot >= 0: + vector[width + slot] += 1.0 + + fact_totals = self.fact_totals + goal_totals = self.goal_totals + for i in range(width): + vector[i] /= fact_totals[i] + vector[width + i] /= goal_totals[i] + + base = 2 * width + vector[base] = open_goals / self.num_goals + # Counts span orders of magnitude across a scaling ladder; the log keeps + # the tail from swamping every other feature. + vector[base + 1] = math.log1p(open_goals) + vector[base + 2] = len(facts) / self.num_facts + vector[base + 3] = 1.0 if open_goals < self.num_goals else 0.0 + return vector + + +def save_space(space: FeatureSpace, path: str) -> None: # pragma: no cover - thin + with open(path, "w", encoding="utf-8") as handle: + json.dump(space.to_dict(), handle) + + +def load_space(path: str) -> FeatureSpace: # pragma: no cover - thin + with open(path, encoding="utf-8") as handle: + return FeatureSpace.from_dict(json.load(handle)) diff --git a/jupyddl/learn/heuristic.py b/jupyddl/learn/heuristic.py new file mode 100644 index 0000000..3c7e4c8 --- /dev/null +++ b/jupyddl/learn/heuristic.py @@ -0,0 +1,116 @@ +"""The learned heuristic itself, and the artefact it is loaded from. + +A trained heuristic is three things that must travel together: the feature +vocabulary, the network, and the scale its targets were normalised by. Ship the +network alone and it silently mis-predicts on the first task whose predicate +symbols hash to different slots — no error, just a bad heuristic, which is the +hardest kind of bug to notice because the planner still returns plans. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field + +from ..heuristics.base import Heuristic +from .features import FeatureSpace +from .model import MLP + +__all__ = ["HeuristicBundle", "LearnedHeuristic"] + +FORMAT_VERSION = 1 + + +@dataclass +class HeuristicBundle: + """Everything needed to evaluate a learned heuristic on a fresh task.""" + + space: FeatureSpace + model: MLP + #: Targets are trained in units of ``scale``; predictions are multiplied + #: back. Regressing raw costs that range over two orders of magnitude + #: leaves the first layer fighting the output layer for the scale. + scale: float = 1.0 + metrics: dict = field(default_factory=dict) + config: dict = field(default_factory=dict) + + def to_dict(self) -> dict: + return { + "format": FORMAT_VERSION, + "space": self.space.to_dict(), + "model": self.model.to_dict(), + "scale": self.scale, + "metrics": self.metrics, + "config": self.config, + } + + @classmethod + def from_dict(cls, data: dict) -> "HeuristicBundle": + version = data.get("format", 0) + if version != FORMAT_VERSION: + raise ValueError( + f"unsupported learned-heuristic format {version}; " + f"this build reads version {FORMAT_VERSION}" + ) + return cls( + FeatureSpace.from_dict(data["space"]), + MLP.from_dict(data["model"]), + float(data.get("scale", 1.0)), + dict(data.get("metrics", {})), + dict(data.get("config", {})), + ) + + def save(self, path: str) -> None: + with open(path, "w", encoding="utf-8") as handle: + json.dump(self.to_dict(), handle, indent=1) + + @classmethod + def load(cls, path: str) -> "HeuristicBundle": + with open(path, encoding="utf-8") as handle: + return cls.from_dict(json.load(handle)) + + def bind(self, task) -> "LearnedHeuristic": + return LearnedHeuristic(task, self) + + +class LearnedHeuristic(Heuristic): + """Evaluates a trained network as a heuristic. + + Never admissible, and it does not pretend otherwise: nothing in the + training objective bounds the prediction from above, and a single + over-estimate costs A* its optimality guarantee. Pair it with ``gbfs``, or + with ``wastar`` if you want the bounded-suboptimality knob. + + Values are cached per state. The planner already caches heuristic values in + its open-list bookkeeping, but a re-opened state is re-evaluated there, and + a network evaluation is far more expensive than a dictionary lookup. + """ + + name = "learned" + admissible = False + + def __init__(self, task, bundle: HeuristicBundle): + super().__init__(task) + self.bundle = bundle + self.bound = bundle.space.bind(task) + self.scale = bundle.scale + self._model = bundle.model + self._cache: dict = {} + self.evaluations = 0 + + def __call__(self, state) -> float: + cached = self._cache.get(state) + if cached is not None: + return cached + self.evaluations += 1 + value = self._model(self.bound(state)) * self.scale + # A goal state must score zero however the network was trained; the + # planner's termination test does not consult the heuristic, but a + # non-zero goal estimate distorts every f-value near the goal. + if self.task.goal_reached(state): + value = 0.0 + self._cache[state] = value + return value + + def __repr__(self) -> str: # pragma: no cover - debugging aid + return f"LearnedHeuristic({self._model!r}, scale={self.scale:.3g})" diff --git a/jupyddl/learn/model.py b/jupyddl/learn/model.py new file mode 100644 index 0000000..5472cc5 --- /dev/null +++ b/jupyddl/learn/model.py @@ -0,0 +1,358 @@ +"""A small feed-forward network, trainable without leaving the standard library. + +The core of this package has no runtime dependencies and this module keeps that +promise: the reference implementation of the forward and backward passes is +plain Python lists. When NumPy happens to be importable the same equations run +batched, which is worth one to two orders of magnitude on a real training run — +so ``pip install jupyddl[learn]`` is a speed option, never a requirement. + +The two paths are held together by ``tests/test_learn.py``, which asserts they +produce the same gradients to floating-point tolerance on a fixed seed. A fast +path that has quietly drifted from the reference is worse than no fast path. + +The network is deliberately small. The bet a learned heuristic makes is that it +recovers more search than it costs to evaluate, and it is evaluated once per +generated state — tens of thousands of times a second. A wide network loses +that bet before it has said anything. +""" + +from __future__ import annotations + +import json +import math +import random +from typing import Optional + +try: # pragma: no cover - exercised by whichever path is installed + import numpy as _np +except ImportError: # pragma: no cover + _np = None + +__all__ = ["MLP", "numpy_available"] + + +def numpy_available() -> bool: + """Whether the batched fast path is usable in this environment.""" + return _np is not None + + +def _he_scale(fan_in: int) -> float: + """He initialisation: keeps activation variance stable through ReLU.""" + return math.sqrt(2.0 / max(1, fan_in)) + + +class MLP: + """A ReLU multilayer perceptron with a non-negative scalar output. + + The output passes through softplus. A heuristic that returns a negative + cost-to-go is not merely inaccurate, it breaks the assumptions of every + planner consuming it — greedy search will chase the negative values and A* + loses its admissibility argument outright. Clamping with ``max(0, x)`` + would do it too, but it has zero gradient below the threshold, so a unit + that lands there early never recovers. + """ + + def __init__(self, sizes, seed: int = 0, weights=None, biases=None): + self.sizes = tuple(sizes) + if len(self.sizes) < 2: + raise ValueError("an MLP needs at least an input and an output size") + if self.sizes[-1] != 1: + raise ValueError("a heuristic model has a scalar output") + if weights is None: + rng = random.Random(seed) + self.weights = [ + [ + [rng.gauss(0.0, _he_scale(fan_in)) for _ in range(fan_in)] + for _ in range(fan_out) + ] + for fan_in, fan_out in zip(self.sizes, self.sizes[1:]) + ] + self.biases = [[0.0] * fan_out for fan_out in self.sizes[1:]] + else: + self.weights = [[list(row) for row in layer] for layer in weights] + self.biases = [list(b) for b in biases] + self._np_cache: Optional[tuple] = None + + # -- shape helpers ----------------------------------------------------- + @property + def input_size(self) -> int: + return self.sizes[0] + + @property + def num_parameters(self) -> int: + return sum(len(row) for layer in self.weights for row in layer) + sum( + len(b) for b in self.biases + ) + + def get_flat(self) -> list: + """Every parameter as one vector (what the derivative-free stage tunes).""" + flat = [] + for layer in self.weights: + for row in layer: + flat.extend(row) + for bias in self.biases: + flat.extend(bias) + return flat + + def set_flat(self, flat) -> None: + """Inverse of :meth:`get_flat`.""" + index = 0 + for layer in self.weights: + for row in layer: + width = len(row) + row[:] = flat[index : index + width] + index += width + for bias in self.biases: + width = len(bias) + bias[:] = flat[index : index + width] + index += width + if index != len(flat): + raise ValueError( + f"expected {self.num_parameters} parameters, got {len(flat)}" + ) + self._np_cache = None + + # -- inference --------------------------------------------------------- + def __call__(self, vector) -> float: + """Predict for a single feature vector. + + This is the path the planner takes, once per generated state, so it + stays a plain loop: for one sample the NumPy call overhead exceeds the + arithmetic it would save. + """ + activation = vector + last = len(self.weights) - 1 + for index, (layer, bias) in enumerate(zip(self.weights, self.biases)): + out = [] + for row, b in zip(layer, bias): + total = b + for w, a in zip(row, activation): + total += w * a + out.append(total if index == last else (total if total > 0.0 else 0.0)) + activation = out + return _softplus(activation[0]) + + # -- training ---------------------------------------------------------- + def forward_batch(self, batch): + """Predictions plus the activations the backward pass needs. + + Returns ``(outputs, cache)``. ``outputs`` are post-softplus. + """ + if _np is not None: + return self._forward_numpy(batch) + return self._forward_python(batch) + + def backward_batch(self, cache, d_outputs): + """Parameter gradients, given d(loss)/d(output) for each sample. + + The loss lives in the caller. Keeping it there is what lets the same + network be trained against a regression target and against a ranking + objective without the model knowing the difference. + """ + if _np is not None: + return self._backward_numpy(cache, d_outputs) + return self._backward_python(cache, d_outputs) + + # -- reference implementation ----------------------------------------- + def _forward_python(self, batch): + activations = [batch] + pre_activations = [] + current = batch + last = len(self.weights) - 1 + for index, (layer, bias) in enumerate(zip(self.weights, self.biases)): + pre = [] + for sample in current: + row_out = [] + for row, b in zip(layer, bias): + total = b + for w, a in zip(row, sample): + total += w * a + row_out.append(total) + pre.append(row_out) + pre_activations.append(pre) + if index == last: + current = pre + else: + current = [[v if v > 0.0 else 0.0 for v in row] for row in pre] + activations.append(current) + raw = [row[0] for row in current] + outputs = [_softplus(v) for v in raw] + return outputs, (activations, pre_activations, raw) + + def _backward_python(self, cache, d_outputs): + activations, pre_activations, raw = cache + # d(softplus)/dx = sigmoid(x) + delta = [[d * _sigmoid(r)] for d, r in zip(d_outputs, raw)] + + grad_w = [[[0.0] * len(row) for row in layer] for layer in self.weights] + grad_b = [[0.0] * len(bias) for bias in self.biases] + + for index in range(len(self.weights) - 1, -1, -1): + layer = self.weights[index] + inputs = activations[index] + gw = grad_w[index] + gb = grad_b[index] + for sample_delta, sample_input in zip(delta, inputs): + for j, d in enumerate(sample_delta): + if d == 0.0: + continue + gb[j] += d + row = gw[j] + for k, a in enumerate(sample_input): + row[k] += d * a + if index == 0: + break + pre = pre_activations[index - 1] + new_delta = [] + for sample_delta, sample_pre in zip(delta, pre): + back = [0.0] * len(sample_pre) + for j, d in enumerate(sample_delta): + if d == 0.0: + continue + row = layer[j] + for k in range(len(back)): + back[k] += d * row[k] + new_delta.append( + [b if p > 0.0 else 0.0 for b, p in zip(back, sample_pre)] + ) + delta = new_delta + return grad_w, grad_b + + # -- batched fast path ------------------------------------------------- + def _numpy_params(self): + if self._np_cache is None: + self._np_cache = ( + [_np.array(layer, dtype=_np.float64) for layer in self.weights], + [_np.array(bias, dtype=_np.float64) for bias in self.biases], + ) + return self._np_cache + + def _forward_numpy(self, batch): + weights, biases = self._numpy_params() + current = _np.array(batch, dtype=_np.float64) + activations = [current] + pre_activations = [] + last = len(weights) - 1 + for index, (layer, bias) in enumerate(zip(weights, biases)): + pre = current @ layer.T + bias + pre_activations.append(pre) + if index == last: + current = pre + else: + current = _np.maximum(pre, 0.0) + activations.append(current) + raw = current[:, 0] + outputs = _np.logaddexp(0.0, raw) + return outputs.tolist(), (activations, pre_activations, raw) + + def _backward_numpy(self, cache, d_outputs): + weights, _ = self._numpy_params() + activations, pre_activations, raw = cache + d_out = _np.array(d_outputs, dtype=_np.float64) + delta = (d_out / (1.0 + _np.exp(-raw)))[:, None] + + grad_w = [None] * len(weights) + grad_b = [None] * len(weights) + for index in range(len(weights) - 1, -1, -1): + inputs = activations[index] + grad_w[index] = (delta.T @ inputs).tolist() + grad_b[index] = delta.sum(axis=0).tolist() + if index == 0: + break + back = delta @ weights[index] + delta = back * (pre_activations[index - 1] > 0.0) + return grad_w, grad_b + + # -- serialisation ----------------------------------------------------- + def to_dict(self) -> dict: + return { + "sizes": list(self.sizes), + "weights": [[list(row) for row in layer] for layer in self.weights], + "biases": [list(b) for b in self.biases], + } + + @classmethod + def from_dict(cls, data: dict) -> "MLP": + return cls(data["sizes"], weights=data["weights"], biases=data["biases"]) + + def copy(self) -> "MLP": + return MLP(self.sizes, weights=self.weights, biases=self.biases) + + def __repr__(self) -> str: # pragma: no cover - debugging aid + shape = "x".join(str(s) for s in self.sizes) + return f"MLP({shape}, {self.num_parameters} parameters)" + + +class Adam: + """Adam, with the usual defaults, over the nested weight/bias structure.""" + + def __init__(self, model: MLP, lr: float = 0.01, betas=(0.9, 0.999), eps=1e-8): + self.model = model + self.lr = lr + self.beta1, self.beta2 = betas + self.eps = eps + self.step_count = 0 + self.m_w = [[[0.0] * len(row) for row in layer] for layer in model.weights] + self.v_w = [[[0.0] * len(row) for row in layer] for layer in model.weights] + self.m_b = [[0.0] * len(b) for b in model.biases] + self.v_b = [[0.0] * len(b) for b in model.biases] + + def step(self, grad_w, grad_b) -> None: + self.step_count += 1 + bias1 = 1.0 - self.beta1**self.step_count + bias2 = 1.0 - self.beta2**self.step_count + for layer_index, layer in enumerate(self.model.weights): + gw = grad_w[layer_index] + mw = self.m_w[layer_index] + vw = self.v_w[layer_index] + for j, row in enumerate(layer): + grad_row = gw[j] + m_row = mw[j] + v_row = vw[j] + for k in range(len(row)): + g = grad_row[k] + m_row[k] = self.beta1 * m_row[k] + (1 - self.beta1) * g + v_row[k] = self.beta2 * v_row[k] + (1 - self.beta2) * g * g + row[k] -= ( + self.lr + * (m_row[k] / bias1) + / (math.sqrt(v_row[k] / bias2) + self.eps) + ) + for layer_index, bias in enumerate(self.model.biases): + gb = grad_b[layer_index] + mb = self.m_b[layer_index] + vb = self.v_b[layer_index] + for j in range(len(bias)): + g = gb[j] + mb[j] = self.beta1 * mb[j] + (1 - self.beta1) * g + vb[j] = self.beta2 * vb[j] + (1 - self.beta2) * g * g + bias[j] -= ( + self.lr * (mb[j] / bias1) / (math.sqrt(vb[j] / bias2) + self.eps) + ) + self.model._np_cache = None + + +def _softplus(x: float) -> float: + # log(1 + e^x), written so a large x does not overflow. + if x > 30.0: + return x + if x < -30.0: + return math.exp(x) + return math.log1p(math.exp(x)) + + +def _sigmoid(x: float) -> float: + if x >= 0.0: + return 1.0 / (1.0 + math.exp(-x)) + z = math.exp(x) + return z / (1.0 + z) + + +def save_model(model: MLP, path: str) -> None: # pragma: no cover - thin + with open(path, "w", encoding="utf-8") as handle: + json.dump(model.to_dict(), handle) + + +def load_model(path: str) -> MLP: # pragma: no cover - thin + with open(path, encoding="utf-8") as handle: + return MLP.from_dict(json.load(handle)) diff --git a/jupyddl/learn/pipeline.py b/jupyddl/learn/pipeline.py new file mode 100644 index 0000000..0831c59 --- /dev/null +++ b/jupyddl/learn/pipeline.py @@ -0,0 +1,296 @@ +"""One call from a domain name to a trained heuristic. + +The stages are useful separately and tedious to wire together, so this is the +assembled version: generate a ladder of instances, solve the small ones +optimally, imitate those plans, then let the reinforcement stages take over. + +The default ladder deliberately trains on instances *smaller* than the ones it +will be used on. That is the whole claim being tested — a heuristic that only +works at the size it was trained on has learned the instance, not the domain — +and :func:`jupyddl.learn.pipeline.evaluate_transfer` is what checks it. +""" + +from __future__ import annotations + +import time +from typing import Optional + +from ..api import solve_task +from ..generator import generate +from ..grounding import ground +from ..parser import parse +from ..search import make_planner +from ..search.result import make_budget +from .dataset import build_corpus +from .features import FeatureSpace +from .heuristic import HeuristicBundle +from .rl import RLConfig, bootstrap, dagger, optimise_search_cost +from .train import TrainConfig, train + +__all__ = [ + "tasks_from_generator", + "solved_corpus", + "learn_heuristic", + "evaluate_transfer", +] + + +def tasks_from_generator(kind: str, sizes, seed: int = 0, seeds_per_size: int = 1): + """Ground a ladder of generated instances into ``(name, task)`` pairs. + + Several seeds per size is usually a better use of a training budget than + more sizes: what varies between two 6-block instances is the goal + structure, which is what has to be learned, whereas what varies between 6 + and 7 blocks is mostly scale, which the features already normalise away. + """ + tasks = [] + for size in sizes: + for offset in range(seeds_per_size): + instance_seed = seed + offset + domain_text, problem_text = generate(kind, size=size, seed=instance_seed) + task = ground(parse(domain_text), parse(problem_text)) + tasks.append((f"{kind}-{size:02d}-{instance_seed}", task)) + return tasks + + +def solved_corpus( + tasks, + space: Optional[FeatureSpace] = None, + planner: str = "astar", + heuristic: str = "lmcut", + optimal: bool = True, + max_expansions: Optional[int] = 20000, + time_limit: Optional[float] = 30.0, + ranking: bool = True, + seed: int = 0, + on_instance=None, +): + """Solve every task and turn the plans into a corpus. + + Defaults to an optimal configuration, so the targets really are ``h*``. On + a ladder where that is out of reach, pass ``planner="gbfs"``, + ``heuristic="hff"`` and ``optimal=False``: the labels become upper bounds, + training down-weights them, and that is far better than an empty corpus. + """ + + def solver(task): + return solve_task( + task, + search=planner, + heuristic=heuristic, + max_expansions=max_expansions, + time_limit=time_limit, + ) + + return build_corpus( + tasks, + solver, + space=space, + ranking=ranking, + optimal=optimal, + seed=seed, + on_instance=on_instance, + ) + + +def learn_heuristic( + kind: str, + sizes=range(3, 7), + seeds_per_size: int = 2, + seed: int = 0, + train_config: Optional[TrainConfig] = None, + rl_config: Optional[RLConfig] = None, + dagger_rounds: int = 0, + bootstrap_sizes=None, + cem_iterations: int = 0, + cem_sizes=None, + verbose: bool = False, +): + """Generate, solve, imitate and (optionally) reinforce. Returns a bundle. + + Every stage after imitation is off by default. They cost minutes rather + than seconds and they are not always worth it — which one pays depends on + the domain, and saying so honestly is more useful than a default that + quietly triples the runtime. + + ``cem_sizes`` is the one setting worth understanding before using it. The + direct search-cost stage needs instances with *headroom*: on the training + ladder the imitated heuristic already expands roughly as many nodes as the + plan is long, so every perturbation scores the same and the objective is + flat. Measured on blocksworld, tuning on the training sizes moved the score + from 12.8 to 12.75 — noise — while tuning on sizes 9-12 moved it from 1605 + to 64, and cut a held-out set the optimiser never saw from 336 to 122. + Default: one rung above the training ladder. + """ + train_config = train_config or TrainConfig(seed=seed) + rl_config = rl_config or RLConfig(seed=seed, verbose=verbose) + started = time.perf_counter() + log = [] + + def note(message): + log.append(message) + if verbose: # pragma: no cover - operator convenience + print(message) + + tasks = tasks_from_generator(kind, sizes, seed=seed, seeds_per_size=seeds_per_size) + note(f"ladder: {len(tasks)} instances of {kind}, sizes {list(sizes)}") + + # The vocabulary must span every task the model will ever see, including + # the larger ones it is only evaluated on: a symbol first encountered at + # evaluation time would otherwise have no slot. + sizes = list(sizes) + if cem_iterations and cem_sizes is None: + # A rung above the training ladder: far enough that search has room to + # be improved, near enough that it can still be searched cheaply. + top = max(sizes) + cem_sizes = range(top + 2, top + 6) + + all_tasks = list(tasks) + if bootstrap_sizes: + all_tasks += tasks_from_generator( + kind, bootstrap_sizes, seed=seed, seeds_per_size=1 + ) + if cem_sizes: + # Two disjoint seed families over the same sizes: one the optimiser + # fits, one it is only ever scored on. Without the split, a thousand + # parameters tuned against eight instances fit those eight instances + # and transfer gets worse, not better. + cem_tasks = tasks_from_generator( + kind, cem_sizes, seed=seed + 1000, seeds_per_size=2 + ) + cem_validation = tasks_from_generator( + kind, cem_sizes, seed=seed + 2000, seeds_per_size=2 + ) + all_tasks += cem_tasks + cem_validation + space = FeatureSpace.from_tasks(task for _, task in all_tasks) + + corpus = solved_corpus(tasks, space=space, seed=seed) + note(f"corpus: {corpus.target_stats()}") + if not corpus.samples: + raise RuntimeError( + f"no instance of '{kind}' was solved optimally within the budget; " + "lower the sizes or pass a satisficing configuration" + ) + + bundle, report = train(corpus, train_config) + note( + f"imitation: {report.metrics.get('mae', 0):.2f} MAE, " + f"top-1 {report.metrics.get('top1', 0):.3f}, " + f"best epoch {report.best_epoch}, {report.seconds:.1f}s" + ) + + if dagger_rounds: + bundle, corpus, history = dagger( + bundle, + tasks, + corpus, + rounds=dagger_rounds, + config=rl_config, + train_config=train_config, + ) + note(f"dagger: {history[-1] if history else 'no rounds ran'}") + + if bootstrap_sizes: + ladder = tasks_from_generator( + kind, bootstrap_sizes, seed=seed, seeds_per_size=1 + ) + bundle, corpus, history = bootstrap( + bundle, + ladder, + corpus, + config=rl_config, + train_config=train_config, + ) + note(f"bootstrap: {history[-1] if history else 'no rounds ran'}") + + if cem_iterations: + note( + f"cem: tuning on {len(cem_tasks)} instances of sizes {list(cem_sizes)}, " + f"selecting on {len(cem_validation)} held out from it" + ) + bundle, history = optimise_search_cost( + bundle, + cem_tasks, + iterations=cem_iterations, + config=rl_config, + validation_tasks=cem_validation, + ) + note(f"cem: {history[-1] if history else 'no iterations ran'}") + + bundle.metrics["pipeline"] = { + "kind": kind, + "sizes": list(sizes), + "instances": len(tasks), + "seconds": time.perf_counter() - started, + "log": log, + } + return bundle + + +def evaluate_transfer( + bundle: HeuristicBundle, + tasks, + baselines=("hff", "goalcount", "blind"), + planner: str = "gbfs", + max_expansions: Optional[int] = 20000, + time_limit: Optional[float] = 30.0, +): + """Compare the learned heuristic against baselines on ``tasks``. + + Reports expansions *and* wall-clock, because they can disagree and the + disagreement is the whole story. A learned heuristic that halves expansions + while tripling time per node has not helped anyone; one that expands + slightly more than ``hff`` but evaluates in a fraction of the time may + still win. Only reporting expansions is how a learned heuristic gets + published as a success without being one. + """ + from ..heuristics import make_heuristic + + rows = [] + for name, task in tasks: + configurations = [("learned", lambda t: bundle.bind(t))] + configurations += [ + (base, (lambda b: lambda t: make_heuristic(b, t))(base)) + for base in baselines + ] + for label, factory in configurations: + planner_instance = make_planner(planner) + budget = make_budget(max_expansions, time_limit) + started = time.perf_counter() + result = planner_instance.search(task, factory(task), budget=budget) + rows.append( + { + "instance": name, + "heuristic": label, + "solved": result.solved, + "cost": result.cost, + "expanded": result.stats.expanded, + "evaluated": result.stats.evaluated, + "seconds": time.perf_counter() - started, + "truncated": result.truncated, + } + ) + return rows + + +def summarise_transfer(rows) -> dict: + """Per-heuristic coverage, mean expansions and mean time over solved runs.""" + summary: dict = {} + for row in rows: + entry = summary.setdefault( + row["heuristic"], + {"solved": 0, "total": 0, "expanded": 0, "seconds": 0.0, "cost": 0}, + ) + entry["total"] += 1 + if row["solved"]: + entry["solved"] += 1 + entry["expanded"] += row["expanded"] + entry["seconds"] += row["seconds"] + entry["cost"] += row["cost"] or 0 + for entry in summary.values(): + solved = entry["solved"] or 1 + entry["mean_expanded"] = entry["expanded"] / solved + entry["mean_seconds"] = entry["seconds"] / solved + entry["mean_cost"] = entry["cost"] / solved + entry["coverage"] = entry["solved"] / (entry["total"] or 1) + return summary diff --git a/jupyddl/learn/rl.py b/jupyddl/learn/rl.py new file mode 100644 index 0000000..2a9bc86 --- /dev/null +++ b/jupyddl/learn/rl.py @@ -0,0 +1,472 @@ +"""The reinforcement stage: stop imitating h*, start minimising search. + +Imitation optimises a proxy. What we want is a heuristic that makes GBFS expand +few nodes, and "nodes expanded" is not a differentiable function of the +weights — it comes out the far side of a priority queue, a goal test and a +successor generator. Three mechanisms close that gap, and they fix different +things. + +**DAgger** (:func:`dagger`) fixes *covariate shift*. Imitation trains on states +lying on optimal plans; search asks about states that do not, including the +dead ends and detours a mediocre heuristic wanders into. Training on your own +search distribution is the standard fix (Ross et al., 2011). Concretely: run +the current heuristic, keep the states it expanded, label them by solving from +each one, retrain on the union. + +**Bootstrapping** (:func:`bootstrap`) fixes *data scarcity at the top of the +ladder*. You cannot label a 20-block instance you cannot solve. But a heuristic +trained on 6 blocks may just crack 8, whose plans then teach it 10 (Arfaee, +Zilles & Holte, 2011). Each round is a policy-improvement step where the policy +is "which instances can I solve at all". + +**Direct search-cost optimisation** (:func:`optimise_search_cost`) attacks the +objective itself, with a derivative-free method, because no gradient exists to +follow. The cross-entropy method over the weight vector treats the planner as a +black box returning expansions. Evolutionary strategies are a real alternative +to policy gradients when rollouts are cheap and the parameter vector is modest +(Salimans et al., 2017), which is exactly this setting. + +The crucial detail for all three: **start from the imitation solution.** Search +cost is a step function of the weights over most of the space — every candidate +that solves nothing scores identically — so a randomly initialised policy gets +no gradient signal, from any method. Imitation is what puts the optimiser +somewhere the objective can distinguish. + +``.docs/rl-for-search.md`` sets out the MDP this corresponds to and why the +obvious policy-gradient formulation is harder than it looks. +""" + +from __future__ import annotations + +import math +import random +import statistics +import time +from dataclasses import dataclass +from typing import Optional + +from ..search import make_planner +from ..search.result import make_budget +from .dataset import Corpus, samples_from_plan +from .heuristic import HeuristicBundle +from .train import TrainConfig, train + +__all__ = [ + "RLConfig", + "SearchCost", + "search_cost", + "dagger", + "bootstrap", + "optimise_search_cost", +] + + +@dataclass +class RLConfig: + """Settings shared by the reinforcement stages.""" + + planner: str = "gbfs" + #: Per-instance expansion budget. Doubles as the penalty for not solving, + #: which is what keeps the objective finite and comparable. + max_expansions: int = 5000 + time_limit: Optional[float] = None + #: A failure is charged this multiple of the budget. Greater than 1 so that + #: solving an instance slowly always beats not solving it — with a factor + #: of exactly 1 the optimiser is indifferent between them. + failure_penalty: float = 2.0 + seed: int = 0 + verbose: bool = False + + +@dataclass +class SearchCost: + """What one evaluation of the objective measured.""" + + score: float + solved: int + total: int + expansions: float + seconds: float + + @property + def coverage(self) -> float: + return self.solved / self.total if self.total else 0.0 + + def to_dict(self) -> dict: + return { + "score": self.score, + "solved": self.solved, + "total": self.total, + "expansions": self.expansions, + "seconds": self.seconds, + "coverage": self.coverage, + } + + +def search_cost(bundle: HeuristicBundle, tasks, config: Optional[RLConfig] = None): + """Run the planner on every task and score the heuristic. Lower is better. + + The score is the mean expansion count, charging an unsolved instance + ``failure_penalty * max_expansions``. Reporting mean expansions over only + the solved instances — the tempting alternative — rewards a heuristic that + solves one easy instance quickly and abandons the rest. + + Because the penalty is a multiple of the budget, scores are comparable + **only at a fixed budget**. Halving ``max_expansions`` halves what a + failure costs, so two runs under different budgets say nothing about each + other. Every optimiser here holds the budget fixed for exactly that reason. + """ + config = config or RLConfig() + tasks = list(tasks) + if not tasks: + raise ValueError("search cost needs at least one task") + planner = make_planner(config.planner) + solved = 0 + total_expansions = 0.0 + started = time.perf_counter() + for _, task in tasks: + heuristic = bundle.bind(task) + budget = make_budget(config.max_expansions, config.time_limit) + result = planner.search(task, heuristic, budget=budget) + if result.solved: + solved += 1 + total_expansions += result.stats.expanded + else: + total_expansions += config.failure_penalty * config.max_expansions + return SearchCost( + score=total_expansions / len(tasks), + solved=solved, + total=len(tasks), + expansions=total_expansions / len(tasks), + seconds=time.perf_counter() - started, + ) + + +# -------------------------------------------------------------------------- +# DAgger: train on the states the search actually visits +# -------------------------------------------------------------------------- +def dagger( + bundle: HeuristicBundle, + tasks, + corpus: Corpus, + labeller=None, + rounds: int = 2, + states_per_task: int = 40, + config: Optional[RLConfig] = None, + train_config: Optional[TrainConfig] = None, +): + """Aggregate data from the current heuristic's own search distribution. + + ``labeller`` maps ``(task, state)`` to a plan from that state, or ``None`` + if it cannot find one; the default solves with greedy search and ``hff``. + Its plans are usually not optimal, so the samples it produces are recorded + as upper bounds and down-weighted during training — an over-estimated label + is still far more informative than no label for a state the imitation + corpus never contained. + + Returns ``(bundle, corpus, history)``. + """ + config = config or RLConfig() + train_config = train_config or TrainConfig() + labeller = labeller or _default_labeller(config) + tasks = list(tasks) + rng = random.Random(config.seed) + history = [] + + for round_index in range(1, rounds + 1): + visited = _collect_visited_states(bundle, tasks, states_per_task, config, rng) + added = 0 + for name, task, state in visited: + plan = labeller(task, state) + if not plan: + continue + samples, groups = samples_from_plan( + task, + plan, + corpus.space.bind(task), + instance=f"{name}#dagger{round_index}", + optimal=False, + ranking=True, + rng=rng, + start_state=state, + ) + corpus.samples.extend(samples) + corpus.groups.extend(groups) + added += len(samples) + bundle, _ = train(corpus, train_config) + cost = search_cost(bundle, tasks, config) + history.append( + { + "round": round_index, + "added_samples": added, + "corpus": len(corpus), + **cost.to_dict(), + } + ) + if config.verbose: # pragma: no cover - operator convenience + print( + f" dagger round {round_index}: +{added} samples, " + f"coverage {cost.coverage:.2f}, expansions {cost.expansions:.0f}" + ) + return bundle, corpus, history + + +def _collect_visited_states(bundle, tasks, states_per_task, config, rng): + """States the current heuristic expands, sampled along each search.""" + from ..trace import SearchObserver + + class _Collector(SearchObserver): + def __init__(self): + self.states = [] + + def on_expand(self, state, **kwargs): + self.states.append(state) + + planner = make_planner(config.planner) + visited = [] + for name, task in tasks: + collector = _Collector() + budget = make_budget(config.max_expansions, config.time_limit) + planner.search(task, bundle.bind(task), observer=collector, budget=budget) + states = collector.states + if not states: + continue + if len(states) > states_per_task: + states = rng.sample(states, states_per_task) + visited.extend((name, task, state) for state in states) + return visited + + +def _default_labeller(config: RLConfig): + """Label a state by solving from it with a fast satisficing configuration. + + The reference heuristic is built once per task and reused across every + state of it. No heuristic in the library reads ``task.init`` — they are all + parameterised by the state they are called on — so the same instance is + valid for every re-rooting of the task, and rebuilding the relaxed-task + tables per label would dominate the cost of labelling. + """ + from ..heuristics import make_heuristic + from .dataset import task_from_state + + # ``Task`` is an unfrozen dataclass, so it is unhashable and cannot key a + # dict directly. Keying on ``id()`` alone would be a trap: CPython reuses + # addresses once an object is collected, so a caller passing tasks this + # cache does not keep alive could get another task's heuristic. Holding the + # task in the value makes the id un-recyclable while the entry lives, and + # the identity check makes it correct even if that reasoning is ever wrong. + cache: dict = {} + + def label(task, state): + key = id(task) + entry = cache.get(key) + if entry is None or entry[0] is not task: + entry = (task, make_heuristic("hff", task)) + cache[key] = entry + reference = entry[1] + planner = make_planner("gbfs") + budget = make_budget(config.max_expansions, config.time_limit) + result = planner.search(task_from_state(task, state), reference, budget=budget) + return result.plan if result.solved else None + + return label + + +# -------------------------------------------------------------------------- +# Bootstrapping: let the heuristic unlock its own training data +# -------------------------------------------------------------------------- +def bootstrap( + bundle: Optional[HeuristicBundle], + ladder, + corpus: Corpus, + rounds: int = 3, + config: Optional[RLConfig] = None, + train_config: Optional[TrainConfig] = None, +): + """Grow the corpus with instances the current heuristic has just cracked. + + ``ladder`` is ``(name, task)`` in increasing difficulty. Each round attempts + every unsolved instance with the current heuristic under a fixed budget, + adds whatever it managed, and retrains. Nothing is ever attempted twice + after it succeeds, so rounds get cheaper as the frontier moves. + + Returns ``(bundle, corpus, history)``. + """ + config = config or RLConfig() + train_config = train_config or TrainConfig() + ladder = list(ladder) + rng = random.Random(config.seed) + remaining = list(ladder) + history = [] + + for round_index in range(1, rounds + 1): + if not remaining: + break + planner = make_planner(config.planner) + newly = [] + still: list = [] + for name, task in remaining: + budget = make_budget(config.max_expansions, config.time_limit) + heuristic = None if bundle is None else bundle.bind(task) + result = planner.search(task, heuristic, budget=budget) + if result.solved and result.plan: + newly.append((name, task, result.plan)) + else: + still.append((name, task)) + for name, task, plan in newly: + samples, groups = samples_from_plan( + task, + plan, + corpus.space.bind(task), + instance=f"{name}#boot{round_index}", + optimal=False, + ranking=True, + rng=rng, + ) + corpus.samples.extend(samples) + corpus.groups.extend(groups) + if newly: + bundle, _ = train(corpus, train_config) + history.append( + { + "round": round_index, + "newly_solved": [name for name, _, _ in newly], + "remaining": len(still), + "corpus": len(corpus), + } + ) + if config.verbose: # pragma: no cover - operator convenience + print( + f" bootstrap round {round_index}: solved {len(newly)}, " + f"{len(still)} left, corpus {len(corpus)}" + ) + if len(still) == len(remaining): + break # the frontier stopped moving; more rounds cost and teach nothing + remaining = still + return bundle, corpus, history + + +# -------------------------------------------------------------------------- +# Direct optimisation of the search cost +# -------------------------------------------------------------------------- +def optimise_search_cost( + bundle: HeuristicBundle, + tasks, + iterations: int = 8, + population: int = 12, + elite_fraction: float = 0.25, + sigma: float = 0.15, + config: Optional[RLConfig] = None, + validation_tasks=None, +): + """Tune the weights against expansions with the cross-entropy method. + + The parameter vector is perturbed, each candidate is scored by actually + planning with it, the best fraction is kept, and the sampling distribution + is refitted to those. No gradient is involved, which is the point: the + objective is the planner. + + **The incumbent is selected on ``validation_tasks``, not on ``tasks``.** A + thousand parameters tuned against eight instances will fit those eight + instances, and this is not hypothetical: a run that reached 108 expansions + on its tuning set scored 1734 on held-out instances, nearly five times + *worse* than the imitated heuristic it started from. Candidates are still + proposed by their tuning score — that is what the sampling distribution + refits on — but nothing is returned unless it also improves on instances + the optimiser is not fitting. Passing no validation set restores the old + behaviour, and with it the old failure mode. + + ``sigma`` is relative to the standard deviation of the incoming weights, so + a sensible perturbation scale does not depend on how the network was + initialised. + + Returns ``(bundle, history)``. + """ + config = config or RLConfig() + tasks = list(tasks) + scoring = list(validation_tasks) if validation_tasks else tasks + rng = random.Random(config.seed) + + base = bundle.model.get_flat() + spread = statistics.pstdev(base) if len(base) > 1 else 1.0 + step = max(1e-6, sigma * (spread or 1.0)) + deviations = [step] * len(base) + mean = list(base) + + elite_count = max(2, int(population * elite_fraction)) + working = bundle.model.copy() + probe = HeuristicBundle( + bundle.space, working, bundle.scale, bundle.metrics, bundle.config + ) + + incumbent = SearchCostCandidate(list(base), search_cost(bundle, scoring, config)) + history = [{"iteration": 0, **incumbent.cost.to_dict()}] + if config.verbose: # pragma: no cover - operator convenience + print( + f" cem start: validation coverage {incumbent.cost.coverage:.2f}, " + f"expansions {incumbent.cost.expansions:.0f}" + ) + + for iteration in range(1, iterations + 1): + candidates = [] + for _ in range(population): + theta = [m + rng.gauss(0.0, d) for m, d in zip(mean, deviations)] + working.set_flat(theta) + candidates.append( + SearchCostCandidate(theta, search_cost(probe, tasks, config)) + ) + candidates.sort(key=lambda c: c.cost.score) + elites = candidates[:elite_count] + + mean = [sum(c.theta[i] for c in elites) / len(elites) for i in range(len(mean))] + deviations = [ + max( + step * 0.1, + math.sqrt( + sum((c.theta[i] - mean[i]) ** 2 for c in elites) / len(elites) + ), + ) + for i in range(len(mean)) + ] + + # CEM's distribution mean is not one of the sampled candidates and can + # be worse than all of them, so it is scored rather than assumed good. + working.set_flat(mean) + centre = SearchCostCandidate(list(mean), search_cost(probe, tasks, config)) + proposal = min([centre] + elites, key=lambda c: c.cost.score) + + working.set_flat(proposal.theta) + validated = ( + proposal.cost if scoring is tasks else search_cost(probe, scoring, config) + ) + if validated.score < incumbent.cost.score: + incumbent = SearchCostCandidate(list(proposal.theta), validated) + history.append( + { + "iteration": iteration, + "tuning_score": proposal.cost.score, + **incumbent.cost.to_dict(), + } + ) + if config.verbose: # pragma: no cover - operator convenience + print( + f" cem iter {iteration}: tuning {proposal.cost.score:.0f}, " + f"validation {validated.score:.0f}, " + f"incumbent {incumbent.cost.score:.0f} " + f"(coverage {incumbent.cost.coverage:.2f})" + ) + + tuned = bundle.model.copy() + tuned.set_flat(incumbent.theta) + metrics = dict(bundle.metrics) + metrics["search_cost"] = incumbent.cost.to_dict() + return ( + HeuristicBundle(bundle.space, tuned, bundle.scale, metrics, bundle.config), + history, + ) + + +@dataclass +class SearchCostCandidate: + theta: list + cost: SearchCost diff --git a/jupyddl/learn/train.py b/jupyddl/learn/train.py new file mode 100644 index 0000000..aa04b7b --- /dev/null +++ b/jupyddl/learn/train.py @@ -0,0 +1,301 @@ +"""The imitation stage: fit a network to the cost-to-go a corpus of plans shows. + +Two objectives are available and the default is the less obvious one. + +**Regression** fits ``h(s) ≈ h*(s)``. It is what "learn a heuristic" usually +means, and it optimises a quantity greedy best-first search never reads. + +**Ranking** fits the *order* instead. At each state on a plan, the successor the +plan took should sort ahead of the ones it passed over. GBFS pops the minimum of +the open list, so a model uniformly 30 too high guides perfectly while scoring +terribly on RMSE, and a model with excellent RMSE that inverts one pair of +siblings sends the search into the wrong subtree. Chrestien et al. (NeurIPS +2023) make this argument at length and measure the gap; the default here +follows them. + +The default is not *pure* ranking. A ranking loss is invariant to any monotone +rescaling of the output, so it pins down no scale at all — which is fine for +GBFS and useless for weighted A*, where the weight multiplies a quantity that +now means nothing. A small regression term anchors it. ``rank_weight`` is that +trade-off, and 0.8 is a reasonable default rather than a tuned optimum. +""" + +from __future__ import annotations + +import math +import random +import time +from dataclasses import dataclass, field +from typing import Optional + +from .dataset import Corpus +from .heuristic import HeuristicBundle +from .model import MLP, Adam + +__all__ = ["TrainConfig", "TrainReport", "train", "evaluate"] + + +@dataclass +class TrainConfig: + """Hyper-parameters for the supervised stage.""" + + hidden: tuple = (32, 16) + epochs: int = 60 + batch_size: int = 64 + learning_rate: float = 0.01 + #: Share of the objective given to ranking; the rest goes to regression. + rank_weight: float = 0.8 + #: Ranking groups considered per step. Each contributes one forward pass + #: per successor, so this is the real cost knob. + rank_batch: int = 16 + #: Samples from satisficing plans carry an upper bound, not ``h*``. Halving + #: their weight is a middle road between trusting and discarding them. + suboptimal_weight: float = 0.5 + validation: float = 0.2 + #: Stop after this many epochs without a new best validation score. Set to + #: 0 to disable and always run every epoch. + patience: int = 12 + seed: int = 0 + verbose: bool = False + + def to_dict(self) -> dict: + return { + "hidden": list(self.hidden), + "epochs": self.epochs, + "batch_size": self.batch_size, + "learning_rate": self.learning_rate, + "rank_weight": self.rank_weight, + "rank_batch": self.rank_batch, + "suboptimal_weight": self.suboptimal_weight, + "validation": self.validation, + "patience": self.patience, + "seed": self.seed, + } + + +@dataclass +class TrainReport: + """What happened during training, for plotting and for the record.""" + + history: list = field(default_factory=list) + best_epoch: int = 0 + metrics: dict = field(default_factory=dict) + seconds: float = 0.0 + + +def train(corpus: Corpus, config: Optional[TrainConfig] = None) -> tuple: + """Fit a heuristic to ``corpus``. Returns ``(bundle, report)``.""" + config = config or TrainConfig() + if not corpus.samples: + raise ValueError("cannot train on an empty corpus") + + rng = random.Random(config.seed) + train_set, val_set = corpus.split(config.validation, seed=config.seed) + if not train_set.samples: # pragma: no cover - guarded by split() + raise ValueError("the training split came out empty") + + # Normalise targets so the network learns a shape rather than a magnitude. + targets = [s.target for s in train_set.samples] + scale = max(1e-6, sum(targets) / len(targets)) + + sizes = [corpus.space.size] + list(config.hidden) + [1] + model = MLP(sizes, seed=config.seed) + optimiser = Adam(model, lr=config.learning_rate) + + started = time.perf_counter() + report = TrainReport() + best_score = math.inf + best_state = model.to_dict() + best_epoch = 0 + stale = 0 + + order = list(range(len(train_set.samples))) + for epoch in range(1, config.epochs + 1): + rng.shuffle(order) + epoch_loss = 0.0 + steps = 0 + for start in range(0, len(order), config.batch_size): + batch = [ + train_set.samples[i] for i in order[start : start + config.batch_size] + ] + loss, grad_w, grad_b = _step( + model, batch, train_set.groups, scale, config, rng + ) + optimiser.step(grad_w, grad_b) + epoch_loss += loss + steps += 1 + + scored = val_set if val_set.samples else train_set + metrics = evaluate(model, scored, scale) + # Selection is on ranking accuracy where there is any to measure, since + # that is the quantity search consumes; MAE only breaks ties. + rank_metrics = evaluate_ranking(model, train_set.groups, scale) + score = -(rank_metrics.get("top1", 0.0)) + 1e-3 * metrics["mae"] + report.history.append( + { + "epoch": epoch, + "loss": epoch_loss / max(1, steps), + "val_mae": metrics["mae"], + "val_rmse": metrics["rmse"], + "train_top1": rank_metrics.get("top1", 0.0), + } + ) + if config.verbose: # pragma: no cover - operator convenience + print( + f" epoch {epoch:3d} loss {epoch_loss / max(1, steps):8.4f}" + f" val MAE {metrics['mae']:7.3f}" + f" top-1 {rank_metrics.get('top1', 0.0):5.3f}" + ) + if score < best_score - 1e-9: + best_score = score + best_state = model.to_dict() + best_epoch = epoch + stale = 0 + else: + stale += 1 + if config.patience and stale >= config.patience: + break + + model = MLP.from_dict(best_state) + report.best_epoch = best_epoch + report.seconds = time.perf_counter() - started + + final = evaluate(model, val_set if val_set.samples else train_set, scale) + final.update( + {"train_" + k: v for k, v in evaluate(model, train_set, scale).items()} + ) + final.update(evaluate_ranking(model, corpus.groups, scale)) + final["held_out_instances"] = len({s.instance for s in val_set.samples}) + report.metrics = final + + bundle = HeuristicBundle( + corpus.space, model, scale, metrics=final, config=config.to_dict() + ) + return bundle, report + + +def _step(model: MLP, batch, groups, scale: float, config: TrainConfig, rng): + """One optimiser step over a regression batch and a sample of ranking groups.""" + reg_weight = 1.0 - config.rank_weight + grad_w = None + grad_b = None + total_loss = 0.0 + + if reg_weight > 0.0 and batch: + features = [s.features for s in batch] + outputs, cache = model.forward_batch(features) + weights = [1.0 if s.optimal else config.suboptimal_weight for s in batch] + norm = sum(weights) or 1.0 + d_out = [] + for out, sample, weight in zip(outputs, batch, weights): + target = sample.target / scale + diff = out - target + total_loss += reg_weight * weight * diff * diff / norm + d_out.append(reg_weight * 2.0 * weight * diff / norm) + grad_w, grad_b = model.backward_batch(cache, d_out) + + if config.rank_weight > 0.0 and groups: + picked = ( + rng.sample(groups, config.rank_batch) + if len(groups) > config.rank_batch + else list(groups) + ) + flat = [] + spans = [] + for group in picked: + start = len(flat) + flat.append(group.chosen) + flat.extend(group.others) + spans.append((start, len(flat))) + outputs, cache = model.forward_batch([f for f, _ in flat]) + d_out = [0.0] * len(flat) + for start, end in spans: + # Score each successor by what search will actually compare: + # the step cost plus the predicted cost-to-go. + scores = [-(flat[i][1] / scale + outputs[i]) for i in range(start, end)] + peak = max(scores) + exps = [math.exp(s - peak) for s in scores] + total = sum(exps) + probs = [e / total for e in exps] + total_loss += ( + -config.rank_weight * math.log(max(probs[0], 1e-12)) / len(spans) + ) + for offset, prob in enumerate(probs): + chosen = 1.0 if offset == 0 else 0.0 + d_out[start + offset] = ( + config.rank_weight * (chosen - prob) / len(spans) + ) + rank_w, rank_b = model.backward_batch(cache, d_out) + if grad_w is None: + grad_w, grad_b = rank_w, rank_b + else: + _accumulate(grad_w, rank_w) + _accumulate(grad_b, rank_b) + + if grad_w is None: # pragma: no cover - both terms disabled + grad_w = [[[0.0] * len(r) for r in layer] for layer in model.weights] + grad_b = [[0.0] * len(b) for b in model.biases] + return total_loss, grad_w, grad_b + + +def _accumulate(into, extra) -> None: + for a, b in zip(into, extra): + if isinstance(a, list) and a and isinstance(a[0], list): + _accumulate(a, b) + else: + for i in range(len(a)): + a[i] += b[i] + + +def evaluate(model: MLP, corpus: Corpus, scale: float) -> dict: + """Regression error of ``model`` on ``corpus``, in original cost units.""" + if not corpus.samples: # pragma: no cover - guarded by callers + return {"mae": 0.0, "rmse": 0.0, "bias": 0.0, "count": 0} + outputs, _ = model.forward_batch([s.features for s in corpus.samples]) + errors = [out * scale - s.target for out, s in zip(outputs, corpus.samples)] + count = len(errors) + return { + "mae": sum(abs(e) for e in errors) / count, + "rmse": math.sqrt(sum(e * e for e in errors) / count), + # Systematically low is dangerous in a different way from + # systematically high: it makes the search look like A* with a weak + # heuristic rather than a greedy one, and expands far more. + "bias": sum(errors) / count, + "count": count, + } + + +def evaluate_ranking(model: MLP, groups, scale: float) -> dict: + """How often the model prefers the successor the plan actually took. + + This tracks search performance far better than RMSE does. ``top1`` is the + fraction of decision points where the plan's successor scores strictly + lowest; ``in_top2`` allows one better-looking sibling, which a search with + any backtracking will usually survive. + """ + groups = list(groups or ()) + if not groups: + return {} + flat = [] + spans = [] + for group in groups: + start = len(flat) + flat.append(group.chosen) + flat.extend(group.others) + spans.append((start, len(flat))) + outputs, _ = model.forward_batch([f for f, _ in flat]) + top1 = 0 + top2 = 0 + for start, end in spans: + scores = [flat[i][1] / scale + outputs[i] for i in range(start, end)] + chosen = scores[0] + better = sum(1 for s in scores[1:] if s < chosen) + if better == 0: + top1 += 1 + if better <= 1: + top2 += 1 + return { + "top1": top1 / len(spans), + "in_top2": top2 / len(spans), + "groups": len(spans), + } diff --git a/promo/jupyddl-rl.mp4 b/promo/jupyddl-rl.mp4 new file mode 100644 index 0000000..e10218d Binary files /dev/null and b/promo/jupyddl-rl.mp4 differ diff --git a/promo/rl-data.json b/promo/rl-data.json new file mode 100644 index 0000000..4c0ea64 --- /dev/null +++ b/promo/rl-data.json @@ -0,0 +1,495 @@ +{ + "space": { + "predicates": [ + "clear", + "handempty", + "holding", + "on", + "ontable" + ], + "features": 14, + "train_instances": 12, + "train_sizes": [ + 3, + 6 + ], + "eval_sizes": [ + 9, + 13 + ] + }, + "corpus": { + "count": 118, + "groups": 104, + "mean": 5.279661016949152, + "std": 3.861776163804393, + "min": 0.0, + "max": 16.0, + "optimal_fraction": 1.0 + }, + "plan": { + "instance": "blocksworld-03-2", + "steps": [ + "unstack(b1,b3)", + "put-down(b1)", + "pick-up(b2)", + "stack(b2,b3)" + ], + "suffix": [ + 4.0, + 3.0, + 2.0, + 1.0 + ] + }, + "imitation": { + "mae": 0.9205159587205748, + "top1": 0.9230769230769231, + "in_top2": 1.0, + "best_epoch": 20, + "seconds": 0.06553611000003912, + "parameters": 1025, + "history": [ + { + "epoch": 1, + "top1": 0.7763157894736842, + "loss": 1.0150187332387282 + }, + { + "epoch": 2, + "top1": 0.7894736842105263, + "loss": 0.9252037718213768 + }, + { + "epoch": 3, + "top1": 0.881578947368421, + "loss": 0.8130478089675709 + }, + { + "epoch": 4, + "top1": 0.8947368421052632, + "loss": 0.775756577786443 + }, + { + "epoch": 5, + "top1": 0.8947368421052632, + "loss": 0.7548487659235616 + }, + { + "epoch": 6, + "top1": 0.8947368421052632, + "loss": 0.7223636728634536 + }, + { + "epoch": 7, + "top1": 0.8947368421052632, + "loss": 0.7723022515427278 + }, + { + "epoch": 8, + "top1": 0.9078947368421053, + "loss": 0.7035696388901393 + }, + { + "epoch": 9, + "top1": 0.9078947368421053, + "loss": 0.7455279763920939 + }, + { + "epoch": 10, + "top1": 0.9210526315789473, + "loss": 0.7359165954972391 + }, + { + "epoch": 11, + "top1": 0.9210526315789473, + "loss": 0.6859913176223392 + }, + { + "epoch": 12, + "top1": 0.881578947368421, + "loss": 0.689344075109841 + }, + { + "epoch": 13, + "top1": 0.881578947368421, + "loss": 0.6480429372974326 + }, + { + "epoch": 14, + "top1": 0.9078947368421053, + "loss": 0.6735856352437681 + }, + { + "epoch": 15, + "top1": 0.9210526315789473, + "loss": 0.665086602116252 + }, + { + "epoch": 16, + "top1": 0.9210526315789473, + "loss": 0.7346798744247327 + }, + { + "epoch": 17, + "top1": 0.9210526315789473, + "loss": 0.7508305741985783 + }, + { + "epoch": 18, + "top1": 0.9210526315789473, + "loss": 0.6577391223695969 + }, + { + "epoch": 19, + "top1": 0.9210526315789473, + "loss": 0.7290821330428392 + }, + { + "epoch": 20, + "top1": 0.9210526315789473, + "loss": 0.6820188370948383 + }, + { + "epoch": 21, + "top1": 0.881578947368421, + "loss": 0.6792723684862068 + }, + { + "epoch": 22, + "top1": 0.9210526315789473, + "loss": 0.7485992569790645 + }, + { + "epoch": 23, + "top1": 0.8947368421052632, + "loss": 0.6309181104792626 + }, + { + "epoch": 24, + "top1": 0.868421052631579, + "loss": 0.6969212951548777 + }, + { + "epoch": 25, + "top1": 0.868421052631579, + "loss": 0.6996913806138843 + }, + { + "epoch": 26, + "top1": 0.8947368421052632, + "loss": 0.6869846718947648 + }, + { + "epoch": 27, + "top1": 0.8947368421052632, + "loss": 0.6906913144263385 + }, + { + "epoch": 28, + "top1": 0.8947368421052632, + "loss": 0.7113118541923102 + }, + { + "epoch": 29, + "top1": 0.881578947368421, + "loss": 0.6900435487807537 + }, + { + "epoch": 30, + "top1": 0.881578947368421, + "loss": 0.6810791261501254 + }, + { + "epoch": 31, + "top1": 0.8552631578947368, + "loss": 0.7706092945366239 + }, + { + "epoch": 32, + "top1": 0.8552631578947368, + "loss": 0.7019795680509671 + } + ] + }, + "transfer_before": { + "learned": { + "solved": 9, + "total": 10, + "expanded": 3290, + "seconds": 1.2946110370000952, + "cost": 442, + "mean_expanded": 365.55555555555554, + "mean_seconds": 0.14384567077778834, + "mean_cost": 49.111111111111114, + "coverage": 0.9 + }, + "hff": { + "solved": 10, + "total": 10, + "expanded": 5183, + "seconds": 5.5783092040001065, + "cost": 518, + "mean_expanded": 518.3, + "mean_seconds": 0.5578309204000107, + "mean_cost": 51.8, + "coverage": 1.0 + }, + "goalcount": { + "solved": 10, + "total": 10, + "expanded": 24834, + "seconds": 1.6121297669999421, + "cost": 510, + "mean_expanded": 2483.4, + "mean_seconds": 0.1612129766999942, + "mean_cost": 51.0, + "coverage": 1.0 + }, + "blind": { + "solved": 0, + "total": 10, + "expanded": 0, + "seconds": 0.0, + "cost": 0, + "mean_expanded": 0.0, + "mean_seconds": 0.0, + "mean_cost": 0.0, + "coverage": 0.0 + } + }, + "cem": { + "seconds": 38.99109729299994, + "iterations": 10, + "population": 12, + "history": [ + { + "iteration": 0, + "tuning": null, + "validation": 276.125, + "coverage": 1.0 + }, + { + "iteration": 1, + "tuning": 106.0, + "validation": 76.125, + "coverage": 1.0 + }, + { + "iteration": 2, + "tuning": 114.5, + "validation": 76.125, + "coverage": 1.0 + }, + { + "iteration": 3, + "tuning": 109.75, + "validation": 76.125, + "coverage": 1.0 + }, + { + "iteration": 4, + "tuning": 112.125, + "validation": 76.125, + "coverage": 1.0 + }, + { + "iteration": 5, + "tuning": 106.25, + "validation": 76.125, + "coverage": 1.0 + }, + { + "iteration": 6, + "tuning": 101.0, + "validation": 73.25, + "coverage": 1.0 + }, + { + "iteration": 7, + "tuning": 101.0, + "validation": 73.25, + "coverage": 1.0 + }, + { + "iteration": 8, + "tuning": 100.75, + "validation": 73.125, + "coverage": 1.0 + }, + { + "iteration": 9, + "tuning": 100.625, + "validation": 73.125, + "coverage": 1.0 + }, + { + "iteration": 10, + "tuning": 99.875, + "validation": 73.125, + "coverage": 1.0 + } + ] + }, + "transfer_after": { + "learned": { + "solved": 10, + "total": 10, + "expanded": 1371, + "seconds": 0.36463156399986474, + "cost": 482, + "mean_expanded": 137.1, + "mean_seconds": 0.03646315639998647, + "mean_cost": 48.2, + "coverage": 1.0 + }, + "hff": { + "solved": 10, + "total": 10, + "expanded": 5183, + "seconds": 5.3366881800002375, + "cost": 518, + "mean_expanded": 518.3, + "mean_seconds": 0.5336688180000237, + "mean_cost": 51.8, + "coverage": 1.0 + }, + "goalcount": { + "solved": 10, + "total": 10, + "expanded": 24834, + "seconds": 1.5423849710005015, + "cost": 510, + "mean_expanded": 2483.4, + "mean_seconds": 0.15423849710005016, + "mean_cost": 51.0, + "coverage": 1.0 + }, + "blind": { + "solved": 0, + "total": 10, + "expanded": 0, + "seconds": 0.0, + "cost": 0, + "mean_expanded": 0.0, + "mean_seconds": 0.0, + "mean_cost": 0.0, + "coverage": 0.0 + } + }, + "flat": { + "easy_before": 12.833333333333334, + "easy_after": 12.25, + "hard_before": 152.125, + "hard_after": 100.75 + }, + "spread": { + "instances": [ + "blocksworld-09-7777", + "blocksworld-09-7778", + "blocksworld-10-7777", + "blocksworld-10-7778", + "blocksworld-11-7777", + "blocksworld-11-7778", + "blocksworld-12-7777", + "blocksworld-12-7778", + "blocksworld-13-7777", + "blocksworld-13-7778" + ], + "imitation": { + "blocksworld-09-7777": 77, + "blocksworld-09-7778": 126, + "blocksworld-10-7777": 70, + "blocksworld-10-7778": 287, + "blocksworld-11-7777": 113, + "blocksworld-11-7778": 150, + "blocksworld-12-7777": 349, + "blocksworld-12-7778": 114, + "blocksworld-13-7777": 30000, + "blocksworld-13-7778": 2004 + }, + "imitation_solved": { + "blocksworld-09-7777": true, + "blocksworld-09-7778": true, + "blocksworld-10-7777": true, + "blocksworld-10-7778": true, + "blocksworld-11-7777": true, + "blocksworld-11-7778": true, + "blocksworld-12-7777": true, + "blocksworld-12-7778": true, + "blocksworld-13-7777": false, + "blocksworld-13-7778": true + }, + "lo": { + "sigma": 0.05, + "tuning": 108.5, + "per_instance": { + "blocksworld-09-7777": 69, + "blocksworld-09-7778": 67, + "blocksworld-10-7777": 33, + "blocksworld-10-7778": 229, + "blocksworld-11-7777": 90, + "blocksworld-11-7778": 121, + "blocksworld-12-7777": 228, + "blocksworld-12-7778": 72, + "blocksworld-13-7777": 16121, + "blocksworld-13-7778": 273 + }, + "solved": { + "blocksworld-09-7777": true, + "blocksworld-09-7778": true, + "blocksworld-10-7777": true, + "blocksworld-10-7778": true, + "blocksworld-11-7777": true, + "blocksworld-11-7778": true, + "blocksworld-12-7777": true, + "blocksworld-12-7778": true, + "blocksworld-13-7777": true, + "blocksworld-13-7778": true + }, + "mean": 1730.3 + }, + "hi": { + "sigma": 0.15, + "tuning": 99.875, + "per_instance": { + "blocksworld-09-7777": 70, + "blocksworld-09-7778": 68, + "blocksworld-10-7777": 58, + "blocksworld-10-7778": 217, + "blocksworld-11-7777": 81, + "blocksworld-11-7778": 111, + "blocksworld-12-7777": 227, + "blocksworld-12-7778": 70, + "blocksworld-13-7777": 214, + "blocksworld-13-7778": 255 + }, + "solved": { + "blocksworld-09-7777": true, + "blocksworld-09-7778": true, + "blocksworld-10-7777": true, + "blocksworld-10-7778": true, + "blocksworld-11-7777": true, + "blocksworld-11-7778": true, + "blocksworld-12-7777": true, + "blocksworld-12-7778": true, + "blocksworld-13-7777": true, + "blocksworld-13-7778": true + }, + "mean": 137.1 + }, + "budget": 30000 + }, + "logistics": { + "predicates": [ + "at", + "in" + ], + "features": 8, + "top1": 0.6560509554140127, + "learned": 203.83333333333334, + "hff": 34.666666666666664, + "blocks_top1": 0.9230769230769231 + }, + "total_seconds": 229.68699452200008 +} \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 2783311..59e1bd6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,6 +24,10 @@ dependencies = [] # 3.6 is the floor for the constrained-layout engine API the charts rely on # (``layout="constrained"`` plus ``Figure.get_layout_engine``). viz = ["matplotlib>=3.6"] +# Purely a speed option. jupyddl.learn trains and evaluates on the standard +# library alone; NumPy just makes the batched passes one to two orders of +# magnitude faster, and tests/test_learn.py pins the two paths together. +learn = ["numpy>=1.21"] dev = ["pytest>=7", "pytest-cov>=4", "flake8>=6", "black>=23"] [project.scripts] diff --git a/tests/test_learn.py b/tests/test_learn.py new file mode 100644 index 0000000..1b6e121 --- /dev/null +++ b/tests/test_learn.py @@ -0,0 +1,653 @@ +"""Learned heuristics: features, model, corpus, training and the RL stages. + +Everything here runs on tiny budgets. The point is that the machinery is +correct, not that a 20-epoch model on six instances is any good — the quality +claims live in ``.docs/learned-heuristics.md`` where they can be stated with +the numbers that produced them. +""" + +from __future__ import annotations + +import json +import math +import random + +import pytest + +from jupyddl.heuristics import make_heuristic +from jupyddl.learn import ( + Corpus, + FeatureSpace, + HeuristicBundle, + LearnedHeuristic, + MLP, + TrainConfig, + build_corpus, + numpy_available, + samples_from_plan, + train, +) +from jupyddl.learn.dataset import task_from_state +from jupyddl.learn.features import predicate_of +from jupyddl.learn.model import Adam +from jupyddl.learn.pipeline import ( + evaluate_transfer, + solved_corpus, + summarise_transfer, + tasks_from_generator, +) +from jupyddl.learn.rl import ( + RLConfig, + bootstrap, + dagger, + optimise_search_cost, + search_cost, +) +from jupyddl.learn.train import evaluate_ranking + + +@pytest.fixture(scope="module") +def ladder(): + return tasks_from_generator("blocksworld", range(3, 5), seed=0, seeds_per_size=2) + + +@pytest.fixture(scope="module") +def space(ladder): + return FeatureSpace.from_tasks(task for _, task in ladder) + + +@pytest.fixture(scope="module") +def corpus(ladder, space): + return solved_corpus(ladder, space=space, time_limit=20.0) + + +# ========================================================================== +# features +# ========================================================================== +@pytest.mark.parametrize( + "fact,symbol", + [ + ("(on b1 b2)", "on"), + ("(handempty)", "handempty"), + ("(at truck1 loc2)", "at"), + ("move(a,b)#2", "move"), + ("__closed", "__closed"), + ], +) +def test_predicate_symbols_are_extracted_from_both_spellings(fact, symbol): + """Facts print Lisp-style and operators functionally; both must parse. + + Getting this wrong does not raise. It gives every ground atom its own + vocabulary slot, and the feature vector quietly stops being size-invariant. + """ + assert predicate_of(fact) == symbol + + +def test_vocabulary_is_predicates_not_atoms(space): + assert set(space.vocabulary) == {"on", "ontable", "clear", "handempty", "holding"} + + +def test_feature_vector_is_the_same_length_for_every_instance_size(): + """The whole transfer claim rests on this.""" + small = tasks_from_generator("blocksworld", [3], seed=0)[0][1] + large = tasks_from_generator("blocksworld", [11], seed=0)[0][1] + space = FeatureSpace.from_tasks([small, large]) + a = space.bind(small)(small.initial_state()) + b = space.bind(large)(large.initial_state()) + assert len(a) == len(b) == space.size + assert space.size == 2 * 5 + len(FeatureSpace.GLOBAL_FEATURES) + + +def test_features_are_bounded_so_scale_does_not_leak_in(): + """Counts are normalised by their per-predicate totals, so 40 blocks and + 4 blocks land in the same range rather than an order of magnitude apart.""" + for size in (3, 8, 14): + task = tasks_from_generator("blocksworld", [size], seed=1)[0][1] + space = FeatureSpace.from_task(task) + vector = space.bind(task)(task.initial_state()) + assert all(0.0 <= v <= 3.0 for v in vector), (size, vector) + + +def test_a_symbol_outside_the_vocabulary_does_not_shift_the_others(): + """A model must survive a task using a predicate it never saw.""" + task = tasks_from_generator("blocksworld", [3], seed=0)[0][1] + space = FeatureSpace(["clear", "on", "totally-made-up"]) + vector = space.bind(task)(task.initial_state()) + assert len(vector) == space.size + assert vector[space.vocabulary.index("totally-made-up")] == 0.0 + + +def test_feature_names_line_up_with_the_vector(space): + task = tasks_from_generator("blocksworld", [3], seed=0)[0][1] + assert len(space.names()) == len(space.bind(task)(task.initial_state())) + + +def test_feature_space_round_trips(space): + assert FeatureSpace.from_dict(json.loads(json.dumps(space.to_dict()))) == space + + +# ========================================================================== +# model +# ========================================================================== +def test_gradients_match_finite_differences(): + """The one test that would catch a wrong backward pass.""" + rng = random.Random(3) + model = MLP([6, 5, 4, 1], seed=2) + batch = [[rng.gauss(0, 1) for _ in range(6)] for _ in range(9)] + targets = [abs(rng.gauss(2, 1)) for _ in range(9)] + + def loss_of(m): + out, _ = m.forward_batch(batch) + return sum((o - t) ** 2 for o, t in zip(out, targets)) / len(targets) + + outputs, cache = model.forward_batch(batch) + d_out = [2 * (o - t) / len(targets) for o, t in zip(outputs, targets)] + grad_w, grad_b = model.backward_batch(cache, d_out) + analytic = [v for layer in grad_w for row in layer for v in row] + analytic += [v for layer in grad_b for v in layer] + + flat = model.get_flat() + eps = 1e-6 + for index in range(len(flat)): + high = list(flat) + high[index] += eps + model.set_flat(high) + upper = loss_of(model) + low = list(flat) + low[index] -= eps + model.set_flat(low) + lower = loss_of(model) + model.set_flat(flat) + assert abs((upper - lower) / (2 * eps) - analytic[index]) < 1e-6 + + +@pytest.mark.skipif(not numpy_available(), reason="numpy is not installed") +def test_the_numpy_path_agrees_with_the_reference(): + """A fast path that has drifted is worse than no fast path.""" + rng = random.Random(0) + model = MLP([7, 6, 1], seed=1) + batch = [[rng.gauss(0, 1) for _ in range(7)] for _ in range(11)] + d_out = [rng.gauss(0, 1) for _ in range(11)] + + fast_out, fast_cache = model._forward_numpy(batch) + slow_out, slow_cache = model._forward_python(batch) + assert fast_out == pytest.approx(slow_out, abs=1e-12) + + fw, fb = model._backward_numpy(fast_cache, d_out) + sw, sb = model._backward_python(slow_cache, d_out) + flat = lambda g: [v for layer in g for row in layer for v in row] # noqa: E731 + assert flat(fw) == pytest.approx(flat(sw), abs=1e-12) + assert [v for layer in fb for v in layer] == pytest.approx( + [v for layer in sb for v in layer], abs=1e-12 + ) + + +def test_single_sample_call_matches_the_batch(space): + rng = random.Random(5) + model = MLP([4, 3, 1], seed=0) + batch = [[rng.gauss(0, 1) for _ in range(4)] for _ in range(3)] + outputs, _ = model.forward_batch(batch) + assert [model(v) for v in batch] == pytest.approx(outputs, abs=1e-12) + + +def test_the_output_is_never_negative(): + """A negative cost-to-go breaks every planner that consumes it.""" + model = MLP([3, 3, 1], seed=0) + model.set_flat([-50.0] * model.num_parameters) + assert model([1.0, 1.0, 1.0]) >= 0.0 + + +def test_flat_round_trip_preserves_predictions(): + model = MLP([5, 4, 1], seed=7) + vector = [0.3, -0.2, 0.9, 0.1, 0.0] + before = model(vector) + model.set_flat(model.get_flat()) + assert model(vector) == pytest.approx(before) + + +def test_set_flat_rejects_the_wrong_length(): + model = MLP([3, 2, 1], seed=0) + with pytest.raises(ValueError): + model.set_flat([0.0]) + + +def test_adam_reduces_a_loss_it_can_reach(): + rng = random.Random(1) + model = MLP([3, 4, 1], seed=0) + optimiser = Adam(model, lr=0.05) + batch = [[rng.gauss(0, 1) for _ in range(3)] for _ in range(16)] + targets = [2.0] * 16 + + def loss_of(): + out, _ = model.forward_batch(batch) + return sum((o - t) ** 2 for o, t in zip(out, targets)) / len(targets) + + first = loss_of() + for _ in range(60): + outputs, cache = model.forward_batch(batch) + d_out = [2 * (o - t) / len(targets) for o, t in zip(outputs, targets)] + optimiser.step(*model.backward_batch(cache, d_out)) + assert loss_of() < first * 0.1 + + +# ========================================================================== +# corpus +# ========================================================================== +def test_plan_suffix_costs_are_the_targets(ladder, space): + name, task = ladder[0] + from jupyddl.api import solve_task + + result = solve_task(task, "astar", "lmcut", time_limit=20) + assert result.solved + samples, groups = samples_from_plan(task, result.plan, space.bind(task), name) + assert len(samples) == len(result.plan) + 1 + assert samples[-1].target == 0.0 + assert samples[0].target == pytest.approx(result.cost) + # Monotonically decreasing towards the goal, by construction. + assert all(a.target >= b.target for a, b in zip(samples, samples[1:])) + assert groups, "a plan through a branching state should yield ranking groups" + + +def test_ranking_groups_hold_the_chosen_successor_first(ladder, space): + from jupyddl.api import solve_task + + name, task = ladder[0] + result = solve_task(task, "astar", "lmcut", time_limit=20) + _, groups = samples_from_plan(task, result.plan, space.bind(task), name) + for group in groups: + assert len(group.chosen) == 2 + assert group.others + assert all(len(other) == 2 for other in group.others) + + +def test_corpus_splits_by_instance_not_by_sample(corpus): + """Two states on the same plan are near-duplicates; splitting them across + train and validation reports memorisation as generalisation.""" + train_set, val_set = corpus.split(validation=0.34, seed=0) + assert train_set.samples and val_set.samples + assert not ( + {s.instance for s in train_set.samples} & {s.instance for s in val_set.samples} + ) + + +def test_corpus_round_trips(corpus, tmp_path): + path = tmp_path / "corpus.json" + corpus.save(str(path)) + again = Corpus.load(str(path)) + assert len(again) == len(corpus) + assert again.space == corpus.space + assert len(again.groups) == len(corpus.groups) + + +def test_corpus_refuses_to_merge_mismatched_spaces(corpus): + other = Corpus(FeatureSpace(["nonsense"])) + with pytest.raises(ValueError): + other.extend(corpus) + + +def test_build_corpus_survives_an_unsolved_instance(space, ladder): + from jupyddl.search.result import SearchResult + + def refuse(_task): + return SearchResult(False, None, None) + + empty = build_corpus(ladder, refuse, space=space) + assert len(empty) == 0 + + +def test_task_from_state_reroots_without_touching_anything_else(ladder): + from jupyddl.api import solve_task, validate_plan + + _, task = ladder[0] + result = solve_task(task, "astar", "lmcut", time_limit=20) + midpoint = task.initial_state() + for operator in result.plan[: len(result.plan) // 2]: + midpoint = task.apply(operator, midpoint) + + rerooted = task_from_state(task, midpoint) + assert rerooted.goals == task.goals + assert rerooted.operators is task.operators + tail = solve_task(rerooted, "astar", "lmcut", time_limit=20) + assert tail.solved + assert validate_plan(rerooted, tail.plan) + + +# ========================================================================== +# training +# ========================================================================== +def test_training_beats_predicting_the_mean(corpus): + """The floor any regressor must clear to have learned anything.""" + bundle, report = train(corpus, TrainConfig(epochs=40, seed=0)) + targets = [s.target for s in corpus.samples] + mean = sum(targets) / len(targets) + baseline = sum(abs(t - mean) for t in targets) / len(targets) + assert bundle.metrics["train_mae"] < baseline + assert report.best_epoch >= 1 + assert report.history + + +def test_ranking_accuracy_is_reported_and_sane(corpus): + bundle, _ = train(corpus, TrainConfig(epochs=40, seed=0)) + metrics = evaluate_ranking(bundle.model, corpus.groups, bundle.scale) + assert 0.0 <= metrics["top1"] <= 1.0 + assert metrics["in_top2"] >= metrics["top1"] + + +def test_pure_regression_and_pure_ranking_both_train(corpus): + for weight in (0.0, 1.0): + bundle, _ = train(corpus, TrainConfig(epochs=15, rank_weight=weight, seed=0)) + assert math.isfinite(bundle.metrics["mae"]) + + +def test_training_refuses_an_empty_corpus(space): + with pytest.raises(ValueError): + train(Corpus(space), TrainConfig(epochs=1)) + + +def test_training_is_deterministic_given_a_seed(corpus): + a, _ = train(corpus, TrainConfig(epochs=10, seed=3)) + b, _ = train(corpus, TrainConfig(epochs=10, seed=3)) + assert a.model.get_flat() == pytest.approx(b.model.get_flat()) + + +# ========================================================================== +# the heuristic itself +# ========================================================================== +def test_bundle_round_trips_and_predicts_identically(corpus, ladder, tmp_path): + bundle, _ = train(corpus, TrainConfig(epochs=20, seed=0)) + path = tmp_path / "model.heur.json" + bundle.save(str(path)) + again = HeuristicBundle.load(str(path)) + _, task = ladder[0] + state = task.initial_state() + assert again.bind(task)(state) == pytest.approx(bundle.bind(task)(state)) + + +def test_a_future_format_is_refused_rather_than_misread(corpus, tmp_path): + bundle, _ = train(corpus, TrainConfig(epochs=5, seed=0)) + data = bundle.to_dict() + data["format"] = 99 + path = tmp_path / "future.json" + path.write_text(json.dumps(data), encoding="utf-8") + with pytest.raises(ValueError, match="unsupported"): + HeuristicBundle.load(str(path)) + + +def test_the_heuristic_scores_a_goal_state_zero(corpus, ladder): + bundle, _ = train(corpus, TrainConfig(epochs=20, seed=0)) + from jupyddl.api import solve_task + + _, task = ladder[0] + result = solve_task(task, "astar", "lmcut", time_limit=20) + goal = task.initial_state() + for operator in result.plan: + goal = task.apply(operator, goal) + assert task.goal_reached(goal) + assert bundle.bind(task)(goal) == 0.0 + + +def test_the_heuristic_is_never_negative(corpus, ladder): + bundle, _ = train(corpus, TrainConfig(epochs=20, seed=0)) + _, task = ladder[0] + heuristic = bundle.bind(task) + state = task.initial_state() + assert heuristic(state) >= 0.0 + for operator in task.operators[:20]: + if operator.applicable(state): + assert heuristic(task.apply(operator, state)) >= 0.0 + + +def test_values_are_cached_per_state(corpus, ladder): + bundle, _ = train(corpus, TrainConfig(epochs=5, seed=0)) + _, task = ladder[0] + heuristic = bundle.bind(task) + state = task.initial_state() + heuristic(state) + heuristic(state) + assert heuristic.evaluations == 1 + + +def test_it_does_not_claim_to_be_admissible(): + """Nothing in the objective bounds the prediction from above.""" + assert LearnedHeuristic.admissible is False + + +# ========================================================================== +# integration: the registry, the API and the CLI all speak `learned:` +# ========================================================================== +def test_learned_spec_resolves_through_the_registry(corpus, ladder, tmp_path): + bundle, _ = train(corpus, TrainConfig(epochs=10, seed=0)) + path = tmp_path / "reg.heur.json" + bundle.save(str(path)) + _, task = ladder[0] + heuristic = make_heuristic(f"learned:{path}", task) + assert isinstance(heuristic, LearnedHeuristic) + assert heuristic(task.initial_state()) >= 0.0 + + +def test_an_already_built_heuristic_passes_through(ladder): + _, task = ladder[0] + built = make_heuristic("hff", task) + assert make_heuristic(built, task) is built + + +def test_an_unknown_loader_is_rejected(ladder): + _, task = ladder[0] + with pytest.raises(ValueError, match="parameterised"): + make_heuristic("psychic:model.json", task) + + +def test_solve_task_accepts_a_learned_spec(corpus, ladder, tmp_path): + from jupyddl.api import solve_task, validate_plan + + bundle, _ = train(corpus, TrainConfig(epochs=20, seed=0)) + path = tmp_path / "solve.heur.json" + bundle.save(str(path)) + _, task = ladder[-1] + result = solve_task(task, "gbfs", f"learned:{path}", time_limit=20) + assert result.solved + assert validate_plan( + task, result.plan + ), "a learned heuristic must not break soundness" + + +def test_the_cli_accepts_and_rejects_heuristic_specs(): + from jupyddl.cli import heuristic_spec + import argparse + + assert heuristic_spec("lmcut") == "lmcut" + assert heuristic_spec("none") == "none" + assert heuristic_spec("learned:x.json") == "learned:x.json" + for bad in ("lmcutt", "learned:", "psychic:x"): + with pytest.raises(argparse.ArgumentTypeError): + heuristic_spec(bad) + + +def test_a_learned_heuristic_still_yields_valid_plans(corpus, ladder): + """The heuristic may be wrong; the planner may not become unsound.""" + from jupyddl.api import solve_task, validate_plan + + bundle, _ = train(corpus, TrainConfig(epochs=20, seed=0)) + for _, task in ladder: + result = solve_task(task, "gbfs", bundle.bind(task), time_limit=20) + assert result.solved + assert validate_plan(task, result.plan) + + +# ========================================================================== +# the reinforcement stages +# ========================================================================== +def test_search_cost_charges_a_failure_more_than_any_possible_success(corpus, ladder): + """At a fixed budget, giving up must never look cheaper than finishing. + + Scores are only comparable *at the same budget* — the penalty is a multiple + of it — so this compares a failure against the worst success the same + budget allows, not against a run under a different one. + """ + bundle, _ = train(corpus, TrainConfig(epochs=10, seed=0)) + config = RLConfig(max_expansions=1) # nothing is solvable in one expansion + starved = search_cost(bundle, ladder, config) + assert starved.coverage == 0.0 + assert starved.score == pytest.approx( + config.failure_penalty * config.max_expansions + ) + assert starved.score > config.max_expansions + + +def test_search_cost_reports_full_coverage_when_everything_solves(corpus, ladder): + bundle, _ = train(corpus, TrainConfig(epochs=10, seed=0)) + cost = search_cost(bundle, ladder, RLConfig(max_expansions=5000)) + assert cost.coverage == 1.0 + assert cost.score < 5000 + + +def test_search_cost_needs_tasks(corpus): + bundle, _ = train(corpus, TrainConfig(epochs=5, seed=0)) + with pytest.raises(ValueError): + search_cost(bundle, []) + + +def test_dagger_adds_samples_from_the_search_distribution(corpus, ladder): + bundle, _ = train(corpus, TrainConfig(epochs=10, seed=0)) + before = len(corpus) + tuned, grown, history = dagger( + bundle, + ladder, + corpus, + rounds=1, + states_per_task=4, + config=RLConfig(max_expansions=500, seed=0), + train_config=TrainConfig(epochs=5, seed=0), + ) + assert len(grown) > before + assert history[0]["round"] == 1 + assert any(s.instance.endswith("#dagger1") for s in grown.samples) + # Aggregated labels come from a satisficing solver, so they are bounds. + assert any(not s.optimal for s in grown.samples) + assert tuned.bind(ladder[0][1])(ladder[0][1].initial_state()) >= 0.0 + + +def test_the_default_labeller_keeps_tasks_apart(ladder): + """Its heuristic cache is keyed on ``id()``, which CPython recycles. + + The entry holds the task so the id cannot be reused underneath it, and + checks identity anyway. Without both, a caller whose tasks this cache does + not keep alive could be handed another task's relaxed-task tables — which + would not raise, just quietly mislabel. + """ + from jupyddl.learn.rl import _default_labeller + + label = _default_labeller(RLConfig(max_expansions=2000, seed=0)) + for _, task in ladder[:2]: + plan = label(task, task.initial_state()) + assert plan, "a solvable task should get a plan from its own heuristic" + state = task.initial_state() + for operator in plan: + assert operator.applicable(state), "plan came from the wrong task" + state = task.apply(operator, state) + assert task.goal_reached(state) + + +def test_bootstrap_records_what_it_newly_solved(corpus, ladder): + bundle, _ = train(corpus, TrainConfig(epochs=10, seed=0)) + harder = tasks_from_generator("blocksworld", [5], seed=11, seeds_per_size=2) + _, grown, history = bootstrap( + bundle, + harder, + corpus, + rounds=1, + config=RLConfig(max_expansions=2000, seed=0), + train_config=TrainConfig(epochs=5, seed=0), + ) + assert history and "newly_solved" in history[0] + assert history[0]["remaining"] + len(history[0]["newly_solved"]) == len(harder) + + +def test_cem_never_returns_a_worse_incumbent(corpus, ladder): + """The incumbent is only displaced by a strictly better mean score.""" + bundle, _ = train(corpus, TrainConfig(epochs=15, seed=0)) + config = RLConfig(max_expansions=500, seed=0) + before = search_cost(bundle, ladder, config).score + tuned, history = optimise_search_cost( + bundle, ladder, iterations=2, population=4, config=config + ) + assert search_cost(tuned, ladder, config).score <= before + assert history[0]["iteration"] == 0 + assert tuned.space == bundle.space + + +def test_cem_selects_the_incumbent_on_the_validation_set(corpus, ladder): + """Tuning a thousand parameters against a handful of instances fits those + instances. The returned model must improve on ones it was not fitted to.""" + bundle, _ = train(corpus, TrainConfig(epochs=15, seed=0)) + held_out = tasks_from_generator("blocksworld", [4], seed=321, seeds_per_size=2) + config = RLConfig(max_expansions=800, seed=0) + before = search_cost(bundle, held_out, config).score + + tuned, history = optimise_search_cost( + bundle, + ladder, + iterations=2, + population=4, + config=config, + validation_tasks=held_out, + ) + # The reported score is the validation score, and the incumbent only ever + # moves when that improves — so it cannot come back worse. + assert search_cost(tuned, held_out, config).score <= before + assert history[-1]["score"] <= before + # Both scores are reported, so a run that is fitting its tuning set while + # losing validation is visible rather than hidden. + assert "tuning_score" in history[-1] + + +def test_cem_without_a_validation_set_scores_on_the_tuning_tasks(corpus, ladder): + """The old behaviour stays reachable, and stays honest about what it means.""" + bundle, _ = train(corpus, TrainConfig(epochs=10, seed=0)) + config = RLConfig(max_expansions=500, seed=0) + before = search_cost(bundle, ladder, config).score + tuned, _ = optimise_search_cost( + bundle, ladder, iterations=2, population=4, config=config + ) + assert search_cost(tuned, ladder, config).score <= before + + +def test_cem_leaves_the_original_bundle_untouched(corpus, ladder): + bundle, _ = train(corpus, TrainConfig(epochs=10, seed=0)) + before = list(bundle.model.get_flat()) + optimise_search_cost( + bundle, ladder, iterations=1, population=4, config=RLConfig(max_expansions=300) + ) + assert bundle.model.get_flat() == pytest.approx(before) + + +# ========================================================================== +# transfer +# ========================================================================== +def test_transfer_evaluation_reports_time_as_well_as_expansions(corpus, ladder): + """Expansions alone can hide a heuristic that is slower per node than it + is smart, which is the standard way to publish a non-result.""" + bundle, _ = train(corpus, TrainConfig(epochs=20, seed=0)) + rows = evaluate_transfer( + bundle, ladder[:2], baselines=("goalcount",), max_expansions=2000 + ) + assert {row["heuristic"] for row in rows} == {"learned", "goalcount"} + for row in rows: + assert row["seconds"] >= 0.0 + assert "expanded" in row + summary = summarise_transfer(rows) + assert summary["learned"]["coverage"] == 1.0 + + +def test_a_model_transfers_to_an_instance_far_larger_than_it_trained_on(corpus): + """Not a quality claim — a claim that it runs at all, which is the thing + a fixed-length one-hot encoding would have made impossible.""" + from jupyddl.api import solve_task, validate_plan + + bundle, _ = train(corpus, TrainConfig(epochs=30, seed=0)) + big = tasks_from_generator("blocksworld", [12], seed=99)[0][1] + heuristic = bundle.bind(big) + assert heuristic(big.initial_state()) >= 0.0 + result = solve_task(big, "gbfs", heuristic, max_expansions=20000, time_limit=30) + if result.solved: + assert validate_plan(big, result.plan) diff --git a/tools/build_web.py b/tools/build_web.py index c94bd1f..c8ced65 100644 --- a/tools/build_web.py +++ b/tools/build_web.py @@ -28,7 +28,10 @@ OUT = os.path.join(ROOT, "web", "dist") # viz/ needs matplotlib, which the playground deliberately does not ship. -SKIP_DIRS = {"viz", "__pycache__"} +# learn/ is stdlib-only and would work in the browser, but training there is +# not the point and no trained model ships with the page, so bundling it would +# add weight to every visitor's download for a feature none of them can use. +SKIP_DIRS = {"viz", "learn", "__pycache__"} # id, title, blurb, feature tags, suggested (planner, heuristic). # diff --git a/tools/make_learn_promo.py b/tools/make_learn_promo.py new file mode 100644 index 0000000..59cc7c7 --- /dev/null +++ b/tools/make_learn_promo.py @@ -0,0 +1,1370 @@ +#!/usr/bin/env python3 +"""Render the promo video for the learned-heuristic / RL work. + +Like ``make_promo.py``, nothing on screen is typed in by hand: this trains a +heuristic, runs the reinforcement stage, reproduces both of the failure modes +that shaped its design, and animates whatever came back. If the method gets +better or worse, so does the video. + + python tools/make_learn_promo.py -o promo/jupyddl-rl.mp4 + +Collection takes a few minutes, so it is cached:: + + python tools/make_learn_promo.py --cache promo/rl-data.json # measure once + python tools/make_learn_promo.py --cache promo/rl-data.json # reuse + +Needs the ``viz`` extra plus an ffmpeg binary; ``learn`` (NumPy) makes the +collection pass much faster but is not required. +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +import time + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 +from matplotlib.animation import FFMpegWriter, PillowWriter # noqa: E402 +from matplotlib.path import Path # noqa: E402 +from matplotlib.patches import PathPatch # noqa: E402 + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from make_promo import ( # noqa: E402 + AQUA, + BG, + BLUE, + DIM, + DPI, + FPS, + GOOD, + INK, + MONO, + MUTED, + ORANGE, + SIZE, + YELLOW, + count_up, + ease_in_out, + ease_out, + fade_in, + fade_window, + fmt, + panel, + text, + typewriter, +) + +BAD = "#e05252" + +# -------------------------------------------------------------------------- +# data collection — every number the video shows is measured here +# -------------------------------------------------------------------------- +TRAIN_SIZES = range(3, 7) +HARD_SIZES = range(9, 13) +EVAL_SIZES = range(9, 14) + + +BUDGET = 30000 +TIME_LIMIT = 30.0 + + +def _transfer(model, tasks, baselines=("hff", "goalcount", "blind")): + """Benchmark ``model`` against ``baselines`` under a fixed budget.""" + from jupyddl.learn.pipeline import evaluate_transfer + + return evaluate_transfer( + model, tasks, baselines=baselines, max_expansions=BUDGET, time_limit=TIME_LIMIT + ) + + +def _families(note): + """The four disjoint seed families, and the vocabulary spanning them all. + + Four rather than three, and disjoint on purpose: training fits one, the + optimiser tunes on a second and is *selected* on a third, and the numbers + that appear on screen come from a fourth that no stage ever touched. + """ + from jupyddl.learn import tasks_from_generator + from jupyddl.learn.features import FeatureSpace + + families = { + "train": tasks_from_generator( + "blocksworld", TRAIN_SIZES, seed=0, seeds_per_size=3 + ), + "tune": tasks_from_generator( + "blocksworld", HARD_SIZES, seed=1000, seeds_per_size=2 + ), + "val": tasks_from_generator( + "blocksworld", HARD_SIZES, seed=2000, seeds_per_size=2 + ), + "eval": tasks_from_generator( + "blocksworld", EVAL_SIZES, seed=7777, seeds_per_size=2 + ), + } + every = [task for group in families.values() for _, task in group] + space = FeatureSpace.from_tasks(every) + note(f"ladder: {len(families['train'])} train, {len(families['eval'])} eval") + return families, space + + +def _demo_plan(task, name): + """One plan, annotated with the cost-to-go each state on it reveals.""" + from jupyddl.api import solve_task + + result = solve_task(task, "astar", "lmcut", time_limit=TIME_LIMIT) + steps = [op.base_name for op in task.visible_plan(result.plan)] + suffix = [0.0] * (len(result.plan) + 1) + for index in range(len(result.plan) - 1, -1, -1): + suffix[index] = suffix[index + 1] + result.plan[index].cost + return {"instance": name, "steps": steps, "suffix": suffix[: len(steps)]} + + +def _imitate(families, space, note): + """Corpus, one annotated plan, and the supervised fit. Returns the bundle.""" + from jupyddl.learn import TrainConfig, solved_corpus, train + from jupyddl.learn.pipeline import summarise_transfer + + corpus = solved_corpus(families["train"], space=space) + stats = corpus.target_stats() + note(f"corpus: {stats['count']} samples, {stats['groups']} groups") + + bundle, report = train(corpus, TrainConfig(epochs=60, seed=0)) + note( + f"imitation: top-1 {report.metrics['top1']:.3f}, MAE {report.metrics['mae']:.2f}" + ) + + name, task = families["train"][2] + section = { + "corpus": stats, + "plan": _demo_plan(task, name), + "imitation": { + "mae": report.metrics["mae"], + "top1": report.metrics["top1"], + "in_top2": report.metrics.get("in_top2", 0.0), + "best_epoch": report.best_epoch, + "seconds": report.seconds, + "parameters": bundle.model.num_parameters, + "history": [ + {"epoch": h["epoch"], "top1": h["train_top1"], "loss": h["loss"]} + for h in report.history + ], + }, + "transfer_before": summarise_transfer(_transfer(bundle, families["eval"])), + } + note( + "imitation on eval: " + f"{section['transfer_before']['learned']['mean_expanded']:.0f} expanded" + ) + return bundle, section + + +def _reinforce(bundle, families, rl, note): + """The reinforcement stage as it is meant to be run. Returns the model.""" + from jupyddl.learn.pipeline import summarise_transfer + from jupyddl.learn.rl import optimise_search_cost + + started = time.perf_counter() + tuned, history = optimise_search_cost( + bundle, + families["tune"], + iterations=10, + population=12, + config=rl, + validation_tasks=families["val"], + ) + section = { + "cem": { + "seconds": time.perf_counter() - started, + "iterations": len(history) - 1, + "population": 12, + "history": [ + { + "iteration": h["iteration"], + "tuning": h.get("tuning_score"), + "validation": h["score"], + "coverage": h["coverage"], + } + for h in history + ], + }, + "transfer_after": summarise_transfer(_transfer(tuned, families["eval"])), + } + note( + f"cem: {section['cem']['seconds']:.0f}s, eval " + f"{section['transfer_after']['learned']['mean_expanded']:.0f} expanded" + ) + return tuned, section + + +def _flat_objective(bundle, tuned, families, rl, note): + """Trap one: nothing to optimise where the search is already good.""" + from jupyddl.learn.rl import optimise_search_cost, search_cost + + before = search_cost(bundle, families["train"], rl).score + retuned, _ = optimise_search_cost( + bundle, families["train"], iterations=6, population=8, config=rl + ) + flat = { + "easy_before": before, + "easy_after": search_cost(retuned, families["train"], rl).score, + "hard_before": search_cost(bundle, families["tune"], rl).score, + "hard_after": search_cost(tuned, families["tune"], rl).score, + } + note( + f"flat: easy {flat['easy_before']:.1f} -> {flat['easy_after']:.1f}, " + f"hard {flat['hard_before']:.0f} -> {flat['hard_after']:.0f}" + ) + return flat + + +def _spread(bundle, families, rl, note): + """Trap two: the per-instance distribution behind the headline mean. + + This exists because the headline nearly went out wrong. Two things changed + in one edit — a validation split, and sigma 0.05 -> 0.15 — and the + improvement was credited to the first. Varying one at a time shows sigma + doing all of it, and the per-instance breakdown shows why: nine of ten + instances barely move, and the tenth decides the average. Both arms are + measured here so the video cannot drift from the claim. + """ + from jupyddl.learn.pipeline import summarise_transfer + from jupyddl.learn.rl import optimise_search_cost + + variants = {} + for label, sigma in (("lo", 0.05), ("hi", 0.15)): + model, history = optimise_search_cost( + bundle, + families["tune"], + iterations=10, + population=12, + sigma=sigma, + config=rl, + validation_tasks=families["val"], + ) + rows = _transfer(model, families["eval"], baselines=()) + variants[label] = { + "sigma": sigma, + "tuning": history[-1].get("tuning_score"), + "per_instance": {r["instance"]: r["expanded"] for r in rows}, + "solved": {r["instance"]: r["solved"] for r in rows}, + "mean": summarise_transfer(rows)["learned"]["mean_expanded"], + } + base = _transfer(bundle, families["eval"], baselines=()) + note( + f"spread: sigma 0.05 mean {variants['lo']['mean']:.0f}, " + f"sigma 0.15 mean {variants['hi']['mean']:.0f}" + ) + return { + "instances": [r["instance"] for r in base], + "imitation": {r["instance"]: r["expanded"] for r in base}, + "imitation_solved": {r["instance"]: r["solved"] for r in base}, + "lo": variants["lo"], + "hi": variants["hi"], + "budget": BUDGET, + } + + +def _logistics(blocks_top1, note): + """The domain where this loses, and the feature space that explains it.""" + from jupyddl.learn import TrainConfig, solved_corpus, tasks_from_generator, train + from jupyddl.learn.features import FeatureSpace + from jupyddl.learn.pipeline import summarise_transfer + + train_tasks = tasks_from_generator( + "logistics", range(2, 5), seed=0, seeds_per_size=3 + ) + eval_tasks = tasks_from_generator( + "logistics", range(5, 8), seed=7777, seeds_per_size=2 + ) + space = FeatureSpace.from_tasks(t for _, t in train_tasks + eval_tasks) + corpus = solved_corpus(train_tasks, space=space, time_limit=TIME_LIMIT) + bundle, report = train(corpus, TrainConfig(epochs=60, seed=0)) + summary = summarise_transfer(_transfer(bundle, eval_tasks)) + note( + f"logistics: {len(space.vocabulary)} predicates, top-1 " + f"{report.metrics['top1']:.3f}, {summary['learned']['mean_expanded']:.0f} " + f"vs hff {summary['hff']['mean_expanded']:.0f}" + ) + return { + "predicates": list(space.vocabulary), + "features": space.size, + "top1": report.metrics["top1"], + "learned": summary["learned"]["mean_expanded"], + "hff": summary["hff"]["mean_expanded"], + "blocks_top1": blocks_top1, + } + + +def collect() -> dict: + """Train, reinforce, and reproduce both failure modes. Returns everything.""" + from jupyddl.learn.rl import RLConfig + + started = time.perf_counter() + + def note(message): + print(f" [{time.perf_counter() - started:5.1f}s] {message}") + + families, space = _families(note) + rl = RLConfig(max_expansions=BUDGET, seed=0) + + data: dict = { + "space": { + "predicates": list(space.vocabulary), + "features": space.size, + "train_instances": len(families["train"]), + "train_sizes": [min(TRAIN_SIZES), max(TRAIN_SIZES)], + "eval_sizes": [min(EVAL_SIZES), max(EVAL_SIZES)], + } + } + + bundle, imitation = _imitate(families, space, note) + data.update(imitation) + + tuned, reinforced = _reinforce(bundle, families, rl, note) + data.update(reinforced) + + data["flat"] = _flat_objective(bundle, tuned, families, rl, note) + data["spread"] = _spread(bundle, families, rl, note) + data["logistics"] = _logistics(data["imitation"]["top1"], note) + + data["total_seconds"] = time.perf_counter() - started + return data + + +# -------------------------------------------------------------------------- +# drawing helpers specific to this video +# -------------------------------------------------------------------------- +def _header(fig, title, subtitle, alpha, y=0.90): + text(fig, 0.08, y, title, size=42, weight="bold", alpha=alpha) + if subtitle: + text(fig, 0.08, y - 0.075, subtitle, size=21, color=DIM, alpha=alpha * 0.95) + + +def _axes(fig, rect, alpha): + ax = fig.add_axes(rect) + ax.set_facecolor("none") + for spine in ("top", "right"): + ax.spines[spine].set_visible(False) + for spine in ("left", "bottom"): + ax.spines[spine].set_color("#3a3a3d") + ax.spines[spine].set_alpha(alpha) + ax.spines[spine].set_linewidth(1.0) + ax.tick_params(colors=MUTED, labelsize=13, length=3, width=1.0) + for label in ax.get_xticklabels() + ax.get_yticklabels(): + label.set_alpha(alpha) + ax.grid(True, axis="y", color="#26262a", linewidth=1.0, alpha=alpha * 0.7) + ax.set_axisbelow(True) + return ax + + +def _rounded_bar(ax, x, width, height, color, alpha, radius_frac=0.35): + """A bar with rounded data-end and a square foot on the baseline. + + Rounding both ends turns a bar into a pill and detaches it from the axis it + is measured against; only the free end gets a radius. + """ + if height <= 0: + return + span = ax.get_ylim()[1] - ax.get_ylim()[0] or 1.0 + radius = min(width * radius_frac, height * 0.5, span * 0.03) + left, right = x - width / 2, x + width / 2 + verts = [ + (left, 0.0), + (left, height - radius), + (left, height - radius * 0.45), + (left + radius * 0.45, height), + (left + radius, height), + (right - radius, height), + (right - radius * 0.45, height), + (right, height - radius * 0.45), + (right, height - radius), + (right, 0.0), + (left, 0.0), + ] + codes = [ + Path.MOVETO, + Path.LINETO, + Path.CURVE4, + Path.CURVE4, + Path.CURVE4, + Path.LINETO, + Path.CURVE4, + Path.CURVE4, + Path.CURVE4, + Path.LINETO, + Path.CLOSEPOLY, + ] + ax.add_patch( + PathPatch(Path(verts, codes), facecolor=color, edgecolor="none", alpha=alpha) + ) + + +def _stat(fig, x, y, value, label, colour=INK, alpha=1.0, size=64, sub=None): + text(fig, x, y, value, size=size, weight="bold", color=colour, alpha=alpha) + text(fig, x, y - 0.085, label, size=19, color=DIM, alpha=alpha * 0.95) + if sub: + text(fig, x, y - 0.135, sub, size=16, color=MUTED, alpha=alpha * 0.9) + + +# -------------------------------------------------------------------------- +# scenes +# -------------------------------------------------------------------------- +def scene_hook(fig, t, data): + alpha = fade_window(t, 0.0, 1.0, 0.14) + text( + fig, + 0.5, + 0.60, + "Your planner already wrote", + size=58, + weight="bold", + alpha=alpha * fade_in(t, 0.05), + ha="center", + ) + text( + fig, + 0.5, + 0.47, + "the training data.", + size=58, + weight="bold", + color=BLUE, + alpha=alpha * fade_in(t, 0.22), + ha="center", + ) + text( + fig, + 0.5, + 0.32, + "every solved plan is a labelled trajectory", + size=24, + color=DIM, + alpha=alpha * fade_in(t, 0.45), + ha="center", + ) + + +def scene_labels(fig, t, data): + alpha = fade_window(t, 0.0, 1.0, 0.10) + _header( + fig, + "A plan is supervision you already paid for.", + "the cost of the suffix from any state on it is that state's cost-to-go", + alpha, + ) + + plan = data["plan"] + steps = plan["steps"][:8] + suffix = plan["suffix"][:8] + panel(fig, 0.08, 0.16, 0.40, 0.58, alpha * fade_in(t, 0.12)) + + for index, (step, cost) in enumerate(zip(steps, suffix)): + row_alpha = alpha * fade_in(t, 0.18 + index * 0.045, 0.12) + y = 0.67 - index * 0.062 + text( + fig, + 0.115, + y, + f"{index + 1}.", + size=17, + color=MUTED, + alpha=row_alpha, + family=MONO, + ) + text(fig, 0.155, y, step, size=17, color=INK, alpha=row_alpha, family=MONO) + # The label the state carries, appearing beside the step that reveals it. + text( + fig, + 0.445, + y, + f"h* = {cost:.0f}", + size=17, + color=AQUA, + alpha=alpha * fade_in(t, 0.42 + index * 0.045, 0.12), + family=MONO, + ha="right", + ) + + right = alpha * fade_in(t, 0.62) + text(fig, 0.56, 0.62, "One instance,", size=30, color=DIM, alpha=right) + _stat( + fig, + 0.56, + 0.50, + f"{count_up(len(steps) + 1, t, 0.66, 0.3)}", + "labelled states", + colour=AQUA, + alpha=right, + size=58, + ) + stats = data["corpus"] + _stat( + fig, + 0.56, + 0.28, + fmt(stats["count"]), + f"across {data['space']['train_instances']} instances of 3-6 blocks", + colour=INK, + alpha=alpha * fade_in(t, 0.78), + size=58, + ) + text( + fig, + 0.56, + 0.13, + "no labelling, no annotation, no human", + size=19, + color=MUTED, + alpha=alpha * fade_in(t, 0.86), + ) + + +def scene_imitation(fig, t, data): + alpha = fade_window(t, 0.0, 1.0, 0.10) + _header( + fig, + "Fit a network to it. That is imitation.", + "a thousand parameters over fourteen features, trained in a tenth of a second", + alpha, + ) + + history = data["imitation"]["history"] + if not history: # pragma: no cover - defensive + return + reveal = min(1.0, max(0.0, (t - 0.20) / 0.45)) + shown = max(2, int(len(history) * ease_out(reveal))) + + ax = _axes(fig, [0.09, 0.20, 0.48, 0.46], alpha * fade_in(t, 0.15)) + xs = [h["epoch"] for h in history[:shown]] + ys = [h["top1"] for h in history[:shown]] + ax.plot(xs, ys, color=BLUE, linewidth=2.0, solid_capstyle="round") + ax.scatter([xs[-1]], [ys[-1]], s=64, color=BLUE, zorder=5) + ax.set_xlim(0, len(history) + 1) + ax.set_ylim(0, 1.02) + ax.set_xlabel("epoch", color=MUTED, fontsize=14) + ax.set_ylabel("top-1 ranking accuracy", color=MUTED, fontsize=14) + # One series: the axis label names it, so no legend box. + ax.annotate( + f"{ys[-1]:.3f}", + (xs[-1], ys[-1]), + textcoords="offset points", + xytext=(12, -4), + color=BLUE, + fontsize=17, + fontweight="bold", + alpha=alpha, + ) + + right = alpha * fade_in(t, 0.55) + _stat( + fig, + 0.65, + 0.56, + f"{data['imitation']['mae']:.2f}", + "mean absolute error, held out", + colour=INK, + alpha=right, + size=56, + ) + _stat( + fig, + 0.65, + 0.32, + f"{data['imitation']['top1']:.0%}", + "of decisions ranked correctly", + colour=BLUE, + alpha=alpha * fade_in(t, 0.70), + size=56, + ) + text( + fig, + 0.65, + 0.16, + f"{data['imitation']['seconds']:.1f}s of training", + size=18, + color=MUTED, + alpha=alpha * fade_in(t, 0.82), + ) + + +def scene_order(fig, t, data): + alpha = fade_window(t, 0.0, 1.0, 0.10) + _header( + fig, + "But search never reads the number.", + "greedy best-first search pops the minimum — it reads the order", + alpha, + ) + + # Two heuristics, same decision point. Deliberately schematic: the claim is + # about what the search consumes, not about any one measured value. + left_alpha = alpha * fade_in(t, 0.16) + panel(fig, 0.07, 0.20, 0.40, 0.48, left_alpha) + text( + fig, 0.10, 0.615, "h + 30, everywhere", size=24, weight="bold", alpha=left_alpha + ) + text(fig, 0.10, 0.565, "terrible RMSE", size=17, color=BAD, alpha=left_alpha) + for index, (label, value) in enumerate([("successor A", 32), ("successor B", 35)]): + row = left_alpha * fade_in(t, 0.26 + index * 0.06, 0.12) + y = 0.47 - index * 0.085 + text(fig, 0.10, y, label, size=19, color=DIM, alpha=row, family=MONO) + text( + fig, + 0.42, + y, + str(value), + size=19, + color=INK, + alpha=row, + family=MONO, + ha="right", + ) + text( + fig, + 0.10, + 0.26, + "→ picks A. Perfect guidance.", + size=20, + color=GOOD, + alpha=left_alpha * fade_in(t, 0.42), + ) + + right_alpha = alpha * fade_in(t, 0.50) + panel(fig, 0.53, 0.20, 0.40, 0.48, right_alpha) + text( + fig, + 0.56, + 0.615, + "accurate, two siblings swapped", + size=24, + weight="bold", + alpha=right_alpha, + ) + text(fig, 0.56, 0.565, "excellent RMSE", size=17, color=GOOD, alpha=right_alpha) + for index, (label, value) in enumerate( + [("successor A", 5.2), ("successor B", 4.8)] + ): + row = right_alpha * fade_in(t, 0.58 + index * 0.06, 0.12) + y = 0.47 - index * 0.085 + text(fig, 0.56, y, label, size=19, color=DIM, alpha=row, family=MONO) + text( + fig, + 0.88, + y, + f"{value}", + size=19, + color=INK, + alpha=row, + family=MONO, + ha="right", + ) + text( + fig, + 0.56, + 0.26, + "→ picks B. Wrong subtree.", + size=20, + color=BAD, + alpha=right_alpha * fade_in(t, 0.72), + ) + + text( + fig, + 0.5, + 0.10, + "so the loss optimises the ordering, not the estimate", + size=23, + color=DIM, + alpha=alpha * fade_in(t, 0.84), + ha="center", + style="italic", + ) + + +def scene_turn(fig, t, data): + alpha = fade_window(t, 0.0, 1.0, 0.12) + text( + fig, + 0.5, + 0.78, + "Imitation optimises a proxy.", + size=44, + weight="bold", + alpha=alpha * fade_in(t, 0.04), + ha="center", + ) + text( + fig, + 0.5, + 0.68, + "What we want is the heuristic that expands the fewest nodes.", + size=24, + color=DIM, + alpha=alpha * fade_in(t, 0.18), + ha="center", + ) + + rows = [ + ("state", "the open list, the closed set"), + ("action", "which node to expand next"), + ("policy", "argmin over h"), + ("reward", "−1 per expansion"), + ] + panel(fig, 0.22, 0.20, 0.56, 0.38, alpha * fade_in(t, 0.32)) + for index, (key, value) in enumerate(rows): + row = alpha * fade_in(t, 0.38 + index * 0.075, 0.14) + y = 0.50 - index * 0.075 + colour = YELLOW if key == "reward" else INK + text(fig, 0.26, y, key, size=21, color=DIM, alpha=row, family=MONO) + text(fig, 0.40, y, value, size=21, color=colour, alpha=row, family=MONO) + + text( + fig, + 0.5, + 0.11, + "That is not a proxy. That is reinforcement learning.", + size=25, + color=BLUE, + weight="bold", + alpha=alpha * fade_in(t, 0.76), + ha="center", + ) + + +def scene_cem(fig, t, data): + alpha = fade_window(t, 0.0, 1.0, 0.10) + _header( + fig, + "No gradient. So don't use one.", + "perturb the weights, run the planner, keep what searched least, repeat", + alpha, + ) + + history = data["cem"]["history"] + reveal = min(1.0, max(0.0, (t - 0.18) / 0.55)) + shown = max(1, int(round(len(history) * ease_in_out(reveal)))) + + ax = _axes(fig, [0.09, 0.20, 0.52, 0.46], alpha * fade_in(t, 0.14)) + xs = [h["iteration"] for h in history[:shown]] + tuning = [ + h["tuning"] if h["tuning"] is not None else h["validation"] + for h in history[:shown] + ] + validation = [h["validation"] for h in history[:shown]] + # Both series are mean expansions — the same unit, so one axis is correct. + ax.plot(xs, tuning, color=ORANGE, linewidth=2.0, solid_capstyle="round") + ax.plot(xs, validation, color=BLUE, linewidth=2.0, solid_capstyle="round") + ax.scatter([xs[-1]], [tuning[-1]], s=64, color=ORANGE, zorder=5) + ax.scatter([xs[-1]], [validation[-1]], s=64, color=BLUE, zorder=5) + ax.set_xlim(-0.3, len(history) - 0.4) + top = max(max(tuning), max(validation)) * 1.15 + ax.set_ylim(0, top) + ax.set_xlabel("CEM iteration", color=MUTED, fontsize=14) + ax.set_ylabel("mean nodes expanded", color=MUTED, fontsize=14) + # Two series: direct-labelled rather than boxed, so identity is never colour alone. + ax.annotate( + f"tuning {tuning[-1]:.0f}", + (xs[-1], tuning[-1]), + textcoords="offset points", + xytext=(10, 6), + color=ORANGE, + fontsize=16, + fontweight="bold", + alpha=alpha, + ) + ax.annotate( + f"validation {validation[-1]:.0f}", + (xs[-1], validation[-1]), + textcoords="offset points", + xytext=(10, -14), + color=BLUE, + fontsize=16, + fontweight="bold", + alpha=alpha, + ) + + right = alpha * fade_in(t, 0.60) + text(fig, 0.68, 0.60, "cross-entropy method", size=24, weight="bold", alpha=right) + text( + fig, + 0.68, + 0.545, + f"{data['cem']['population']} candidates x {data['cem']['iterations']} rounds", + size=19, + color=DIM, + alpha=right, + ) + _stat( + fig, + 0.68, + 0.38, + f"{data['cem']['seconds']:.0f}s", + "of planning, as the objective", + colour=AQUA, + alpha=alpha * fade_in(t, 0.74), + size=50, + ) + text( + fig, + 0.68, + 0.19, + "the planner is the environment", + size=19, + color=MUTED, + alpha=alpha * fade_in(t, 0.86), + ) + + +def scene_flat(fig, t, data): + alpha = fade_window(t, 0.0, 1.0, 0.10) + _header( + fig, + "Trap one: the objective is flat.", + "on instances search already solves well, every perturbation scores the same", + alpha, + ) + + flat = data["flat"] + blocks = [ + ( + "tuned on the training ladder", + flat["easy_before"], + flat["easy_after"], + BAD, + "no headroom left", + 0.16, + ), + ( + "tuned a rung higher", + flat["hard_before"], + flat["hard_after"], + GOOD, + "room to be wrong in", + 0.46, + ), + ] + for label, before, after, colour, caption, start in blocks: + block = alpha * fade_in(t, start) + y = 0.56 if start < 0.3 else 0.30 + text(fig, 0.09, y + 0.06, label, size=23, color=DIM, alpha=block) + text( + fig, + 0.09, + y - 0.02, + f"{before:,.0f}", + size=46, + weight="bold", + color=MUTED, + alpha=block, + family=MONO, + ) + text( + fig, + 0.28, + y - 0.02, + "→", + size=34, + color=MUTED, + alpha=block * fade_in(t, start + 0.08), + ) + text( + fig, + 0.36, + y - 0.02, + f"{after:,.0f}", + size=46, + weight="bold", + color=colour, + alpha=block * fade_in(t, start + 0.12), + family=MONO, + ) + text( + fig, + 0.56, + y - 0.02, + caption, + size=20, + color=colour, + alpha=block * fade_in(t, start + 0.16), + ) + + text( + fig, + 0.5, + 0.10, + "optimise where the search is still bad", + size=25, + color=INK, + weight="bold", + alpha=alpha * fade_in(t, 0.80), + ha="center", + ) + + +def scene_spread(fig, t, data): + """The distribution behind the mean — and the instance that decides it.""" + alpha = fade_window(t, 0.0, 1.0, 0.10) + _header( + fig, + "Trap two: the mean was doing the lying.", + "same run, one setting changed — and nine of ten instances barely notice", + alpha, + ) + + spread = data["spread"] + names = spread["instances"] + lo, hi = spread["lo"], spread["hi"] + budget = spread["budget"] + + ax = _axes(fig, [0.085, 0.20, 0.60, 0.47], alpha * fade_in(t, 0.14)) + xs = list(range(len(names))) + base = [spread["imitation"][n] for n in names] + lo_vals = [lo["per_instance"][n] for n in names] + hi_vals = [hi["per_instance"][n] for n in names] + + reveal = ease_out(min(1.0, max(0.0, (t - 0.18) / 0.34))) + shown = max(1, int(round(len(names) * reveal))) + ax.set_yscale("log") + ax.set_ylim(30, budget * 2.2) + ax.set_xlim(-0.6, len(names) - 0.4) + # Log scale because the tail spans three orders of magnitude; a linear axis + # would render nine of the ten instances as a flat line on the floor. + ax.plot( + xs[:shown], base[:shown], color=MUTED, linewidth=2.0, marker="o", markersize=7 + ) + ax.plot( + xs[:shown], + lo_vals[:shown], + color=ORANGE, + linewidth=2.0, + marker="o", + markersize=7, + ) + ax.plot( + xs[:shown], hi_vals[:shown], color=BLUE, linewidth=2.0, marker="o", markersize=7 + ) + ax.set_xticks(xs) + ax.set_xticklabels([n.split("-")[1] for n in names], color=DIM, fontsize=13) + ax.set_xlabel("held-out instance, by number of blocks", color=MUTED, fontsize=14) + ax.set_ylabel("nodes expanded (log)", color=MUTED, fontsize=14) + + legend = alpha * fade_in(t, 0.30) + for index, (label, colour) in enumerate( + [ + ("imitation", MUTED), + (f"sigma {lo['sigma']}", ORANGE), + (f"sigma {hi['sigma']}", BLUE), + ] + ): + y = 0.615 - index * 0.045 + ax.figure.patches.append( + plt.Rectangle( + (0.115, y - 0.008), + 0.016, + 0.016, + transform=fig.transFigure, + facecolor=colour, + edgecolor="none", + alpha=legend, + ) + ) + text(fig, 0.140, y, label, size=15, color=DIM, alpha=legend) + + # The outlier, called out where it sits. + worst = max(names, key=lambda n: spread["imitation"][n]) + if t > 0.52: + marker = alpha * fade_in(t, 0.52) + index = names.index(worst) + ax.annotate( + "imitation never solved this one", + (index, spread["imitation"][worst]), + textcoords="offset points", + xytext=(-150, 14), + color=BAD, + fontsize=15, + fontweight="bold", + alpha=marker, + ) + + right = alpha * fade_in(t, 0.64) + text(fig, 0.72, 0.60, "the two means", size=21, color=DIM, alpha=right) + for row, (value, colour, tag) in enumerate( + [ + (lo["mean"], ORANGE, f"sigma {lo['sigma']}"), + (hi["mean"], BLUE, f"sigma {hi['sigma']}"), + ] + ): + y = 0.535 - row * 0.062 + text( + fig, + 0.72, + y, + f"{value:>6,.0f}", + size=30, + weight="bold", + color=INK, + alpha=right, + family=MONO, + ) + text(fig, 0.845, y, tag, size=16, color=colour, alpha=right) + text( + fig, + 0.72, + 0.395, + f"differ by {lo['mean'] / max(1.0, hi['mean']):.0f}x.", + size=22, + color=DIM, + alpha=alpha * fade_in(t, 0.72), + ) + text( + fig, + 0.72, + 0.32, + "One instance out of ten\nis the whole gap.", + size=22, + color=BAD, + weight="bold", + alpha=alpha * fade_in(t, 0.80), + ) + text( + fig, + 0.72, + 0.19, + "look at the distribution\nbefore believing the mean", + size=19, + color=GOOD, + alpha=alpha * fade_in(t, 0.88), + ) + + +def scene_result(fig, t, data): + alpha = fade_window(t, 0.0, 1.0, 0.10) + space = data["space"] + _header( + fig, + "Trained on 3-6 blocks. Judged on 9-13.", + "greedy best-first search, on instances no stage of training ever saw", + alpha, + ) + + after = data["transfer_after"] + order = ["learned", "hff", "goalcount"] + colours = {"learned": BLUE, "hff": MUTED, "goalcount": "#4a4a4e"} + + # Two measures on different scales get two charts. Never two y-axes. + for panel_index, (key, title, fmt_value) in enumerate( + [ + ("mean_expanded", "nodes expanded", lambda v: f"{v:,.0f}"), + ("mean_seconds", "seconds", lambda v: f"{v:.3f}"), + ] + ): + left = 0.09 + panel_index * 0.47 + ax = _axes( + fig, [left, 0.22, 0.36, 0.44], alpha * fade_in(t, 0.14 + panel_index * 0.06) + ) + values = [after[name][key] for name in order] + ax.set_xlim(-0.65, 2.65) + ax.set_ylim(0, max(values) * 1.24) + for index, name in enumerate(order): + grow = ease_out( + min( + 1.0, max(0.0, (t - 0.22 - index * 0.09 - panel_index * 0.05) / 0.32) + ) + ) + _rounded_bar(ax, index, 0.52, values[index] * grow, colours[name], alpha) + if grow > 0.6: + ax.text( + index, + values[index] * grow + max(values) * 0.04, + fmt_value(values[index]), + ha="center", + color=colours[name] if name == "learned" else DIM, + fontsize=19, + fontweight="bold", + alpha=alpha * ease_out((grow - 0.6) / 0.4), + ) + ax.set_xticks(range(3)) + ax.set_xticklabels(order, color=DIM, fontsize=15) + ax.set_title(title, color=MUTED, fontsize=16, pad=12) + + ratio = after["hff"]["mean_expanded"] / max(1e-9, after["learned"]["mean_expanded"]) + speed = after["hff"]["mean_seconds"] / max(1e-9, after["learned"]["mean_seconds"]) + text( + fig, + 0.5, + 0.11, + f"{ratio:.1f}x fewer expansions than h_ff, and {speed:.0f}x faster", + size=27, + weight="bold", + color=INK, + alpha=alpha * fade_in(t, 0.72), + ha="center", + ) + text( + fig, + 0.5, + 0.05, + f"trained on {space['train_instances']} instances in " + f"{data['imitation']['seconds']:.1f}s", + size=17, + color=MUTED, + alpha=alpha * fade_in(t, 0.84), + ha="center", + ) + + +def scene_honest(fig, t, data): + alpha = fade_window(t, 0.0, 1.0, 0.10) + log = data["logistics"] + _header( + fig, + "And on logistics, it loses.", + "worth saying, because the reason is exact rather than mysterious", + alpha, + ) + + left = alpha * fade_in(t, 0.14) + text(fig, 0.09, 0.60, "the domain has", size=22, color=DIM, alpha=left) + text( + fig, + 0.09, + 0.50, + " ".join(f"({p} ...)" for p in log["predicates"]), + size=30, + color=INK, + alpha=left, + family=MONO, + ) + text( + fig, + 0.09, + 0.41, + f"two predicates → {log['features']} features", + size=21, + color=DIM, + alpha=alpha * fade_in(t, 0.30), + ) + text( + fig, + 0.09, + 0.30, + "which cannot say which package is where —", + size=23, + color=BAD, + alpha=alpha * fade_in(t, 0.44), + ) + text( + fig, + 0.09, + 0.235, + "only how many are somewhere at all.", + size=23, + color=BAD, + alpha=alpha * fade_in(t, 0.52), + ) + + right = alpha * fade_in(t, 0.64) + _stat( + fig, + 0.68, + 0.56, + f"{log['top1']:.0%}", + "of decisions ranked correctly", + colour=BAD, + alpha=right, + size=54, + sub=f"against {log['blocks_top1']:.0%} on blocksworld", + ) + _stat( + fig, + 0.68, + 0.30, + f"{log['learned'] / max(1.0, log['hff']):.0f}x", + "more expansions than h_ff", + colour=DIM, + alpha=alpha * fade_in(t, 0.78), + size=54, + ) + text( + fig, + 0.5, + 0.09, + "counting is blind to topology — that is the case for relational features", + size=21, + color=MUTED, + alpha=alpha * fade_in(t, 0.88), + ha="center", + style="italic", + ) + + +def scene_cta(fig, t, data): + alpha = fade_window(t, 0.0, 1.0, 0.14) + text( + fig, + 0.5, + 0.70, + "jupyddl learn", + size=60, + weight="bold", + alpha=alpha * fade_in(t, 0.05), + ha="center", + family=MONO, + ) + command = "jupyddl learn blocksworld --cem 10 --evaluate 9-13" + panel(fig, 0.14, 0.44, 0.72, 0.11, alpha * fade_in(t, 0.22)) + text( + fig, + 0.5, + 0.495, + typewriter(command, t, 0.26, 0.34), + size=24, + color=AQUA, + alpha=alpha * fade_in(t, 0.26), + ha="center", + family=MONO, + ) + text( + fig, + 0.5, + 0.33, + "-H learned:model.json · anywhere a heuristic name goes", + size=21, + color=DIM, + alpha=alpha * fade_in(t, 0.62), + ha="center", + family=MONO, + ) + text( + fig, + 0.5, + 0.22, + "github.com/APLA-Toolbox/PythonPDDL", + size=22, + color=INK, + alpha=alpha * fade_in(t, 0.72), + ha="center", + family=MONO, + ) + text( + fig, + 0.5, + 0.14, + "zero dependencies · the research notes are in .docs/", + size=18, + color=MUTED, + alpha=alpha * fade_in(t, 0.80), + ha="center", + ) + + +SCENES = [ + (scene_hook, 5.0), + (scene_labels, 9.0), + (scene_imitation, 8.5), + (scene_order, 9.5), + (scene_turn, 8.0), + (scene_cem, 11.0), + (scene_flat, 9.0), + (scene_spread, 11.0), + (scene_result, 10.0), + (scene_honest, 10.0), + (scene_cta, 6.5), +] + + +def render(data, out: str, fps: int = FPS, dpi: int = DPI): + fig = plt.figure(figsize=SIZE, facecolor=BG) + total = sum(int(seconds * fps) for _, seconds in SCENES) + + if out.lower().endswith(".gif"): + writer = PillowWriter(fps=fps) + else: + if not FFMpegWriter.isAvailable(): + try: + import imageio_ffmpeg + + matplotlib.rcParams["animation.ffmpeg_path"] = ( + imageio_ffmpeg.get_ffmpeg_exe() + ) + except Exception: + pass + if not FFMpegWriter.isAvailable(): + raise RuntimeError( + "ffmpeg not found. pip install imageio-ffmpeg, or render a .gif" + ) + writer = FFMpegWriter( + fps=fps, + bitrate=-1, + extra_args=["-pix_fmt", "yuv420p", "-crf", "20", "-preset", "medium"], + ) + + started = time.perf_counter() + done = 0 + print(f"Rendering {total} frames ({total / fps:.1f}s) -> {out}") + with writer.saving(fig, out, dpi=dpi): + for draw, seconds in SCENES: + frames = int(seconds * fps) + for frame in range(frames): + fig.clear() + fig.patches.clear() + draw(fig, frame / max(1, frames - 1), data) + writer.grab_frame(facecolor=BG) + done += 1 + if done % 60 == 0: + rate = done / (time.perf_counter() - started) + print(f" {done}/{total} frames ({rate:.0f} fps render)") + plt.close(fig) + print(f"Wrote {out} in {time.perf_counter() - started:.0f}s") + + +def main() -> int: + parser = argparse.ArgumentParser( + description="Render the learned-heuristic / RL promo video." + ) + parser.add_argument("-o", "--output", default="promo/jupyddl-rl.mp4") + parser.add_argument("--fps", type=int, default=FPS) + parser.add_argument("--dpi", type=int, default=DPI) + parser.add_argument( + "--cache", + default=None, + help="read measurements from this JSON if it exists, else write them to it", + ) + parser.add_argument( + "--collect-only", + action="store_true", + help="measure and write the cache without rendering", + ) + args = parser.parse_args() + + if args.cache and os.path.exists(args.cache): + print(f"Reusing measurements from {args.cache}") + with open(args.cache, encoding="utf-8") as handle: + data = json.load(handle) + else: + print("Measuring (training, reinforcing, and reproducing both traps)...") + data = collect() + if args.cache: + os.makedirs(os.path.dirname(os.path.abspath(args.cache)), exist_ok=True) + with open(args.cache, "w", encoding="utf-8") as handle: + json.dump(data, handle, indent=1) + print(f"Wrote {args.cache} ({data['total_seconds']:.0f}s of measurement)") + + if args.collect_only: + return 0 + + folder = os.path.dirname(os.path.abspath(args.output)) + os.makedirs(folder, exist_ok=True) + render(data, args.output, fps=args.fps, dpi=args.dpi) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/web/dist/jupyddl-sources.json b/web/dist/jupyddl-sources.json index c1f12b8..9443289 100644 --- a/web/dist/jupyddl-sources.json +++ b/web/dist/jupyddl-sources.json @@ -1 +1 @@ -{"jupyddl/__init__.py": "\"\"\"jupyddl: a pure-Python PDDL planning framework.\n\nQuickstart::\n\n from jupyddl import solve, build_task, trace_search, validate_plan\n\n result = solve(\"domain.pddl\", \"problem.pddl\", search=\"astar\", heuristic=\"lmcut\")\n print(result.solved, result.cost, result.plan_names())\n\n # ...and watch the search itself\n task = build_task(\"domain.pddl\", \"problem.pddl\")\n result, trace = trace_search(task, \"astar\", \"lmcut\")\n trace.save(\"run.json\")\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom .api import build_task, solve, solve_task, trace_search, validate_plan\nfrom .grounding import ground, ground_files\nfrom .heuristics import HEURISTICS, make_heuristic\nfrom .parser import PDDLError, UnsupportedFeatureError, parse\nfrom .search import PLANNERS, SearchResult, make_planner\nfrom .task import Operator, Task\nfrom .trace import (\n MultiObserver,\n SearchEvent,\n SearchObserver,\n SearchTrace,\n TraceRecorder,\n)\n\n__version__ = \"2.2.0\"\n\n__all__ = [\n \"solve\",\n \"solve_task\",\n \"build_task\",\n \"trace_search\",\n \"validate_plan\",\n \"ground\",\n \"ground_files\",\n \"make_planner\",\n \"make_heuristic\",\n \"PLANNERS\",\n \"HEURISTICS\",\n \"SearchResult\",\n \"SearchTrace\",\n \"SearchEvent\",\n \"SearchObserver\",\n \"MultiObserver\",\n \"TraceRecorder\",\n \"Task\",\n \"Operator\",\n \"parse\",\n \"PDDLError\",\n \"UnsupportedFeatureError\",\n \"__version__\",\n]\n", "jupyddl/api.py": "\"\"\"High-level convenience API: parse + ground + plan + validate.\"\"\"\n\nfrom __future__ import annotations\n\nfrom .grounding import ground_files\nfrom .heuristics import make_heuristic\nfrom .search import make_planner\nfrom .search.result import SearchResult, make_budget\nfrom .task import Task\nfrom .trace import TraceRecorder\n\n\ndef build_task(domain_path: str, problem_path: str) -> Task:\n \"\"\"Parse and ground a domain/problem pair into a :class:`Task`.\"\"\"\n return ground_files(domain_path, problem_path)\n\n\ndef solve_task(\n task: Task,\n search: str = \"astar\",\n heuristic=None,\n observer=None,\n max_expansions=None,\n time_limit=None,\n **planner_kwargs,\n) -> SearchResult:\n \"\"\"Run ``search`` (optionally with ``heuristic``) on an already-ground task.\n\n Pass ``observer`` (see :mod:`jupyddl.trace`) to record or live-render the\n search as it runs. ``max_expansions`` and ``time_limit`` bound the run; when\n either is hit the planner stops and sets ``result.truncated``, so an empty\n result means \"gave up\", not \"proved unsolvable\".\n \"\"\"\n planner = make_planner(search, **planner_kwargs)\n heur = None\n name = heuristic if heuristic else (\"hff\" if planner.requires_heuristic else None)\n if name is not None:\n heur = make_heuristic(name, task)\n budget = make_budget(max_expansions, time_limit)\n return planner.search(task, heur, observer=observer, budget=budget)\n\n\ndef solve(\n domain_path: str,\n problem_path: str,\n search: str = \"astar\",\n heuristic=\"lmcut\",\n observer=None,\n max_expansions=None,\n time_limit=None,\n **planner_kwargs,\n) -> SearchResult:\n \"\"\"Parse, ground and solve a PDDL instance in one call.\"\"\"\n task = build_task(domain_path, problem_path)\n return solve_task(\n task,\n search=search,\n heuristic=heuristic,\n observer=observer,\n max_expansions=max_expansions,\n time_limit=time_limit,\n **planner_kwargs,\n )\n\n\ndef trace_search(\n task: Task,\n search: str = \"astar\",\n heuristic=None,\n max_events: int = 20000,\n record_generated: bool = False,\n observer=None,\n max_expansions=None,\n time_limit=None,\n **planner_kwargs,\n):\n \"\"\"Solve ``task`` while recording the search.\n\n Returns ``(result, trace)`` where ``trace`` is a\n :class:`~jupyddl.trace.SearchTrace` ready to plot, save or replay. An extra\n ``observer`` (a live dashboard, say) is notified alongside the recorder.\n \"\"\"\n recorder = TraceRecorder(max_events=max_events, record_generated=record_generated)\n if observer is not None:\n from .trace import MultiObserver\n\n sink = MultiObserver(recorder, observer)\n else:\n sink = recorder\n result = solve_task(\n task,\n search=search,\n heuristic=heuristic,\n observer=sink,\n max_expansions=max_expansions,\n time_limit=time_limit,\n **planner_kwargs,\n )\n return result, recorder.trace\n\n\ndef validate_plan(task: Task, plan) -> bool:\n \"\"\"Return ``True`` iff applying ``plan`` from the initial state reaches the goal.\n\n Replays through the task rather than the raw operators, so derived\n predicates are closed and numeric fluents are carried at every step \u2014\n validating against ``task.init`` directly would see neither.\n \"\"\"\n state = task.initial_state()\n for op in plan or ():\n if not op.applicable(state):\n return False\n state = task.apply(op, state)\n return task.goal_reached(state)\n", "jupyddl/benchmark.py": "\"\"\"Comparative benchmarking of planners/heuristics over PDDL instances.\n\nExample::\n\n from jupyddl.benchmark import discover_instances, run_benchmark, to_csv\n rows = run_benchmark(discover_instances(\"pddl-examples\"),\n [(\"astar\", \"lmcut\"), (\"gbfs\", \"hff\")])\n to_csv(rows, \"results.csv\")\n\"\"\"\n\nfrom __future__ import annotations\n\nimport csv\nimport glob\nimport os\nfrom dataclasses import asdict, dataclass\nfrom typing import Optional\n\nfrom .api import build_task, solve_task, validate_plan\nfrom .parser import PDDLError\n\n\n@dataclass\nclass Instance:\n name: str\n domain: str\n problem: str\n\n\n@dataclass\nclass BenchmarkRow:\n instance: str\n planner: str\n heuristic: str\n solved: bool\n valid: bool\n cost: Optional[int]\n plan_length: Optional[int]\n expanded: int\n generated: int\n evaluated: int\n runtime: float\n error: str = \"\"\n truncated: bool = False\n\n\ndef discover_instances(root: str) -> list:\n \"\"\"Find ``/*/`` folders containing both ``domain.pddl`` and ``problem.pddl``.\"\"\"\n instances = []\n for domain in sorted(glob.glob(os.path.join(root, \"*\", \"domain.pddl\"))):\n folder = os.path.dirname(domain)\n problem = os.path.join(folder, \"problem.pddl\")\n if os.path.exists(problem):\n instances.append(Instance(os.path.basename(folder), domain, problem))\n return instances\n\n\ndef run_benchmark(instances, configs, max_expansions=None, time_limit=None) -> list:\n \"\"\"Run each ``(planner, heuristic[, kwargs])`` config on each instance.\n\n ``configs`` entries are ``(planner_name, heuristic_name_or_None)`` or\n ``(planner_name, heuristic_name_or_None, planner_kwargs)``.\n\n ``max_expansions`` and ``time_limit`` bound every individual run, which is\n what keeps one pathological instance from stalling a whole benchmark. A run\n that stops on its budget is recorded with ``truncated=True`` and does not\n count towards coverage -- \"we stopped looking\" is not \"no plan exists\".\n \"\"\"\n rows: list = []\n for inst in instances:\n try:\n task = build_task(inst.domain, inst.problem)\n except (PDDLError, ValueError) as exc:\n for cfg in configs:\n planner, heuristic = cfg[0], (cfg[1] or \"\")\n rows.append(\n BenchmarkRow(\n inst.name,\n planner,\n heuristic,\n False,\n False,\n None,\n None,\n 0,\n 0,\n 0,\n 0.0,\n f\"{type(exc).__name__}: {exc}\",\n )\n )\n continue\n for cfg in configs:\n planner = cfg[0]\n heuristic = cfg[1]\n kwargs = cfg[2] if len(cfg) > 2 else {}\n try:\n result = solve_task(\n task,\n planner,\n heuristic,\n max_expansions=max_expansions,\n time_limit=time_limit,\n **kwargs,\n )\n valid = bool(result.solved and validate_plan(task, result.plan))\n rows.append(\n BenchmarkRow(\n inst.name,\n planner,\n heuristic or \"\",\n result.solved,\n valid,\n result.cost,\n result.plan_length,\n result.stats.expanded,\n result.stats.generated,\n result.stats.evaluated,\n round(result.stats.runtime, 6),\n \"\",\n result.stats.truncated,\n )\n )\n except Exception as exc: # keep the benchmark going on a single failure\n rows.append(\n BenchmarkRow(\n inst.name,\n planner,\n heuristic or \"\",\n False,\n False,\n None,\n None,\n 0,\n 0,\n 0,\n 0.0,\n f\"{type(exc).__name__}: {exc}\",\n )\n )\n return rows\n\n\ndef to_csv(rows, path: str) -> None:\n fieldnames = (\n list(asdict(rows[0]).keys())\n if rows\n else [f.name for f in BenchmarkRow.__dataclass_fields__.values()]\n )\n with open(path, \"w\", newline=\"\", encoding=\"utf-8\") as handle:\n writer = csv.DictWriter(handle, fieldnames=fieldnames)\n writer.writeheader()\n for row in rows:\n writer.writerow(asdict(row))\n\n\ndef summarize(rows) -> dict:\n \"\"\"Aggregate coverage and totals per ``planner/heuristic`` configuration.\"\"\"\n summary: dict = {}\n for row in rows:\n key = f\"{row.planner}/{row.heuristic}\" if row.heuristic else row.planner\n agg = summary.setdefault(\n key, {\"coverage\": 0, \"expanded\": 0, \"runtime\": 0.0, \"instances\": 0}\n )\n agg[\"instances\"] += 1\n agg[\"coverage\"] += int(row.valid)\n agg[\"expanded\"] += row.expanded\n agg[\"runtime\"] += row.runtime\n return summary\n\n\ndef plot_summary(rows, path: str, metric: str = \"expanded\") -> None:\n \"\"\"Bar chart of a metric per configuration (requires the ``viz`` extra).\"\"\"\n import matplotlib\n\n matplotlib.use(\"Agg\")\n import matplotlib.pyplot as plt\n\n summary = summarize(rows)\n labels = list(summary.keys())\n values = [summary[k][metric] for k in labels]\n fig, ax = plt.subplots(figsize=(max(6, len(labels) * 0.9), 4))\n ax.bar(labels, values, color=\"#4C72B0\")\n ax.set_ylabel(f\"total {metric}\")\n ax.set_title(f\"Planner comparison ({metric})\")\n plt.xticks(rotation=45, ha=\"right\")\n plt.tight_layout()\n fig.savefig(path, dpi=120)\n plt.close(fig)\n", "jupyddl/cli.py": "\"\"\"Command-line interface.\n\n``solve``, ``benchmark``, ``animate``, ``demo``, ``requirements`` and\n``generate``. Every long-running command accepts ``--max-expansions`` and\n``--time-limit``; when a search stops on one of those it says so rather than\nreporting the instance unsolvable.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport argparse\nimport os\nimport sys\n\nfrom .api import build_task, solve_task, trace_search, validate_plan\nfrom .benchmark import (\n discover_instances,\n plot_summary,\n run_benchmark,\n summarize,\n to_csv,\n)\nfrom .heuristics import HEURISTICS\nfrom .search import INFORMED_PLANNERS, PLANNERS\n\nINFORMED = set(INFORMED_PLANNERS)\n\n\ndef _add_solve(sub):\n p = sub.add_parser(\"solve\", help=\"solve a single PDDL instance\")\n p.add_argument(\"domain\")\n p.add_argument(\"problem\")\n p.add_argument(\"-s\", \"--search\", default=\"astar\", choices=sorted(PLANNERS))\n p.add_argument(\n \"-H\", \"--heuristic\", default=\"lmcut\", choices=sorted(HEURISTICS) + [\"none\"]\n )\n p.add_argument(\n \"-w\", \"--weight\", type=float, default=2.0, help=\"weight for weighted A*\"\n )\n p.add_argument(\n \"--live\",\n action=\"store_true\",\n help=\"watch the search live in the terminal (no dependencies)\",\n )\n p.add_argument(\"--trace\", default=None, help=\"write the search trace as JSON\")\n p.add_argument(\n \"--plot\",\n default=None,\n help=\"write a four-panel search-progress chart (PNG; needs the viz extra)\",\n )\n p.add_argument(\n \"--tree\", default=None, help=\"write the radial search-wavefront chart (PNG)\"\n )\n p.add_argument(\n \"--plan-plot\", default=None, help=\"write the plan timeline chart (PNG)\"\n )\n p.add_argument(\"--dark\", action=\"store_true\", help=\"render charts for dark mode\")\n p.add_argument(\n \"--max-expansions\",\n type=int,\n default=None,\n help=\"stop after this many node expansions and report what was found\",\n )\n p.add_argument(\n \"--time-limit\",\n type=float,\n default=None,\n help=\"stop after this many seconds and report what was found\",\n )\n p.add_argument(\"--quiet\", action=\"store_true\", help=\"do not print the plan\")\n p.set_defaults(func=_cmd_solve)\n\n\ndef _add_benchmark(sub):\n p = sub.add_parser(\"benchmark\", help=\"compare planners over a folder of instances\")\n p.add_argument(\"root\", help=\"folder containing /domain.pddl + problem.pddl\")\n p.add_argument(\"--planners\", default=\"bfs,dijkstra,astar,gbfs,wastar,ehc\")\n p.add_argument(\"--heuristic\", default=\"hff\", help=\"heuristic for informed planners\")\n p.add_argument(\"--csv\", default=None, help=\"write per-run results to this CSV\")\n p.add_argument(\"--plot\", default=None, help=\"write a comparison bar chart (PNG)\")\n p.add_argument(\n \"--dashboard\",\n default=None,\n help=\"write the full benchmark dashboard: coverage, effort, time, heatmap\",\n )\n p.add_argument(\"--metric\", default=\"expanded\")\n p.add_argument(\"--dark\", action=\"store_true\", help=\"render charts for dark mode\")\n p.add_argument(\n \"--max-expansions\",\n type=int,\n default=None,\n help=\"stop after this many node expansions and report what was found\",\n )\n p.add_argument(\n \"--time-limit\",\n type=float,\n default=None,\n help=\"stop after this many seconds and report what was found\",\n )\n p.set_defaults(func=_cmd_benchmark)\n\n\ndef _add_animate(sub):\n p = sub.add_parser(\"animate\", help=\"replay a search as an animation (MP4 or GIF)\")\n p.add_argument(\"domain\")\n p.add_argument(\"problem\")\n p.add_argument(\"-o\", \"--output\", default=\"search.mp4\")\n p.add_argument(\"-s\", \"--search\", default=\"astar\", choices=sorted(PLANNERS))\n p.add_argument(\n \"-H\", \"--heuristic\", default=\"lmcut\", choices=sorted(HEURISTICS) + [\"none\"]\n )\n p.add_argument(\"--fps\", type=int, default=30)\n p.add_argument(\"--seconds\", type=float, default=8.0)\n p.add_argument(\"--dark\", action=\"store_true\")\n p.set_defaults(func=_cmd_animate)\n\n\ndef _add_demo(sub):\n p = sub.add_parser(\n \"demo\",\n help=\"run the bundled demo instances and write every chart to a folder\",\n )\n p.add_argument(\"-o\", \"--output\", default=\"gallery\", help=\"output folder\")\n p.add_argument(\"--root\", default=\"demos\", help=\"folder of demo instances\")\n p.add_argument(\n \"--both-modes\",\n action=\"store_true\",\n help=\"render a light and a dark variant of every chart\",\n )\n p.add_argument(\n \"--animate\", action=\"store_true\", help=\"also render the search animations\"\n )\n p.set_defaults(func=_cmd_demo)\n\n\n# --------------------------------------------------------------------------\ndef _observers(args):\n \"\"\"Build the observer for a solve, honouring --live/--trace/--plot.\"\"\"\n from .trace import MultiObserver, TraceRecorder\n\n wants_trace = bool(\n args.trace\n or args.plot\n or getattr(args, \"tree\", None)\n or getattr(args, \"plan_plot\", None)\n )\n recorder = TraceRecorder() if wants_trace else None\n dashboard = None\n if args.live:\n from .live import TerminalDashboard\n\n dashboard = TerminalDashboard()\n if recorder and dashboard:\n return MultiObserver(recorder, dashboard), recorder\n return (dashboard or recorder), recorder\n\n\ndef _cmd_solve(args) -> int:\n task = build_task(args.domain, args.problem)\n heuristic = None if args.heuristic == \"none\" else args.heuristic\n kwargs = {\"weight\": args.weight} if args.search == \"wastar\" else {}\n observer, recorder = _observers(args)\n result = solve_task(\n task,\n args.search,\n heuristic,\n observer=observer,\n max_expansions=args.max_expansions,\n time_limit=args.time_limit,\n **kwargs,\n )\n\n trace = recorder.trace if recorder else None\n if trace is not None:\n if args.trace:\n trace.save(args.trace)\n print(f\"Wrote trace to {args.trace}\")\n _write_charts(args, trace)\n\n if not result.solved:\n if not args.live:\n if result.truncated:\n print(\n \"No plan found within the budget \"\n \"(the instance may still be solvable).\"\n )\n else:\n print(\"No plan found.\")\n _print_stats(result)\n return 1\n valid = validate_plan(task, result.plan)\n visible = task.visible_plan(result.plan)\n if not args.quiet:\n header = f\"Plan ({len(visible)} steps, cost {result.cost}\"\n if task.temporal:\n header += f\", makespan {task.makespan(result.plan):g}\"\n print(header + \"):\")\n for i, op in enumerate(visible):\n print(f\" {i + 1:3d}. {op.base_name}\")\n print(f\"Valid: {valid}\")\n if not args.live:\n _print_stats(result)\n return 0 if valid else 2\n\n\ndef _write_charts(args, trace) -> None:\n targets = [\n (args.plot, \"plot_search_progress\"),\n (getattr(args, \"tree\", None), \"plot_search_tree\"),\n (getattr(args, \"plan_plot\", None), \"plot_plan_timeline\"),\n ]\n if not any(path for path, _ in targets):\n return\n try:\n from . import viz\n except ImportError as exc:\n print(f\"Charts need the viz extra: {exc}\", file=sys.stderr)\n return\n for path, function in targets:\n if not path:\n continue\n getattr(viz, function)(trace, path, dark=args.dark)\n print(f\"Wrote {path}\")\n\n\ndef _cmd_benchmark(args) -> int:\n instances = discover_instances(args.root)\n if not instances:\n print(f\"No instances found under {args.root}\", file=sys.stderr)\n return 1\n configs = []\n for planner in args.planners.split(\",\"):\n planner = planner.strip()\n configs.append((planner, args.heuristic if planner in INFORMED else None))\n\n rows = run_benchmark(\n instances,\n configs,\n max_expansions=args.max_expansions,\n time_limit=args.time_limit,\n )\n _print_summary(summarize(rows))\n if args.csv:\n to_csv(rows, args.csv)\n print(f\"\\nWrote per-run results to {args.csv}\")\n if args.plot:\n plot_summary(rows, args.plot, metric=args.metric)\n print(f\"Wrote plot to {args.plot}\")\n if args.dashboard:\n from .viz import plot_benchmark_dashboard\n\n plot_benchmark_dashboard(rows, args.dashboard, dark=args.dark)\n print(f\"Wrote dashboard to {args.dashboard}\")\n return 0\n\n\ndef _cmd_animate(args) -> int:\n from .viz import animate_search\n\n task = build_task(args.domain, args.problem)\n heuristic = None if args.heuristic == \"none\" else args.heuristic\n _, trace = trace_search(task, args.search, heuristic)\n animate_search(\n trace, args.output, dark=args.dark, fps=args.fps, seconds=args.seconds\n )\n print(f\"Wrote animation to {args.output}\")\n return 0\n\n\ndef _add_requirements(sub):\n p = sub.add_parser(\n \"requirements\",\n help=\"show which PDDL requirement flags are supported, and how\",\n )\n p.add_argument(\n \"--support\",\n default=None,\n help=\"filter by support level: native, compiled, partial or rejected\",\n )\n p.add_argument(\"--json\", action=\"store_true\", help=\"emit machine-readable JSON\")\n p.add_argument(\"--verbose\", action=\"store_true\", help=\"include the full notes\")\n p.set_defaults(func=_cmd_requirements)\n\n\ndef _add_generate(sub):\n from .generator import GENERATORS\n\n p = sub.add_parser(\n \"generate\", help=\"generate PDDL instances reproducibly from a seed\"\n )\n p.add_argument(\"kind\", choices=sorted(GENERATORS))\n p.add_argument(\"-o\", \"--output\", default=None, help=\"write into this folder\")\n p.add_argument(\"-n\", \"--size\", type=int, default=4, help=\"instance size\")\n p.add_argument(\"--seed\", type=int, default=0)\n p.add_argument(\n \"--count\",\n type=int,\n default=1,\n help=\"generate a ladder of instances with increasing size\",\n )\n p.add_argument(\n \"--step\", type=int, default=1, help=\"size increment between ladder rungs\"\n )\n p.set_defaults(func=_cmd_generate)\n\n\ndef _cmd_requirements(args) -> int:\n from .requirements import as_rows, summary\n\n rows = as_rows()\n if args.support:\n rows = [row for row in rows if row[\"support\"] == args.support]\n if not rows:\n print(\n f\"No requirements with support level '{args.support}'.\", file=sys.stderr\n )\n return 1\n\n if args.json:\n import json\n\n print(json.dumps({\"requirements\": rows, \"summary\": summary()}, indent=2))\n return 0\n\n counts = summary()\n print(\n f\"jupyddl PDDL support: {counts['native']} native, \"\n f\"{counts['compiled']} compiled, {counts['partial']} partial, \"\n f\"{counts['rejected']} rejected\\n\"\n )\n print(f\"{'requirement':<30}{'PDDL':<7}{'support':<11}summary\")\n print(\"-\" * 100)\n for row in rows:\n print(f\"{row['name']:<30}{row['pddl']:<7}{row['support']:<11}{row['summary']}\")\n if args.verbose and row[\"note\"]:\n for line in _wrap(row[\"note\"], 92):\n print(f\"{'':<48}{line}\")\n if not args.verbose:\n print(\"\\nRun with --verbose for the details of each compilation.\")\n return 0\n\n\ndef _wrap(text: str, width: int) -> list:\n words = text.split()\n lines, current = [], \"\"\n for word in words:\n if len(current) + len(word) + 1 > width:\n lines.append(current)\n current = word\n else:\n current = f\"{current} {word}\".strip()\n if current:\n lines.append(current)\n return lines\n\n\ndef _cmd_generate(args) -> int:\n from .generator import generate, write_instance\n\n sizes = [args.size + i * args.step for i in range(max(1, args.count))]\n if args.output is None:\n if len(sizes) > 1:\n print(\n \"Generating a ladder needs --output; a single instance can go \"\n \"to stdout but several cannot.\",\n file=sys.stderr,\n )\n return 1\n domain, problem = generate(args.kind, size=args.size, seed=args.seed)\n print(\";; ---------- domain.pddl ----------\")\n print(domain)\n print(\";; ---------- problem.pddl ----------\")\n print(problem)\n return 0\n\n for size in sizes:\n folder = write_instance(args.kind, args.output, size=size, seed=args.seed)\n print(f\"Wrote {folder}\")\n return 0\n\n\ndef _cmd_demo(args) -> int:\n \"\"\"Render the whole gallery: per-instance charts plus a benchmark dashboard.\"\"\"\n from .viz import (\n plot_benchmark_dashboard,\n plot_plan_timeline,\n plot_planner_comparison,\n plot_search_progress,\n plot_search_tree,\n )\n\n instances = discover_instances(args.root)\n if not instances:\n print(f\"No instances found under {args.root}\", file=sys.stderr)\n return 1\n os.makedirs(args.output, exist_ok=True)\n modes = [False, True] if args.both_modes else [False]\n\n configs = [(\"astar\", \"lmcut\"), (\"astar\", \"hmax\"), (\"gbfs\", \"hff\"), (\"bfs\", None)]\n for instance in instances:\n print(f\"-- {instance.name}\")\n try:\n task = build_task(instance.domain, instance.problem)\n except Exception as exc:\n print(f\" skipped: {type(exc).__name__}: {exc}\")\n continue\n traces = []\n for planner, heuristic in configs:\n try:\n _, trace = trace_search(task, planner, heuristic)\n traces.append(trace)\n except Exception as exc:\n print(f\" {planner}/{heuristic}: {type(exc).__name__}: {exc}\")\n if not traces:\n continue\n for dark in modes:\n suffix = \"-dark\" if dark else \"\"\n base = os.path.join(args.output, instance.name)\n plot_search_progress(traces[0], f\"{base}-progress{suffix}.png\", dark=dark)\n plot_search_tree(traces[0], f\"{base}-tree{suffix}.png\", dark=dark)\n plot_plan_timeline(traces[0], f\"{base}-plan{suffix}.png\", dark=dark)\n plot_planner_comparison(traces, f\"{base}-compare{suffix}.png\", dark=dark)\n if args.animate:\n from .viz import animate_search\n\n animate_search(\n traces[0], os.path.join(args.output, f\"{instance.name}-search.mp4\")\n )\n\n rows = run_benchmark(\n instances,\n [\n (\"astar\", \"lmcut\"),\n (\"astar\", \"hmax\"),\n (\"gbfs\", \"hff\"),\n (\"ehc\", \"hff\"),\n (\"bfs\", None),\n (\"dijkstra\", None),\n ],\n )\n to_csv(rows, os.path.join(args.output, \"benchmark.csv\"))\n for dark in modes:\n suffix = \"-dark\" if dark else \"\"\n plot_benchmark_dashboard(\n rows, os.path.join(args.output, f\"benchmark{suffix}.png\"), dark=dark\n )\n print(f\"\\nGallery written to {args.output}/\")\n return 0\n\n\ndef _print_stats(result) -> None:\n s = result.stats\n print(\n f\"Stats: expanded={s.expanded} generated={s.generated} \"\n f\"evaluated={s.evaluated} reopened={s.reopened} \"\n f\"deadends={s.deadends} runtime={s.runtime:.4f}s\"\n )\n\n\ndef _print_summary(summary) -> None:\n print(f\"{'config':<20}{'coverage':>12}{'expanded':>12}{'runtime(s)':>12}\")\n print(\"-\" * 56)\n for key, agg in summary.items():\n cov = f\"{agg['coverage']}/{agg['instances']}\"\n print(f\"{key:<20}{cov:>12}{agg['expanded']:>12}{agg['runtime']:>12.3f}\")\n\n\ndef main(argv=None) -> int:\n parser = argparse.ArgumentParser(\n prog=\"jupyddl\",\n description=\"Pure-Python PDDL planning framework.\",\n )\n sub = parser.add_subparsers(dest=\"command\", required=True)\n _add_solve(sub)\n _add_benchmark(sub)\n _add_animate(sub)\n _add_demo(sub)\n _add_requirements(sub)\n _add_generate(sub)\n args = parser.parse_args(argv)\n return args.func(args)\n\n\nif __name__ == \"__main__\": # pragma: no cover\n sys.exit(main())\n", "jupyddl/compile.py": "\"\"\"Compile PDDL 3 constructs down to the classical core.\n\nPreferences, trajectory constraints, timed initial literals and object fluents\nare all rewritten here, at the AST level, *before* grounding. The grounder and\nthe search never learn they existed \u2014 which is the point: one representation to\noptimise, and every new front-end feature is a source-to-source transformation\nrather than another special case in the hot loop.\n\nEach compilation is documented in :mod:`jupyddl.requirements`, including what it\ncosts you. The synthetic actions introduced along the way are all named with a\nleading ``__``; :attr:`jupyddl.task.Task.synthetic` collects them so a printed\nplan shows only the actions the domain author wrote.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom dataclasses import replace\n\nfrom .parser.ast import (\n Action,\n AddEffect,\n And,\n Atom,\n Comparison,\n ConjunctiveEffect,\n DelEffect,\n Domain,\n EqualityConstraint,\n Exists,\n Forall,\n ForallEffect,\n FluentRef,\n IncreaseCostEffect,\n Literal,\n Number,\n NumericEffect,\n Or,\n PDDLError,\n Predicate,\n Preference,\n Problem,\n Truth,\n UnsupportedFeatureError,\n WhenEffect,\n)\n\n# Every fact and action this module invents is prefixed, so they cannot collide\n# with a domain's own names and are easy to filter out of a plan.\nPREFIX = \"__\"\nCLOCK = FluentRef(\"__time\", ())\n\n__all__ = [\"compile_problem\", \"SYNTHETIC_PREFIX\"]\n\nSYNTHETIC_PREFIX = PREFIX\n\n\n# --------------------------------------------------------------------------\n# small AST helpers\n# --------------------------------------------------------------------------\ndef _fact(name: str, *args) -> Atom:\n return Atom(name, tuple(args))\n\n\ndef _holds(name: str, *args) -> Literal:\n return Literal(_fact(name, *args), True)\n\n\ndef _absent(name: str, *args) -> Literal:\n return Literal(_fact(name, *args), False)\n\n\ndef _conjoin(*parts):\n \"\"\"And() over the parts, dropping trivially-true ones.\"\"\"\n kept = [p for p in parts if p is not None and not isinstance(p, Truth)]\n if not kept:\n return Truth(True)\n if len(kept) == 1:\n return kept[0]\n return And(tuple(kept))\n\n\ndef _free():\n \"\"\"An explicit zero cost.\n\n Grounding charges 1 for any action without a cost effect, which is right for\n a domain action and wrong for the bookkeeping this module invents: closing\n the plan or observing that a constraint held must not show up in the metric.\n \"\"\"\n return IncreaseCostEffect(0.0)\n\n\ndef _also(effect, *extra):\n \"\"\"Append effects to an existing effect tree.\"\"\"\n parts = list(effect.parts) if isinstance(effect, ConjunctiveEffect) else [effect]\n parts.extend(extra)\n return ConjunctiveEffect(parts)\n\n\ndef _negate(formula):\n \"\"\"Negation normal form of ``not formula``.\"\"\"\n if isinstance(formula, Truth):\n return Truth(not formula.value)\n if isinstance(formula, Literal):\n return Literal(formula.atom, not formula.positive)\n if isinstance(formula, EqualityConstraint):\n return EqualityConstraint(formula.left, formula.right, not formula.positive)\n if isinstance(formula, Comparison):\n flip = {\"<\": \">=\", \"<=\": \">\", \">\": \"<=\", \">=\": \"<\", \"=\": \"!=\", \"!=\": \"=\"}\n return Comparison(flip[formula.op], formula.left, formula.right)\n if isinstance(formula, And):\n return Or(tuple(_negate(p) for p in formula.parts))\n if isinstance(formula, Or):\n return And(tuple(_negate(p) for p in formula.parts))\n if isinstance(formula, Forall):\n return Exists(formula.params, _negate(formula.body))\n if isinstance(formula, Exists):\n return Forall(formula.params, _negate(formula.body))\n raise PDDLError(f\"cannot negate {formula!r}\")\n\n\ndef _is_domain_action(action: Action) -> bool:\n return not action.name.startswith(PREFIX)\n\n\ndef _declare(domain: Domain, name: str, params=()) -> None:\n if all(p.name != name for p in domain.predicates):\n domain.predicates.append(Predicate(name, list(params)))\n\n\ndef _add_precondition(domain: Domain, formula, only_domain_actions=True) -> None:\n \"\"\"Conjoin ``formula`` onto every action's precondition.\"\"\"\n for index, action in enumerate(domain.actions):\n if only_domain_actions and not _is_domain_action(action):\n continue\n domain.actions[index] = replace(\n action, precondition=_conjoin(action.precondition, formula)\n )\n\n\ndef _add_conditional_effect(domain: Domain, condition, body) -> None:\n \"\"\"Give every domain action a ``when condition body`` effect.\n\n Used for monitors that must not be optional: the planner cannot decline to\n notice that a constraint's trigger became true.\n \"\"\"\n for index, action in enumerate(domain.actions):\n if not _is_domain_action(action):\n continue\n domain.actions[index] = replace(\n action, effect=_also(action.effect, WhenEffect(condition, body))\n )\n\n\n# --------------------------------------------------------------------------\n# object fluents\n# --------------------------------------------------------------------------\ndef compile_object_fluents(domain: Domain, problem: Problem):\n \"\"\"Turn ``(location ?p) - place`` into a predicate plus a uniqueness rule.\n\n ``(= (location ?p) ?x)`` becomes ``(__fn-location ?p ?x)``, and\n ``(assign (location ?p) ?y)`` clears the old value before setting the new\n one, so the predicate stays single-valued. Object fluents used as *nested\n terms* \u2014 ``(at ?t (location ?p))`` \u2014 are refused: flattening those needs a\n fresh existential per occurrence, and the equality form above expresses the\n same thing without the guesswork.\n \"\"\"\n if not domain.object_fluents:\n return domain, problem\n\n fluents = {fluent.name: fluent for fluent in domain.object_fluents}\n\n def predicate_name(name: str) -> str:\n return f\"{PREFIX}fn-{name}\"\n\n for fluent in domain.object_fluents:\n params = list(fluent.params) + [(\"?__value\", fluent.result_type)]\n _declare(domain, predicate_name(fluent.name), params)\n\n def rewrite_condition(formula):\n if isinstance(formula, Comparison) and formula.op in (\"=\", \"!=\"):\n left, right = formula.left, formula.right\n for a, b in ((left, right), (right, left)):\n if isinstance(a, FluentRef) and a.name in fluents:\n if isinstance(b, Number):\n raise UnsupportedFeatureError(\n f\"'{a.name}' returns an object, so it cannot be \"\n \"compared with a number\"\n )\n value = b.name if isinstance(b, FluentRef) else str(b)\n if isinstance(b, FluentRef) and b.name in fluents:\n raise UnsupportedFeatureError(\n \"comparing two object fluents directly is not \"\n \"supported; introduce a variable for one of them\"\n )\n atom = _fact(predicate_name(a.name), *a.args, value)\n return Literal(atom, formula.op == \"=\")\n return formula\n if isinstance(formula, And):\n return And(tuple(rewrite_condition(p) for p in formula.parts))\n if isinstance(formula, Or):\n return Or(tuple(rewrite_condition(p) for p in formula.parts))\n if isinstance(formula, Forall):\n return Forall(formula.params, rewrite_condition(formula.body))\n if isinstance(formula, Exists):\n return Exists(formula.params, rewrite_condition(formula.body))\n return formula\n\n def rewrite_effect(effect):\n if isinstance(effect, ConjunctiveEffect):\n return ConjunctiveEffect([rewrite_effect(p) for p in effect.parts])\n if isinstance(effect, ForallEffect):\n return ForallEffect(effect.params, rewrite_effect(effect.body))\n if isinstance(effect, WhenEffect):\n return WhenEffect(\n rewrite_condition(effect.condition), rewrite_effect(effect.body)\n )\n if isinstance(effect, NumericEffect) and effect.target.name in fluents:\n if effect.op != \"assign\":\n raise UnsupportedFeatureError(\n f\"'{effect.op}' is arithmetic, but '{effect.target.name}' \"\n \"returns an object; only 'assign' is meaningful\"\n )\n fluent = fluents[effect.target.name]\n name = predicate_name(fluent.name)\n value = effect.value\n new = value.name if isinstance(value, FluentRef) else str(value)\n # Clear whatever the function used to return, then set the new\n # value: that is what keeps it a function rather than a relation.\n clear = ForallEffect(\n [(\"?__old\", fluent.result_type)],\n WhenEffect(\n _holds(name, *effect.target.args, \"?__old\"),\n DelEffect(_fact(name, *effect.target.args, \"?__old\")),\n ),\n )\n return ConjunctiveEffect(\n [clear, AddEffect(_fact(name, *effect.target.args, new))]\n )\n return effect\n\n for index, action in enumerate(domain.actions):\n domain.actions[index] = replace(\n action,\n precondition=rewrite_condition(action.precondition),\n effect=rewrite_effect(action.effect),\n )\n for index, rule in enumerate(domain.derived):\n domain.derived[index] = replace(rule, body=rewrite_condition(rule.body))\n\n problem.goal = rewrite_condition(problem.goal)\n for index, preference in enumerate(problem.preferences):\n problem.preferences[index] = Preference(\n preference.name, rewrite_condition(preference.body)\n )\n\n # `(= (location p1) depot)` in :init becomes a plain fact.\n for fluent_ref, value in problem.init_objects.items():\n if fluent_ref.name not in fluents:\n raise PDDLError(\n f\"'{fluent_ref.name}' is assigned an object in :init but is not \"\n \"declared as an object fluent in :functions\"\n )\n problem.init.append(\n _fact(predicate_name(fluent_ref.name), *fluent_ref.args, value)\n )\n problem.init_objects = {}\n return domain, problem\n\n\n# --------------------------------------------------------------------------\n# timed initial literals\n# --------------------------------------------------------------------------\ndef compile_timed_initials(domain: Domain, problem: Problem):\n \"\"\"Give the model a clock, and make each timed literal fire off it.\n\n Elapsed time becomes the numeric fluent ``(__time)``, advanced by each\n durative action's duration. Every timed literal gets a zero-cost\n ``__fire-til-k`` action guarded by ``(>= (__time) t)``, and a\n ``__wait-til-k`` action that lets the planner advance the clock to ``t``\n when it wants the literal to happen.\n\n The literal must not be *skipped*: every domain action therefore carries\n ``(or (< (__time) t) (__til-k))`` for each timed literal, so once the clock\n passes ``t`` nothing else may happen until the literal has fired.\n\n Literals must also fire **in time order**, which is a separate constraint:\n ``(at 0 (open))`` and ``(at 3 (not (open)))`` describe a shop that opens then\n shuts, but firing them the other way round would leave it open forever. Each\n firing therefore requires every earlier literal to have fired already.\n\n Because actions do not overlap, a literal scheduled strictly inside an\n action's duration fires immediately after it rather than during it.\n \"\"\"\n if not problem.timed_initials:\n return domain, problem\n\n problem.init_numeric = dict(problem.init_numeric)\n problem.init_numeric.setdefault(CLOCK, 0.0)\n\n # The clock only moves if something moves it.\n for index, action in enumerate(domain.actions):\n if action.duration is None or not _is_domain_action(action):\n continue\n domain.actions[index] = replace(\n action,\n effect=_also(\n action.effect, NumericEffect(\"increase\", CLOCK, action.duration)\n ),\n )\n\n guards = []\n earlier: list = []\n for k, timed in enumerate(sorted(problem.timed_initials, key=lambda t: t.time)):\n marker = f\"{PREFIX}til-{k}\"\n _declare(domain, marker)\n due = Comparison(\">=\", CLOCK, Number(timed.time))\n not_due = Comparison(\"<\", CLOCK, Number(timed.time))\n # Everything scheduled before this must already have happened.\n in_order = _conjoin(*[_holds(name) for name in earlier])\n\n body = (\n AddEffect(timed.literal.atom)\n if timed.literal.positive\n else DelEffect(timed.literal.atom)\n )\n domain.actions.append(\n Action(\n f\"{PREFIX}fire-til-{k}\",\n [],\n _conjoin(due, _absent(marker), in_order),\n ConjunctiveEffect([body, AddEffect(_fact(marker)), _free()]),\n )\n )\n domain.actions.append(\n Action(\n f\"{PREFIX}wait-til-{k}\",\n [],\n # Waiting past an event that has not happened yet would skip it.\n _conjoin(not_due, in_order),\n ConjunctiveEffect(\n [NumericEffect(\"assign\", CLOCK, Number(timed.time)), _free()]\n ),\n )\n )\n guards.append(Or((not_due, _holds(marker))))\n earlier.append(marker)\n\n # A due literal blocks everything else until it has fired.\n for guard in guards:\n _add_precondition(domain, guard)\n problem.timed_initials = []\n return domain, problem\n\n\n# --------------------------------------------------------------------------\n# trajectory constraints\n# --------------------------------------------------------------------------\ndef compile_constraints(domain: Domain, problem: Problem):\n \"\"\"Compile ``(:constraints ...)`` into preconditions, monitors and goals.\n\n * ``always phi`` \u2014 conjoined onto every action's precondition and onto the\n goal. Every state on a plan's trajectory is either the initial state, a\n state an action is taken from, or the final state, so those three checks\n cover all of them.\n * ``at-end phi`` \u2014 conjoined onto the goal.\n * ``sometime phi`` \u2014 a zero-cost ``__observe`` action, applicable exactly\n when ``phi`` holds, sets a monitor fact the goal then requires.\n * ``sometime-before phi psi`` \u2014 the same monitor for ``psi``, plus\n ``always (phi implies monitor)``.\n * ``sometime-after phi psi`` \u2014 a *forced* monitor: every action records an\n outstanding obligation when ``phi`` holds without ``psi``, and discharges\n it when ``psi`` holds. The goal requires nothing outstanding.\n * ``at-most-once phi`` \u2014 forced monitors for \"phi has held\" and \"phi has\n since stopped\", plus ``always not (phi and stopped)``.\n\n Forced monitors ride on conditional effects, which the planner cannot\n decline; optional ones ride on actions, which it applies when convenient.\n The difference matters: a constraint the planner could satisfy by *not\n looking* would not be a constraint.\n \"\"\"\n constraints = list(domain.constraints) + list(problem.constraints)\n if not constraints:\n return domain, problem\n\n for constraint in constraints:\n if isinstance(constraint, Preference):\n raise UnsupportedFeatureError(\n f\"the soft constraint '{constraint.name}' is not supported: \"\n \"preferences are supported over goals, not over trajectory \"\n \"constraints\"\n )\n\n goal_parts = [problem.goal]\n invariants = []\n\n for index, constraint in enumerate(constraints):\n kind = constraint.kind\n if kind == \"always\":\n invariants.append(constraint.args[0])\n elif kind == \"at-end\":\n goal_parts.append(constraint.args[0])\n elif kind == \"sometime\":\n marker = _observation_monitor(domain, index, constraint.args[0])\n goal_parts.append(_holds(marker))\n elif kind == \"sometime-before\":\n trigger, earlier = constraint.args\n marker = _observation_monitor(domain, index, earlier)\n # \"phi implies the monitor\" as an invariant: by the time phi holds,\n # psi must already have been observed.\n invariants.append(Or((_negate(trigger), _holds(marker))))\n elif kind == \"sometime-after\":\n trigger, follower = constraint.args\n marker = f\"{PREFIX}pending-{index}\"\n _declare(domain, marker)\n _add_conditional_effect(\n domain,\n _conjoin(trigger, _negate(follower)),\n AddEffect(_fact(marker)),\n )\n _add_conditional_effect(domain, follower, DelEffect(_fact(marker)))\n goal_parts.append(_absent(marker))\n # ...and the final state must not leave a fresh obligation either.\n goal_parts.append(Or((_negate(trigger), follower)))\n elif kind == \"at-most-once\":\n phi = constraint.args[0]\n seen = f\"{PREFIX}amo-seen-{index}\"\n closed = f\"{PREFIX}amo-closed-{index}\"\n _declare(domain, seen)\n _declare(domain, closed)\n _add_conditional_effect(domain, phi, AddEffect(_fact(seen)))\n _add_conditional_effect(\n domain,\n _conjoin(_negate(phi), _holds(seen)),\n AddEffect(_fact(closed)),\n )\n # Holding again after an interval has closed is the second interval.\n invariants.append(Or((_negate(phi), _absent(closed))))\n else: # pragma: no cover - the parser rejects anything else\n raise UnsupportedFeatureError(f\"unsupported constraint '{kind}'\")\n\n for invariant in invariants:\n # Every action, not just the domain's own. The coverage argument above\n # holds only if each state on the trajectory is either taken from by an\n # action that checks the invariant or is the final state \u2014 and the\n # actions timed initial literals compile to *change facts*. Exempting\n # them let a plan step through a state the invariant forbade: a literal\n # that clears `(safe)` at t=3 and one that restores it at t=4 fire\n # back-to-back with nothing checking the state in between.\n _add_precondition(domain, invariant, only_domain_actions=False)\n goal_parts.append(invariant)\n\n problem.goal = _conjoin(*goal_parts)\n domain.constraints = []\n problem.constraints = []\n return domain, problem\n\n\ndef _observation_monitor(domain: Domain, index: int, formula) -> str:\n \"\"\"A fact the planner can set, for free, whenever ``formula`` holds.\"\"\"\n marker = f\"{PREFIX}seen-{index}\"\n _declare(domain, marker)\n domain.actions.append(\n Action(\n f\"{PREFIX}observe-{index}\",\n [],\n _conjoin(formula, _absent(marker)),\n ConjunctiveEffect([AddEffect(_fact(marker)), _free()]),\n )\n )\n return marker\n\n\n# --------------------------------------------------------------------------\n# preferences\n# --------------------------------------------------------------------------\ndef compile_preferences(domain: Domain, problem: Problem):\n \"\"\"Turn each soft goal into a priced choice: satisfy it, or pay for it.\n\n A closing phase makes this sound. ``__close`` ends the plan \u2014 every domain\n action requires that it has *not* happened \u2014 after which each preference is\n resolved by one of two zero-parameter actions: a free one that requires the\n preference to hold, or one costing the metric's ``(is-violated p)`` weight\n that does not. Cost-optimal search then picks whichever is cheaper, which is\n exactly what minimising the metric means.\n\n Freezing the state first is the point: without it the planner could satisfy\n a preference halfway through and then break it, and still be paid for it.\n \"\"\"\n if not problem.preferences:\n return domain, problem\n\n closed = f\"{PREFIX}closed\"\n _declare(domain, closed)\n _add_precondition(domain, _absent(closed))\n\n domain.actions.append(\n Action(\n f\"{PREFIX}close\",\n [],\n _absent(closed),\n ConjunctiveEffect([AddEffect(_fact(closed)), _free()]),\n )\n )\n\n goal_parts = [problem.goal, _holds(closed)]\n for index, preference in enumerate(problem.preferences):\n done = f\"{PREFIX}pref-{index}\"\n _declare(domain, done)\n weight = float(problem.violation_weights.get(preference.name, 1.0))\n if weight < 0:\n raise PDDLError(\n f\"preference '{preference.name}' has a negative violation \"\n \"weight, which would reward breaking it\"\n )\n domain.actions.append(\n Action(\n f\"{PREFIX}satisfy-{preference.name}\",\n [],\n _conjoin(_holds(closed), _absent(done), preference.body),\n ConjunctiveEffect([AddEffect(_fact(done)), _free()]),\n )\n )\n domain.actions.append(\n Action(\n f\"{PREFIX}violate-{preference.name}\",\n [],\n _conjoin(_holds(closed), _absent(done)),\n ConjunctiveEffect([AddEffect(_fact(done))]), # priced below\n )\n )\n # The penalty rides on the action cost, so plain cost-optimal search\n # optimises the metric without knowing what a preference is.\n domain.actions[-1] = replace(\n domain.actions[-1],\n effect=_also(domain.actions[-1].effect, _increase_cost(weight)),\n )\n goal_parts.append(_holds(done))\n\n problem.goal = _conjoin(*goal_parts)\n problem.preferences = []\n return domain, problem\n\n\ndef _increase_cost(amount: float):\n return IncreaseCostEffect(amount)\n\n\n# --------------------------------------------------------------------------\n# entry point\n# --------------------------------------------------------------------------\ndef compile_problem(domain: Domain, problem: Problem):\n \"\"\"Apply every PDDL 3 compilation, in the order they depend on each other.\n\n Object fluents go first because they rewrite terms everywhere else reads;\n preferences go last because their closing phase must sit outside the\n machinery the other compilations add.\n \"\"\"\n domain, problem = compile_object_fluents(domain, problem)\n domain, problem = compile_timed_initials(domain, problem)\n domain, problem = compile_constraints(domain, problem)\n domain, problem = compile_preferences(domain, problem)\n return domain, problem\n", "jupyddl/generator.py": "\"\"\"Generate PDDL domains and problems, reproducibly.\n\nBenchmarks live or die on their instance set, and hand-writing a scaling ladder\nis tedious. Each generator here takes a size and a ``seed`` and emits PDDL text,\nso a whole difficulty curve is one comprehension \u2014 and re-running the same seed\ngives byte-identical files, which is what makes a published experiment\nreproducible.\n\n::\n\n from jupyddl.generator import generate, GENERATORS\n\n domain, problem = generate(\"blocksworld\", size=10, seed=7)\n ladder = [generate(\"gripper\", size=n, seed=1) for n in range(2, 12)]\n\nEvery generator is exercised by the test suite, which grounds and solves what it\nproduces \u2014 a generator that emits unsolvable or unparseable PDDL is a bug.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport random\nfrom typing import Optional\n\n__all__ = [\n \"GENERATORS\",\n \"generate\",\n \"describe_generators\",\n \"generate_blocksworld\",\n \"generate_gripper\",\n \"generate_logistics\",\n \"generate_rovers\",\n \"generate_random_strips\",\n \"generate_numeric_transport\",\n \"generate_temporal_workshop\",\n]\n\n\ndef _objects(prefix: str, count: int) -> list:\n return [f\"{prefix}{i + 1}\" for i in range(count)]\n\n\n# --------------------------------------------------------------------------\n# blocksworld\n# --------------------------------------------------------------------------\nBLOCKSWORLD_DOMAIN = \"\"\"\\\n;; Blocksworld -- generated by jupyddl.generator\n(define (domain blocksworld)\n (:requirements :strips :typing)\n (:types block - object)\n (:predicates\n (on ?x - block ?y - block)\n (ontable ?x - block)\n (clear ?x - block)\n (handempty)\n (holding ?x - block))\n\n (:action pick-up\n :parameters (?x - block)\n :precondition (and (clear ?x) (ontable ?x) (handempty))\n :effect (and (not (ontable ?x)) (not (clear ?x))\n (not (handempty)) (holding ?x)))\n\n (:action put-down\n :parameters (?x - block)\n :precondition (holding ?x)\n :effect (and (not (holding ?x)) (clear ?x) (handempty) (ontable ?x)))\n\n (:action stack\n :parameters (?x - block ?y - block)\n :precondition (and (holding ?x) (clear ?y))\n :effect (and (not (holding ?x)) (not (clear ?y))\n (clear ?x) (handempty) (on ?x ?y)))\n\n (:action unstack\n :parameters (?x - block ?y - block)\n :precondition (and (on ?x ?y) (clear ?x) (handempty))\n :effect (and (holding ?x) (clear ?y) (not (clear ?x))\n (not (handempty)) (not (on ?x ?y)))))\n\"\"\"\n\n\ndef _random_towers(blocks, rng) -> list:\n \"\"\"Partition ``blocks`` into a random set of towers (bottom-first).\"\"\"\n shuffled = list(blocks)\n rng.shuffle(shuffled)\n towers: list = []\n index = 0\n while index < len(shuffled):\n height = rng.randint(1, max(1, len(shuffled) - index))\n towers.append(shuffled[index : index + height])\n index += height\n return towers\n\n\ndef _tower_facts(towers) -> list:\n facts = []\n for tower in towers:\n facts.append(f\"(ontable {tower[0]})\")\n for lower, upper in zip(tower, tower[1:]):\n facts.append(f\"(on {upper} {lower})\")\n facts.append(f\"(clear {tower[-1]})\")\n return facts\n\n\ndef generate_blocksworld(size: int = 6, seed: int = 0):\n \"\"\"``size`` blocks in random towers, to be rearranged into other towers.\"\"\"\n rng = random.Random(seed)\n blocks = _objects(\"b\", max(2, size))\n start = _random_towers(blocks, rng)\n goal_towers = _random_towers(blocks, rng)\n while goal_towers == start and len(blocks) > 2:\n goal_towers = _random_towers(blocks, rng)\n\n goal_facts = []\n for tower in goal_towers:\n goal_facts.append(f\"(ontable {tower[0]})\")\n for lower, upper in zip(tower, tower[1:]):\n goal_facts.append(f\"(on {upper} {lower})\")\n\n problem = f\"\"\"\\\n;; {len(blocks)} blocks, seed {seed} -- generated by jupyddl.generator\n(define (problem blocksworld-{len(blocks)}-{seed})\n (:domain blocksworld)\n (:objects {' '.join(blocks)} - block)\n (:init\n (handempty)\n {chr(10).join(' ' + f for f in _tower_facts(start)).strip()})\n (:goal (and\n {chr(10).join(' ' + f for f in goal_facts).strip()})))\n\"\"\"\n return BLOCKSWORLD_DOMAIN, problem\n\n\n# --------------------------------------------------------------------------\n# gripper\n# --------------------------------------------------------------------------\nGRIPPER_DOMAIN = \"\"\"\\\n;; Gripper -- generated by jupyddl.generator\n(define (domain gripper)\n (:requirements :strips :typing)\n (:types room ball gripper - object)\n (:predicates\n (at-robby ?r - room)\n (at ?b - ball ?r - room)\n (free ?g - gripper)\n (carry ?b - ball ?g - gripper))\n\n (:action move\n :parameters (?from - room ?to - room)\n :precondition (at-robby ?from)\n :effect (and (at-robby ?to) (not (at-robby ?from))))\n\n (:action pick\n :parameters (?b - ball ?r - room ?g - gripper)\n :precondition (and (at ?b ?r) (at-robby ?r) (free ?g))\n :effect (and (carry ?b ?g) (not (at ?b ?r)) (not (free ?g))))\n\n (:action drop\n :parameters (?b - ball ?r - room ?g - gripper)\n :precondition (and (carry ?b ?g) (at-robby ?r))\n :effect (and (at ?b ?r) (free ?g) (not (carry ?b ?g)))))\n\"\"\"\n\n\ndef generate_gripper(size: int = 4, seed: int = 0, grippers: int = 2):\n \"\"\"``size`` balls to move from room A to room B. Plan length grows linearly.\"\"\"\n balls = _objects(\"ball\", max(1, size))\n hands = _objects(\"gripper\", max(1, grippers))\n init = [\"(at-robby rooma)\"]\n init += [f\"(free {g})\" for g in hands]\n init += [f\"(at {b} rooma)\" for b in balls]\n\n problem = f\"\"\"\\\n;; {len(balls)} balls, {len(hands)} grippers, seed {seed}\n(define (problem gripper-{len(balls)}-{seed})\n (:domain gripper)\n (:objects\n rooma roomb - room\n {' '.join(balls)} - ball\n {' '.join(hands)} - gripper)\n (:init\n {chr(10).join(' ' + f for f in init).strip()})\n (:goal (and\n {chr(10).join(f' (at {b} roomb)' for b in balls).strip()})))\n\"\"\"\n return GRIPPER_DOMAIN, problem\n\n\n# --------------------------------------------------------------------------\n# logistics\n# --------------------------------------------------------------------------\nLOGISTICS_DOMAIN = \"\"\"\\\n;; Logistics with action costs -- generated by jupyddl.generator\n(define (domain logistics)\n (:requirements :strips :typing :action-costs)\n (:types\n truck airplane - vehicle\n package vehicle - thing\n airport location - place\n city - object)\n (:predicates\n (at ?t - thing ?p - place)\n (in ?p - package ?v - vehicle)\n (in-city ?p - place ?c - city))\n (:functions (total-cost))\n\n (:action load\n :parameters (?p - package ?v - vehicle ?l - place)\n :precondition (and (at ?p ?l) (at ?v ?l))\n :effect (and (not (at ?p ?l)) (in ?p ?v) (increase (total-cost) 1)))\n\n (:action unload\n :parameters (?p - package ?v - vehicle ?l - place)\n :precondition (and (in ?p ?v) (at ?v ?l))\n :effect (and (not (in ?p ?v)) (at ?p ?l) (increase (total-cost) 1)))\n\n (:action drive\n :parameters (?t - truck ?from - place ?to - place ?c - city)\n :precondition (and (at ?t ?from) (in-city ?from ?c) (in-city ?to ?c))\n :effect (and (not (at ?t ?from)) (at ?t ?to) (increase (total-cost) 2)))\n\n (:action fly\n :parameters (?a - airplane ?from - airport ?to - airport)\n :precondition (at ?a ?from)\n :effect (and (not (at ?a ?from)) (at ?a ?to) (increase (total-cost) 6))))\n\"\"\"\n\n\ndef generate_logistics(size: int = 3, seed: int = 0, cities: int = 2):\n \"\"\"``size`` packages across ``cities`` cities, each with an airport and depot.\"\"\"\n rng = random.Random(seed)\n cities = max(2, cities)\n packages = _objects(\"pkg\", max(1, size))\n city_names = _objects(\"city\", cities)\n airports = [f\"apt{i + 1}\" for i in range(cities)]\n depots = [f\"depot{i + 1}\" for i in range(cities)]\n trucks = [f\"truck{i + 1}\" for i in range(cities)]\n\n init = [\"(= (total-cost) 0)\"]\n for i in range(cities):\n init.append(f\"(in-city {airports[i]} {city_names[i]})\")\n init.append(f\"(in-city {depots[i]} {city_names[i]})\")\n init.append(f\"(at {trucks[i]} {depots[i]})\")\n init.append(f\"(at plane1 {airports[0]})\")\n\n goals = []\n for package in packages:\n source = rng.randrange(cities)\n target = rng.randrange(cities)\n while target == source and cities > 1:\n target = rng.randrange(cities)\n init.append(f\"(at {package} {rng.choice([airports[source], depots[source]])})\")\n goals.append(f\"(at {package} {rng.choice([airports[target], depots[target]])})\")\n\n problem = f\"\"\"\\\n;; {len(packages)} packages, {cities} cities, seed {seed}\n(define (problem logistics-{len(packages)}-{seed})\n (:domain logistics)\n (:objects\n {' '.join(city_names)} - city\n {' '.join(trucks)} - truck\n plane1 - airplane\n {' '.join(packages)} - package\n {' '.join(airports)} - airport\n {' '.join(depots)} - location)\n (:init\n {chr(10).join(' ' + f for f in init).strip()})\n (:goal (and\n {chr(10).join(' ' + g for g in goals).strip()}))\n (:metric minimize (total-cost)))\n\"\"\"\n return LOGISTICS_DOMAIN, problem\n\n\n# --------------------------------------------------------------------------\n# rovers (ADL: disjunction + quantification)\n# --------------------------------------------------------------------------\nROVERS_DOMAIN = \"\"\"\\\n;; Rovers -- exercises disjunctive and quantified preconditions.\n(define (domain rovers)\n (:requirements :strips :typing :adl)\n (:types rover waypoint objective - object)\n (:predicates\n (at ?r - rover ?w - waypoint)\n (can-traverse ?a ?b - waypoint)\n (visible ?o - objective ?w - waypoint)\n (imaged ?o - objective)\n (sampled ?w - waypoint)\n (analysed ?w - waypoint)\n (reported ?o - objective))\n\n (:action navigate\n :parameters (?r - rover ?from ?to - waypoint)\n :precondition (and (at ?r ?from) (can-traverse ?from ?to))\n :effect (and (not (at ?r ?from)) (at ?r ?to)))\n\n (:action sample\n :parameters (?r - rover ?w - waypoint)\n :precondition (and (at ?r ?w) (not (sampled ?w)))\n :effect (sampled ?w))\n\n (:action analyse\n :parameters (?w - waypoint)\n :precondition (sampled ?w)\n :effect (analysed ?w))\n\n ;; An objective can be imaged from any waypoint that sees it.\n (:action image\n :parameters (?r - rover ?o - objective)\n :precondition (exists (?w - waypoint) (and (at ?r ?w) (visible ?o ?w)))\n :effect (imaged ?o))\n\n ;; Reporting accepts either a picture or a full sample analysis.\n (:action report\n :parameters (?o - objective)\n :precondition (or (imaged ?o)\n (exists (?w - waypoint)\n (and (visible ?o ?w) (analysed ?w))))\n :effect (reported ?o)))\n\"\"\"\n\n\ndef generate_rovers(size: int = 3, seed: int = 0, waypoints: int = 5):\n \"\"\"``size`` objectives over a connected waypoint graph. Uses `or` and `exists`.\"\"\"\n rng = random.Random(seed)\n waypoints = max(2, waypoints)\n points = _objects(\"wp\", waypoints)\n objectives = _objects(\"obj\", max(1, size))\n\n init = [\"(at rover1 wp1)\"]\n # A spanning path guarantees the graph is connected, so every instance is\n # solvable; the extra edges just give the planner choices.\n for a, b in zip(points, points[1:]):\n init.append(f\"(can-traverse {a} {b})\")\n init.append(f\"(can-traverse {b} {a})\")\n for _ in range(waypoints):\n a, b = rng.sample(points, 2)\n init.append(f\"(can-traverse {a} {b})\")\n\n for objective in objectives:\n for point in rng.sample(points, rng.randint(1, min(2, len(points)))):\n init.append(f\"(visible {objective} {point})\")\n\n problem = f\"\"\"\\\n;; {len(objectives)} objectives, {waypoints} waypoints, seed {seed}\n(define (problem rovers-{len(objectives)}-{seed})\n (:domain rovers)\n (:objects\n rover1 - rover\n {' '.join(points)} - waypoint\n {' '.join(objectives)} - objective)\n (:init\n {chr(10).join(' ' + f for f in sorted(set(init))).strip()})\n (:goal (and\n {chr(10).join(f' (reported {o})' for o in objectives).strip()})))\n\"\"\"\n return ROVERS_DOMAIN, problem\n\n\n# --------------------------------------------------------------------------\n# numeric transport\n# --------------------------------------------------------------------------\nNUMERIC_DOMAIN = \"\"\"\\\n;; Transport with fuel -- exercises numeric fluents.\n(define (domain numeric-transport)\n (:requirements :strips :typing :numeric-fluents)\n (:types truck location package - object)\n (:predicates\n (at ?t - truck ?l - location)\n (road ?a ?b - location)\n (carrying ?p - package ?t - truck)\n (package-at ?p - package ?l - location))\n (:functions\n (fuel ?t - truck)\n (distance ?a ?b - location))\n\n (:action drive\n :parameters (?t - truck ?from ?to - location)\n :precondition (and (at ?t ?from) (road ?from ?to)\n (>= (fuel ?t) (distance ?from ?to)))\n :effect (and (not (at ?t ?from)) (at ?t ?to)\n (decrease (fuel ?t) (distance ?from ?to))))\n\n (:action refuel\n :parameters (?t - truck)\n :precondition (< (fuel ?t) 40)\n :effect (assign (fuel ?t) 60))\n\n (:action load\n :parameters (?p - package ?t - truck ?l - location)\n :precondition (and (at ?t ?l) (package-at ?p ?l))\n :effect (and (not (package-at ?p ?l)) (carrying ?p ?t)))\n\n (:action unload\n :parameters (?p - package ?t - truck ?l - location)\n :precondition (and (at ?t ?l) (carrying ?p ?t))\n :effect (and (not (carrying ?p ?t)) (package-at ?p ?l))))\n\"\"\"\n\n\ndef generate_numeric_transport(size: int = 2, seed: int = 0, locations: int = 4):\n \"\"\"``size`` packages on a fuel-limited road network.\"\"\"\n rng = random.Random(seed)\n locations = max(2, locations)\n places = _objects(\"loc\", locations)\n packages = _objects(\"pkg\", max(1, size))\n\n init = [\"(at truck1 loc1)\", \"(= (fuel truck1) 50)\"]\n for a, b in zip(places, places[1:]):\n distance = rng.choice([10, 15, 20])\n init.append(f\"(road {a} {b})\")\n init.append(f\"(road {b} {a})\")\n init.append(f\"(= (distance {a} {b}) {distance})\")\n init.append(f\"(= (distance {b} {a}) {distance})\")\n\n goals = []\n for package in packages:\n init.append(f\"(package-at {package} {places[0]})\")\n goals.append(f\"(package-at {package} {places[-1]})\")\n\n problem = f\"\"\"\\\n;; {len(packages)} packages, {locations} locations, seed {seed}\n(define (problem numeric-transport-{len(packages)}-{seed})\n (:domain numeric-transport)\n (:objects\n truck1 - truck\n {' '.join(places)} - location\n {' '.join(packages)} - package)\n (:init\n {chr(10).join(' ' + f for f in init).strip()})\n (:goal (and\n {chr(10).join(' ' + g for g in goals).strip()})))\n\"\"\"\n return NUMERIC_DOMAIN, problem\n\n\n# --------------------------------------------------------------------------\n# temporal workshop\n# --------------------------------------------------------------------------\nTEMPORAL_DOMAIN = \"\"\"\\\n;; Workshop -- exercises durative actions (sequential compilation).\n(define (domain workshop)\n (:requirements :strips :typing :durative-actions)\n (:types part - object)\n (:predicates\n (cut ?p - part)\n (drilled ?p - part)\n (painted ?p - part)\n (finished ?p - part))\n\n (:durative-action cut\n :parameters (?p - part)\n :duration (= ?duration 2)\n :condition (and (at start (not (cut ?p))))\n :effect (and (at end (cut ?p))))\n\n (:durative-action drill\n :parameters (?p - part)\n :duration (= ?duration 3)\n :condition (and (over all (cut ?p)))\n :effect (and (at end (drilled ?p))))\n\n (:durative-action paint\n :parameters (?p - part)\n :duration (= ?duration 5)\n :condition (and (over all (drilled ?p)))\n :effect (and (at end (painted ?p))))\n\n (:durative-action inspect\n :parameters (?p - part)\n :duration (= ?duration 1)\n :condition (and (over all (painted ?p)))\n :effect (and (at end (finished ?p)))))\n\"\"\"\n\n\ndef generate_temporal_workshop(size: int = 2, seed: int = 0):\n \"\"\"``size`` parts through a cut/drill/paint/inspect pipeline, with durations.\"\"\"\n parts = _objects(\"part\", max(1, size))\n problem = f\"\"\"\\\n;; {len(parts)} parts, seed {seed}\n(define (problem workshop-{len(parts)}-{seed})\n (:domain workshop)\n (:objects {' '.join(parts)} - part)\n (:init )\n (:goal (and\n {chr(10).join(f' (finished {p})' for p in parts).strip()})))\n\"\"\"\n return TEMPORAL_DOMAIN, problem\n\n\n# --------------------------------------------------------------------------\n# random STRIPS\n# --------------------------------------------------------------------------\ndef generate_random_strips(\n size: int = 8,\n seed: int = 0,\n actions: int = 10,\n goal_size: int = 3,\n):\n \"\"\"A random STRIPS instance built backwards from a guaranteed solution.\n\n Purely random operators almost always give an unsolvable problem, which\n makes for a useless benchmark. Instead this plants a random *chain*: each\n step in the chain has an action that turns the previous state into the next,\n so a plan of known length exists. The remaining actions are noise the\n planner has to search past.\n \"\"\"\n rng = random.Random(seed)\n size = max(4, size)\n facts = [f\"p{i + 1}\" for i in range(size)]\n chain_length = max(2, min(actions // 2, size))\n\n init_facts = set(rng.sample(facts, max(1, size // 3)))\n state = set(init_facts)\n operators = []\n\n # --- the planted solution ---------------------------------------------\n for step in range(chain_length):\n candidates = [f for f in facts if f not in state]\n if not candidates:\n break\n added = rng.choice(candidates)\n precondition = sorted(rng.sample(sorted(state), min(2, len(state))))\n deleted = []\n if len(state) > 1 and rng.random() < 0.4:\n deleted = [\n (\n rng.choice([f for f in sorted(state) if f not in precondition])\n if len(state) > len(precondition)\n else None\n )\n ]\n deleted = [d for d in deleted if d]\n operators.append((f\"solve{step + 1}\", precondition, [added], deleted))\n state.add(added)\n state.difference_update(deleted)\n\n goal_facts = sorted(rng.sample(sorted(state), min(goal_size, len(state))))\n\n # --- distractors -------------------------------------------------------\n for index in range(max(0, actions - len(operators))):\n precondition = sorted(rng.sample(facts, rng.randint(1, 2)))\n added = sorted(rng.sample(facts, rng.randint(1, 2)))\n deleted = [f for f in rng.sample(facts, 1) if f not in added]\n operators.append((f\"noise{index + 1}\", precondition, added, deleted))\n\n def action_text(name, pre, add, delete):\n conditions = \" \".join(f\"({p})\" for p in pre) or \"\"\n effects = \" \".join(f\"({a})\" for a in add)\n effects += \" \" + \" \".join(f\"(not ({d}))\" for d in delete)\n return (\n f\" (:action {name}\\n\"\n f\" :precondition (and {conditions})\\n\"\n f\" :effect (and {effects.strip()}))\"\n )\n\n domain = \"\\n\".join(\n [\n \";; Random STRIPS -- generated by jupyddl.generator\",\n f\"(define (domain random-strips-{size}-{seed})\",\n \" (:requirements :strips)\",\n \" (:predicates \" + \" \".join(f\"({f})\" for f in facts) + \")\",\n \"\",\n \"\\n\\n\".join(action_text(*op) for op in operators),\n \")\",\n ]\n )\n problem = f\"\"\"\\\n;; {size} facts, {len(operators)} actions, seed {seed}\n(define (problem random-strips-{size}-{seed})\n (:domain random-strips-{size}-{seed})\n (:init {' '.join(f'({f})' for f in sorted(init_facts))})\n (:goal (and {' '.join(f'({f})' for f in goal_facts)})))\n\"\"\"\n return domain, problem\n\n\n# --------------------------------------------------------------------------\n# registry\n# --------------------------------------------------------------------------\nGENERATORS = {\n \"blocksworld\": generate_blocksworld,\n \"gripper\": generate_gripper,\n \"logistics\": generate_logistics,\n \"rovers\": generate_rovers,\n \"numeric-transport\": generate_numeric_transport,\n \"workshop\": generate_temporal_workshop,\n \"random-strips\": generate_random_strips,\n}\n\n# What each generator is meant to stress, for the CLI and the playground.\nGENERATOR_NOTES = {\n \"blocksworld\": \"Classic block stacking. Deep plans, heavy plateaus.\",\n \"gripper\": \"Balls between two rooms. High branching factor.\",\n \"logistics\": \"Trucks and a plane with action costs.\",\n \"rovers\": \"ADL: disjunctive and existential preconditions.\",\n \"numeric-transport\": \"Numeric fluents: fuel consumption and refuelling.\",\n \"workshop\": \"Durative actions; plans report a makespan.\",\n \"random-strips\": \"Random operators around a planted solution chain.\",\n}\n\n\ndef generate(kind: str, size: int = 4, seed: int = 0, **kwargs):\n \"\"\"Generate ``(domain_text, problem_text)`` for a named generator.\"\"\"\n try:\n factory = GENERATORS[kind]\n except KeyError:\n raise ValueError(\n f\"Unknown generator '{kind}'. Available: {sorted(GENERATORS)}\"\n ) from None\n return factory(size=size, seed=seed, **kwargs)\n\n\ndef describe_generators() -> list:\n \"\"\"Serialisable metadata for the CLI and the web workbench.\"\"\"\n return [\n {\"name\": name, \"summary\": GENERATOR_NOTES.get(name, \"\")}\n for name in sorted(GENERATORS)\n ]\n\n\ndef write_instance(\n kind: str,\n folder: str,\n size: int = 4,\n seed: int = 0,\n name: Optional[str] = None,\n **kwargs,\n) -> str:\n \"\"\"Generate an instance and write ``//{domain,problem}.pddl``.\"\"\"\n import os\n\n domain, problem = generate(kind, size=size, seed=seed, **kwargs)\n target = os.path.join(folder, name or f\"{kind}-{size:02d}-{seed}\")\n os.makedirs(target, exist_ok=True)\n with open(os.path.join(target, \"domain.pddl\"), \"w\", encoding=\"utf-8\") as handle:\n handle.write(domain)\n with open(os.path.join(target, \"problem.pddl\"), \"w\", encoding=\"utf-8\") as handle:\n handle.write(problem)\n return target\n", "jupyddl/grounding.py": "\"\"\"Grounding: turn a parsed :class:`Domain` + :class:`Problem` into a\ngrounded :class:`~jupyddl.task.Task`.\n\nPipeline:\n\n1. Build the ``type -> objects`` table (with type-hierarchy closure).\n2. Instantiate every action over all type-consistent parameter tuples. The\n precondition formula is expanded (quantifiers over the object pool) and\n distributed into DNF; **each disjunct becomes its own grounded operator**, so\n the search only ever sees conjunctive preconditions.\n3. Detect static predicates (never added, deleted, or derived) and use the\n initial state to prune infeasible action instances and simplify conditions.\n4. Compile negative preconditions/goals into *positive normal form* by\n introducing complement facts ``(not ...)`` and maintaining them on every\n operator that touches the underlying atom.\n5. Ground derived-predicate rules into :class:`~jupyddl.task.Axiom` objects.\n6. Collect numeric fluents, compile their expressions into closures over the\n state's value vector, and encode everything as integer ids.\n\nA disjunctive goal is compiled to a single artificial goal fact achieved by one\nzero-cost operator per disjunct; those operators are recorded in\n``Task.synthetic`` so they can be hidden when a plan is printed.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom collections import defaultdict\nfrom dataclasses import dataclass, field\nfrom itertools import product\n\nfrom .parser.ast import (\n AddEffect,\n And,\n Arithmetic,\n Atom,\n Comparison,\n ConjunctiveEffect,\n DelEffect,\n Domain,\n EqualityConstraint,\n Exists,\n Forall,\n ForallEffect,\n FluentRef,\n IncreaseCostEffect,\n Literal,\n Number,\n NumericEffect,\n Or,\n PDDLError,\n Problem,\n Truth,\n UnsupportedFeatureError,\n WhenEffect,\n)\nfrom .compile import CLOCK as CLOCK_FLUENT\nfrom .compile import SYNTHETIC_PREFIX, compile_problem\nfrom .parser.parser import parse_domain_file, parse_problem_file\nfrom .task import Axiom, CondEffect, Operator, Task\n\n# Distributing a deeply disjunctive precondition into DNF can blow up\n# combinatorially. Stop with a clear message rather than exhausting memory.\nMAX_DISJUNCTS = 20000\n\nGOAL_FACT = Atom(\"__goal-reached__\", ())\n\n\ndef _ground_atom(atom: Atom, subst: dict) -> Atom:\n return Atom(atom.predicate, tuple(subst.get(a, a) for a in atom.args))\n\n\ndef _ground_fluent(ref: FluentRef, subst: dict) -> FluentRef:\n return FluentRef(ref.name, tuple(subst.get(a, a) for a in ref.args))\n\n\ndef _ground_expression(expr, subst: dict):\n if isinstance(expr, Number):\n return expr\n if isinstance(expr, FluentRef):\n return _ground_fluent(expr, subst)\n if isinstance(expr, Arithmetic):\n return Arithmetic(\n expr.op,\n _ground_expression(expr.left, subst),\n _ground_expression(expr.right, subst),\n )\n raise PDDLError(f\"unexpected numeric expression node: {expr!r}\")\n\n\ndef _build_type_objects(domain: Domain, problem: Problem):\n \"\"\"Map every type to the set of objects that inhabit it (hierarchy-aware).\"\"\"\n parent = dict(domain.types)\n\n def ancestors(typ: str):\n chain = [typ]\n seen = {typ}\n cur = typ\n while cur in parent and parent[cur] not in (None, \"object\", cur):\n cur = parent[cur]\n if cur in seen:\n break\n seen.add(cur)\n chain.append(cur)\n return chain\n\n type_objects: dict = defaultdict(set)\n declared: set = set()\n for name, typ in list(domain.constants) + list(problem.objects):\n declared.add(name)\n for anc in ancestors(typ):\n type_objects[anc].add(name)\n type_objects[\"object\"].add(name)\n\n # Robustness: some (toy) problems omit the :objects section and only mention\n # constants in :init/:goal. Treat any such undeclared constant as an object\n # of the root type so untyped domains still ground.\n for name in _harvest_constants(problem):\n if name not in declared:\n type_objects[\"object\"].add(name)\n\n for typ in set(list(parent) + list(parent.values())):\n type_objects.setdefault(typ, set())\n return {typ: tuple(sorted(objs)) for typ, objs in type_objects.items()}\n\n\ndef _harvest_constants(problem: Problem) -> set:\n found: set = set()\n for atom in problem.init:\n found.update(arg for arg in atom.args if not arg.startswith(\"?\"))\n for ref in problem.init_numeric:\n found.update(arg for arg in ref.args if not arg.startswith(\"?\"))\n\n def walk(formula):\n if isinstance(formula, Literal):\n found.update(a for a in formula.atom.args if not a.startswith(\"?\"))\n elif isinstance(formula, (And, Or)):\n for part in formula.parts:\n walk(part)\n elif isinstance(formula, (Exists, Forall)):\n walk(formula.body)\n\n walk(problem.goal)\n return found\n\n\n# --------------------------------------------------------------------------\n# condition formulas -> DNF\n# --------------------------------------------------------------------------\n@dataclass\nclass _Disjunct:\n \"\"\"One ground conjunction: positive atoms, negative atoms, comparisons.\"\"\"\n\n pos: set = field(default_factory=set)\n neg: set = field(default_factory=set)\n comparisons: list = field(default_factory=list)\n\n def merged(self, other: \"_Disjunct\") -> \"_Disjunct\":\n return _Disjunct(\n self.pos | other.pos,\n self.neg | other.neg,\n self.comparisons + other.comparisons,\n )\n\n @property\n def contradictory(self) -> bool:\n return bool(self.pos & self.neg)\n\n\nTRUE_DNF = [_Disjunct()]\nFALSE_DNF: list = []\n\n\ndef _dnf(formula, subst: dict, type_objects: dict) -> list:\n \"\"\"Expand quantifiers and distribute ``formula`` into a list of disjuncts.\n\n An empty list means *unsatisfiable*; a list holding one empty disjunct means\n *trivially true*.\n \"\"\"\n if formula is None:\n return TRUE_DNF\n\n if isinstance(formula, Truth):\n return TRUE_DNF if formula.value else FALSE_DNF\n\n if isinstance(formula, Literal):\n atom = _ground_atom(formula.atom, subst)\n if formula.positive:\n return [_Disjunct(pos={atom})]\n return [_Disjunct(neg={atom})]\n\n if isinstance(formula, EqualityConstraint):\n left = subst.get(formula.left, formula.left)\n right = subst.get(formula.right, formula.right)\n holds = (left == right) == formula.positive\n return TRUE_DNF if holds else FALSE_DNF\n\n if isinstance(formula, Comparison):\n grounded = Comparison(\n formula.op,\n _ground_expression(formula.left, subst),\n _ground_expression(formula.right, subst),\n )\n return [_Disjunct(comparisons=[grounded])]\n\n if isinstance(formula, And):\n result = TRUE_DNF\n for part in formula.parts:\n result = _cross(result, _dnf(part, subst, type_objects))\n if not result:\n return FALSE_DNF\n return result\n\n if isinstance(formula, Or):\n out: list = []\n for part in formula.parts:\n out.extend(_dnf(part, subst, type_objects))\n if len(out) > MAX_DISJUNCTS:\n raise UnsupportedFeatureError(\n \"disjunctive condition expands past \"\n f\"{MAX_DISJUNCTS} cases; simplify the domain or split the action\"\n )\n return out\n\n if isinstance(formula, Forall):\n result = TRUE_DNF\n for sub in _quantifier_substitutions(formula.params, subst, type_objects):\n result = _cross(result, _dnf(formula.body, sub, type_objects))\n if not result:\n return FALSE_DNF\n return result\n\n if isinstance(formula, Exists):\n out = []\n for sub in _quantifier_substitutions(formula.params, subst, type_objects):\n out.extend(_dnf(formula.body, sub, type_objects))\n if len(out) > MAX_DISJUNCTS:\n raise UnsupportedFeatureError(\n \"existential condition expands past \"\n f\"{MAX_DISJUNCTS} cases; the object pool is too large\"\n )\n return out\n\n raise PDDLError(f\"unexpected condition node: {formula!r}\")\n\n\ndef _quantifier_substitutions(params, subst: dict, type_objects: dict):\n pools = [type_objects.get(typ, ()) for (_, typ) in params]\n for combo in product(*pools):\n extended = dict(subst)\n for (var, _), obj in zip(params, combo):\n extended[var] = obj\n yield extended\n\n\ndef _cross(left: list, right: list) -> list:\n \"\"\"Distribute a conjunction of two DNFs, dropping contradictory disjuncts.\"\"\"\n if not left or not right:\n return FALSE_DNF\n out = []\n for a in left:\n for b in right:\n merged = a.merged(b)\n if not merged.contradictory:\n out.append(merged)\n if len(out) > MAX_DISJUNCTS:\n raise UnsupportedFeatureError(\n f\"conjunction of disjunctions expands past {MAX_DISJUNCTS} cases\"\n )\n return out\n\n\n# --------------------------------------------------------------------------\n# effects\n# --------------------------------------------------------------------------\n@dataclass\nclass _RawOp:\n name: str\n pre_pos: set\n pre_neg: set\n comparisons: list\n add: set\n delete: set\n cond: list # (cpos, cneg, cadd, cdel)\n numeric: list # (op, FluentRef, expression)\n cost: float\n duration: float\n synthetic: bool = False\n\n\ndef _collect_effect(eff, subst, type_objects, cpos, cneg, acc):\n if isinstance(eff, ConjunctiveEffect):\n for part in eff.parts:\n _collect_effect(part, subst, type_objects, cpos, cneg, acc)\n elif isinstance(eff, AddEffect):\n atom = _ground_atom(eff.atom, subst)\n if cpos or cneg:\n acc[\"cond\"].append((frozenset(cpos), frozenset(cneg), {atom}, set()))\n else:\n acc[\"add\"].add(atom)\n elif isinstance(eff, DelEffect):\n atom = _ground_atom(eff.atom, subst)\n if cpos or cneg:\n acc[\"cond\"].append((frozenset(cpos), frozenset(cneg), set(), {atom}))\n else:\n acc[\"delete\"].add(atom)\n elif isinstance(eff, IncreaseCostEffect):\n acc[\"cost\"] += eff.amount\n acc[\"has_cost\"] = True\n elif isinstance(eff, NumericEffect):\n if cpos or cneg:\n raise UnsupportedFeatureError(\n \"numeric effects inside a 'when' are not supported\"\n )\n acc[\"numeric\"].append(\n (\n eff.op,\n _ground_fluent(eff.target, subst),\n _ground_expression(eff.value, subst),\n )\n )\n elif isinstance(eff, ForallEffect):\n for sub in _quantifier_substitutions(eff.params, subst, type_objects):\n _collect_effect(eff.body, sub, type_objects, cpos, cneg, acc)\n elif isinstance(eff, WhenEffect):\n # A disjunctive effect condition splits into one conditional effect per\n # disjunct, which is exactly equivalent.\n for disjunct in _dnf(eff.condition, subst, type_objects):\n if disjunct.comparisons:\n raise UnsupportedFeatureError(\n \"numeric comparisons inside a 'when' condition are not supported\"\n )\n _collect_effect(\n eff.body,\n subst,\n type_objects,\n cpos | disjunct.pos,\n cneg | disjunct.neg,\n acc,\n )\n else:\n raise PDDLError(f\"unexpected effect node: {eff!r}\")\n\n\ndef _effect_predicates(domain: Domain) -> set:\n preds: set = set()\n\n def walk(eff):\n if isinstance(eff, ConjunctiveEffect):\n for part in eff.parts:\n walk(part)\n elif isinstance(eff, (AddEffect, DelEffect)):\n preds.add(eff.atom.predicate)\n elif isinstance(eff, (ForallEffect, WhenEffect)):\n walk(eff.body)\n\n for action in domain.actions:\n walk(action.effect)\n return preds\n\n\ndef _ground_raw_operators(domain, problem, type_objects) -> list:\n raw = []\n for action in domain.actions:\n pools = [type_objects.get(typ, ()) for (_, typ) in action.parameters]\n for combo in product(*pools):\n subst = {var: obj for (var, _), obj in zip(action.parameters, combo)}\n disjuncts = _dnf(action.precondition, subst, type_objects)\n if not disjuncts:\n continue # precondition is unsatisfiable for this instance\n\n acc = {\n \"add\": set(),\n \"delete\": set(),\n \"cond\": [],\n \"numeric\": [],\n \"cost\": 0.0,\n \"has_cost\": False,\n }\n _collect_effect(action.effect, subst, type_objects, set(), set(), acc)\n\n args = \",\".join(combo)\n base = f\"{action.name}({args})\" if combo else action.name\n duration = 0.0\n if action.duration is not None:\n duration = _constant_value(action.duration, subst, action.name)\n if acc[\"has_cost\"]:\n cost = acc[\"cost\"]\n elif duration:\n # A temporal action with no explicit cost: optimise makespan.\n cost = duration\n else:\n cost = 1\n\n for index, disjunct in enumerate(disjuncts):\n # Only tag the name when the split is real, so classical domains\n # keep the operator names their users expect.\n name = base if len(disjuncts) == 1 else f\"{base}#{index + 1}\"\n raw.append(\n _RawOp(\n name=name,\n pre_pos=set(disjunct.pos),\n pre_neg=set(disjunct.neg),\n comparisons=list(disjunct.comparisons),\n add=set(acc[\"add\"]),\n delete=set(acc[\"delete\"]),\n cond=list(acc[\"cond\"]),\n numeric=list(acc[\"numeric\"]),\n cost=cost,\n duration=duration,\n )\n )\n return raw\n\n\ndef _constant_value(expr, subst, action_name):\n \"\"\"Evaluate a duration expression that must be constant.\"\"\"\n grounded = _ground_expression(expr, subst)\n if isinstance(grounded, Number):\n return float(grounded.value)\n raise UnsupportedFeatureError(\n f\"the duration of '{action_name}' must be a constant; \"\n \"durations that read numeric fluents are not supported\"\n )\n\n\n# --------------------------------------------------------------------------\n# numeric compilation\n# --------------------------------------------------------------------------\ndef _collect_fluents(expr, out: set) -> None:\n if isinstance(expr, FluentRef):\n out.add(expr)\n elif isinstance(expr, Arithmetic):\n _collect_fluents(expr.left, out)\n _collect_fluents(expr.right, out)\n\n\ndef _compile_expression(expr, index_of: dict):\n \"\"\"Compile a ground numeric expression into a ``values -> float`` closure.\"\"\"\n if isinstance(expr, Number):\n constant = float(expr.value)\n return lambda values: constant\n if isinstance(expr, FluentRef):\n index = index_of[expr]\n return lambda values: values[index]\n if isinstance(expr, Arithmetic):\n left = _compile_expression(expr.left, index_of)\n right = _compile_expression(expr.right, index_of)\n if expr.op == \"+\":\n return lambda values: left(values) + right(values)\n if expr.op == \"-\":\n return lambda values: left(values) - right(values)\n if expr.op == \"*\":\n return lambda values: left(values) * right(values)\n if expr.op == \"/\":\n\n def divide(values):\n denominator = right(values)\n if denominator == 0:\n # Undefined rather than crashing mid-search: an infinite\n # value makes every comparison against it fail.\n return float(\"inf\")\n return left(values) / denominator\n\n return divide\n raise PDDLError(f\"unknown arithmetic operator '{expr.op}'\")\n raise PDDLError(f\"unexpected numeric expression: {expr!r}\")\n\n\ndef _compile_comparison(comparison: Comparison, index_of: dict):\n left = _compile_expression(comparison.left, index_of)\n right = _compile_expression(comparison.right, index_of)\n op = comparison.op\n if op == \"<\":\n return lambda values: left(values) < right(values)\n if op == \"<=\":\n return lambda values: left(values) <= right(values)\n if op == \">\":\n return lambda values: left(values) > right(values)\n if op == \">=\":\n return lambda values: left(values) >= right(values)\n if op == \"=\":\n return lambda values: left(values) == right(values)\n if op == \"!=\":\n return lambda values: left(values) != right(values)\n raise PDDLError(f\"unknown comparison operator '{op}'\")\n\n\ndef _compile_numeric_effect(op: str, index: int, value, index_of: dict):\n compute = _compile_expression(value, index_of)\n if op == \"assign\":\n return (index, compute)\n if op == \"increase\":\n return (index, lambda values: values[index] + compute(values))\n if op == \"decrease\":\n return (index, lambda values: values[index] - compute(values))\n if op == \"scale-up\":\n return (index, lambda values: values[index] * compute(values))\n if op == \"scale-down\":\n\n def scale_down(values):\n divisor = compute(values)\n return float(\"inf\") if divisor == 0 else values[index] / divisor\n\n return (index, scale_down)\n raise PDDLError(f\"unknown numeric assignment '{op}'\")\n\n\n# --------------------------------------------------------------------------\n# encoding\n# --------------------------------------------------------------------------\n@dataclass\nclass _Encoder:\n fact_ids: dict = field(default_factory=dict)\n comp_ids: dict = field(default_factory=dict)\n names: list = field(default_factory=list)\n\n def fact(self, atom: Atom) -> int:\n if atom not in self.fact_ids:\n self.fact_ids[atom] = len(self.names)\n self.names.append(str(atom))\n return self.fact_ids[atom]\n\n def comp(self, atom: Atom) -> int:\n if atom not in self.comp_ids:\n self.comp_ids[atom] = len(self.names)\n self.names.append(f\"(not {atom})\")\n return self.comp_ids[atom]\n\n\ndef _ground_axioms(domain: Domain, type_objects: dict) -> list:\n \"\"\"Ground every derived-predicate rule into (head atom, disjunct) pairs.\"\"\"\n grounded = []\n for rule in domain.derived:\n for subst in _quantifier_substitutions(rule.params, {}, type_objects):\n head = _ground_atom(rule.head, subst)\n for disjunct in _dnf(rule.body, subst, type_objects):\n if disjunct.comparisons:\n raise UnsupportedFeatureError(\n \"numeric comparisons in a derived predicate are not supported\"\n )\n grounded.append((head, disjunct))\n return grounded\n\n\ndef ground(domain: Domain, problem: Problem) -> Task:\n \"\"\"Ground ``domain`` + ``problem`` into a :class:`Task`.\"\"\"\n # PDDL 3 constructs (preferences, trajectory constraints, timed literals,\n # object fluents) are rewritten into the classical core first, so nothing\n # below this line has to know they exist.\n requirements = tuple(domain.requirements)\n domain, problem = compile_problem(domain, problem)\n type_objects = _build_type_objects(domain, problem)\n init_atoms = set(problem.init)\n\n derived_preds = {rule.head.predicate for rule in domain.derived}\n all_preds = {p.name for p in domain.predicates} | derived_preds\n # A derived predicate never appears in an effect, but it is emphatically not\n # static -- the axioms compute it.\n static_preds = all_preds - _effect_predicates(domain) - derived_preds\n\n def is_static(atom: Atom) -> bool:\n return atom.predicate in static_preds\n\n raw_ops = _ground_raw_operators(domain, problem, type_objects)\n axiom_rules = _ground_axioms(domain, type_objects)\n\n enc = _Encoder()\n tracked_neg: set = set()\n\n def resolve_literals(pos_atoms, neg_atoms):\n \"\"\"Drop static literals that hold, fail on ones that do not.\"\"\"\n pos_fluent, neg_fluent = set(), set()\n for atom in pos_atoms:\n if is_static(atom):\n if atom not in init_atoms:\n return None\n else:\n pos_fluent.add(atom)\n for atom in neg_atoms:\n if is_static(atom):\n if atom in init_atoms:\n return None\n else:\n neg_fluent.add(atom)\n tracked_neg.add(atom)\n return pos_fluent, neg_fluent\n\n # --- resolve static literals, drop infeasible operators ------------------\n resolved = []\n for op in raw_ops:\n pre = resolve_literals(op.pre_pos, op.pre_neg)\n if pre is None:\n continue\n pre_pos_fluent, pre_neg_fluent = pre\n\n add = set(op.add)\n delete = set(op.delete)\n cond = []\n for cpos, cneg, cadd, cdel in op.cond:\n trigger = resolve_literals(cpos, cneg)\n if trigger is None:\n continue # this conditional effect can never fire\n cpos_f, cneg_f = trigger\n if not cpos_f and not cneg_f:\n add |= cadd\n delete |= cdel\n else:\n cond.append((cpos_f, cneg_f, cadd, cdel))\n resolved.append((op, pre_pos_fluent, pre_neg_fluent, add, delete, cond))\n\n # --- goal ----------------------------------------------------------------\n goal_disjuncts = _dnf(problem.goal, {}, type_objects)\n if not goal_disjuncts:\n raise ValueError(\"Goal condition is self-contradictory\")\n\n resolved_goals = []\n unsolvable = True\n for disjunct in goal_disjuncts:\n parts = resolve_literals(disjunct.pos, disjunct.neg)\n if parts is None:\n continue # this way of satisfying the goal is statically impossible\n unsolvable = False\n resolved_goals.append((parts[0], parts[1], disjunct.comparisons))\n if not resolved_goals:\n resolved_goals = [(set(), set(), [])]\n\n # --- axioms --------------------------------------------------------------\n resolved_axioms = []\n for head, disjunct in axiom_rules:\n parts = resolve_literals(disjunct.pos, disjunct.neg)\n if parts is None:\n continue\n resolved_axioms.append((head, parts[0], parts[1]))\n\n # --- numeric fluents -----------------------------------------------------\n fluents: set = set(problem.init_numeric)\n for op, *_ in resolved:\n for comparison in op.comparisons:\n _collect_fluents(comparison.left, fluents)\n _collect_fluents(comparison.right, fluents)\n for _, target, value in op.numeric:\n fluents.add(target)\n _collect_fluents(value, fluents)\n for _, _, comparisons in resolved_goals:\n for comparison in comparisons:\n _collect_fluents(comparison.left, fluents)\n _collect_fluents(comparison.right, fluents)\n # `total-cost` is bookkeeping handled by operator costs, not a state variable.\n fluents = {ref for ref in fluents if ref.name != \"total-cost\"}\n\n ordered_fluents = sorted(fluents, key=str)\n index_of = {ref: i for i, ref in enumerate(ordered_fluents)}\n numeric_names = tuple(str(ref) for ref in ordered_fluents)\n init_values = tuple(\n float(problem.init_numeric.get(ref, 0.0)) for ref in ordered_fluents\n )\n\n # --- encode facts --------------------------------------------------------\n init_ids = {enc.fact(a) for a in init_atoms if not is_static(a)}\n for atom in tracked_neg:\n cid = enc.comp(atom)\n if atom not in init_atoms:\n init_ids.add(cid)\n\n def encode_add_del(add_atoms, del_atoms):\n add_ids = {enc.fact(a) for a in add_atoms}\n del_ids = {enc.fact(a) for a in del_atoms}\n for a in add_atoms:\n if a in tracked_neg:\n del_ids.add(enc.comp(a))\n for a in del_atoms:\n if a in tracked_neg:\n add_ids.add(enc.comp(a))\n return frozenset(add_ids), frozenset(del_ids)\n\n operators = []\n for op, pp, pn, add, delete, cond in resolved:\n precond = {enc.fact(a) for a in pp} | {enc.comp(a) for a in pn}\n add_ids, del_ids = encode_add_del(add, delete)\n cond_effects = []\n for cpos_f, cneg_f, cadd, cdel in cond:\n cond_ids = {enc.fact(a) for a in cpos_f} | {enc.comp(a) for a in cneg_f}\n cadd_ids, cdel_ids = encode_add_del(cadd, cdel)\n if cadd_ids or cdel_ids:\n cond_effects.append(CondEffect(frozenset(cond_ids), cadd_ids, cdel_ids))\n numeric_pre = tuple(_compile_comparison(c, index_of) for c in op.comparisons)\n numeric_eff = tuple(\n _compile_numeric_effect(kind, index_of[target], value, index_of)\n for kind, target, value in op.numeric\n )\n operators.append(\n Operator(\n op.name,\n frozenset(precond),\n add_ids,\n del_ids,\n tuple(cond_effects),\n op.cost,\n numeric_pre,\n numeric_eff,\n op.duration,\n )\n )\n\n # --- goal encoding, compiling a disjunction into a synthetic fact --------\n synthetic: set = set()\n goal_numeric: tuple = ()\n if len(resolved_goals) == 1:\n goal_pos, goal_neg, comparisons = resolved_goals[0]\n goals = {enc.fact(a) for a in goal_pos} | {enc.comp(a) for a in goal_neg}\n goal_numeric = tuple(_compile_comparison(c, index_of) for c in comparisons)\n else:\n goal_fact = enc.fact(GOAL_FACT)\n goals = {goal_fact}\n for index, (goal_pos, goal_neg, comparisons) in enumerate(resolved_goals):\n name = f\"__reach-goal__#{index + 1}\"\n synthetic.add(name)\n precond = {enc.fact(a) for a in goal_pos} | {enc.comp(a) for a in goal_neg}\n operators.append(\n Operator(\n name,\n frozenset(precond),\n frozenset({goal_fact}),\n frozenset(),\n (),\n 0,\n tuple(_compile_comparison(c, index_of) for c in comparisons),\n (),\n 0.0,\n )\n )\n\n if unsolvable:\n sentinel = len(enc.names)\n enc.names.append(\"(unsolvable)\")\n goals.add(sentinel) # never produced by any operator\n\n axioms = tuple(\n Axiom(\n enc.fact(head),\n frozenset(\n {enc.fact(a) for a in body_pos} | {enc.comp(a) for a in body_neg}\n ),\n )\n for head, body_pos, body_neg in resolved_axioms\n )\n\n # Anything a compilation introduced is bookkeeping, not something the\n # domain author wrote, so keep it out of printed plans.\n synthetic |= {\n op.name for op in operators if op.base_name.startswith(SYNTHETIC_PREFIX)\n }\n temporal = any(op.duration for op in operators)\n clock_index = index_of.get(CLOCK_FLUENT)\n metric = None\n if problem.metric is not None:\n direction, expression = problem.metric\n metric = f\"{direction} {expression}\"\n\n return Task(\n name=problem.name or domain.name,\n facts=tuple(enc.names),\n init=frozenset(init_ids),\n goals=frozenset(goals),\n operators=tuple(operators),\n metric_cost=problem.metric_minimize_cost,\n axioms=axioms,\n numeric_names=numeric_names,\n init_values=init_values,\n goal_numeric=goal_numeric,\n temporal=temporal,\n clock_index=clock_index,\n metric=metric,\n requirements=requirements,\n synthetic=frozenset(synthetic),\n )\n\n\ndef ground_files(domain_path: str, problem_path: str) -> Task:\n \"\"\"Convenience: parse both files and ground them into a :class:`Task`.\"\"\"\n return ground(parse_domain_file(domain_path), parse_problem_file(problem_path))\n", "jupyddl/heuristics/__init__.py": "\"\"\"Heuristics and a name-based registry.\"\"\"\n\nfrom __future__ import annotations\n\nfrom functools import partial\n\nfrom .base import Heuristic\nfrom .critical_path import CriticalPathHeuristic\nfrom .delete_relaxation import FFHeuristic, HAddHeuristic, HMaxHeuristic\nfrom .lmcut import LMCutHeuristic\nfrom .simple import BlindHeuristic, GoalCountHeuristic\n\n# name -> callable(task) -> Heuristic\nHEURISTICS = {\n \"blind\": BlindHeuristic,\n \"goalcount\": GoalCountHeuristic,\n \"hmax\": HMaxHeuristic,\n \"hadd\": HAddHeuristic,\n \"hff\": FFHeuristic,\n \"lmcut\": LMCutHeuristic,\n \"h1\": partial(CriticalPathHeuristic, m=1),\n \"h2\": partial(CriticalPathHeuristic, m=2),\n \"hm\": CriticalPathHeuristic,\n}\n\n\ndef make_heuristic(name: str, task) -> Heuristic:\n \"\"\"Instantiate a heuristic by name (see :data:`HEURISTICS`).\"\"\"\n try:\n factory = HEURISTICS[name]\n except KeyError:\n raise ValueError(\n f\"Unknown heuristic '{name}'. Available: {sorted(HEURISTICS)}\"\n ) from None\n return factory(task)\n\n\n__all__ = [\n \"Heuristic\",\n \"BlindHeuristic\",\n \"GoalCountHeuristic\",\n \"HMaxHeuristic\",\n \"HAddHeuristic\",\n \"FFHeuristic\",\n \"LMCutHeuristic\",\n \"CriticalPathHeuristic\",\n \"HEURISTICS\",\n \"make_heuristic\",\n]\n", "jupyddl/heuristics/base.py": "\"\"\"Heuristic base class.\"\"\"\n\nfrom __future__ import annotations\n\n\nclass Heuristic:\n \"\"\"A heuristic bound to a task; call it on a state to get an estimate.\n\n Returning ``math.inf`` signals that the state is a (relaxed) dead end.\n \"\"\"\n\n name: str = \"heuristic\"\n admissible: bool = False\n\n def __init__(self, task):\n self.task = task\n\n def __call__(self, state: frozenset) -> float: # pragma: no cover\n raise NotImplementedError\n", "jupyddl/heuristics/critical_path.py": "\"\"\"Critical-path heuristics h^m (Haslum & Geffner, 2000).\n\n``h^m`` estimates the cost of the most expensive size-``m`` subset of atoms.\n``h^1`` equals ``h_max``; higher ``m`` is more informative but costs more.\nAll ``h^m`` are admissible.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport math\nfrom itertools import combinations\n\nfrom .base import Heuristic\n\n\nclass CriticalPathHeuristic(Heuristic):\n name = \"hm\"\n admissible = True\n\n def __init__(self, task, m: int = 2):\n super().__init__(task)\n self.m = m\n self.goal = task.goals\n self.ops = list(task.relaxed_operators()) # (pre, add, cost)\n\n def __call__(self, state) -> float:\n return self._hm(state)\n\n def _hm(self, state) -> float:\n m = self.m\n inf = math.inf\n atoms = set(state) | set(self.goal)\n for pre, add, _ in self.ops:\n atoms |= pre\n atoms |= add\n atoms = sorted(atoms)\n\n # Table of costs for every atom-set of size 1..m.\n table: dict = {}\n sets: list = []\n for size in range(1, m + 1):\n for combo in combinations(atoms, size):\n fs = frozenset(combo)\n table[fs] = 0.0 if fs <= state else inf\n sets.append(fs)\n\n def cost_of(subset: frozenset) -> float:\n # Cost of an arbitrary atom-set: table lookup, or (if larger than m)\n # the max over its size-m subsets.\n if not subset:\n return 0.0\n if len(subset) <= m:\n return table.get(subset, inf)\n best = 0.0\n for combo in combinations(sorted(subset), m):\n val = table.get(frozenset(combo), inf)\n if val > best:\n best = val\n return best\n\n # Value iteration to the fixpoint.\n changed = True\n while changed:\n changed = False\n for target in sets:\n if table[target] == 0.0:\n continue\n best = table[target]\n for pre, add, cost in self.ops:\n if not (add & target):\n continue\n regressed = (target - add) | pre\n val = cost + cost_of(regressed)\n if val < best:\n best = val\n if best < table[target]:\n table[target] = best\n changed = True\n\n return cost_of(frozenset(self.goal))\n", "jupyddl/heuristics/delete_relaxation.py": "\"\"\"Delete-relaxation heuristics: h_max, h_add and the FF heuristic.\"\"\"\n\nfrom __future__ import annotations\n\nimport math\nfrom collections import deque\n\nfrom .base import Heuristic\nfrom .relaxation import RelaxedTask, goal_value, propagate_costs\n\n\nclass HMaxHeuristic(Heuristic):\n \"\"\"h_max: the most expensive relaxed goal fact. Admissible.\"\"\"\n\n name = \"hmax\"\n admissible = True\n\n def __init__(self, task):\n super().__init__(task)\n self.rt = RelaxedTask(task)\n\n def __call__(self, state) -> float:\n cost, _ = propagate_costs(self.rt, state, additive=False)\n return goal_value(cost, self.rt.goal, additive=False)\n\n\nclass HAddHeuristic(Heuristic):\n \"\"\"h_add: sum of relaxed goal-fact costs. Informative, not admissible.\"\"\"\n\n name = \"hadd\"\n\n def __init__(self, task):\n super().__init__(task)\n self.rt = RelaxedTask(task)\n\n def __call__(self, state) -> float:\n cost, _ = propagate_costs(self.rt, state, additive=True)\n return goal_value(cost, self.rt.goal, additive=True)\n\n\nclass FFHeuristic(Heuristic):\n \"\"\"FF heuristic: cost of a relaxed plan extracted from the h_add graph.\"\"\"\n\n name = \"hff\"\n\n def __init__(self, task):\n super().__init__(task)\n self.rt = RelaxedTask(task)\n\n def __call__(self, state) -> float:\n cost, supporter = propagate_costs(self.rt, state, additive=True)\n if any(math.isinf(cost[g]) for g in self.rt.goal):\n return math.inf\n relaxed_plan: set = set()\n seen: set = set()\n queue = deque(self.rt.goal)\n while queue:\n fact = queue.popleft()\n if fact in state or fact in seen:\n continue\n seen.add(fact)\n op_idx = supporter[fact]\n if op_idx < 0:\n return math.inf\n if op_idx not in relaxed_plan:\n relaxed_plan.add(op_idx)\n for pre_fact in self.rt.ops[op_idx].pre:\n if pre_fact not in state:\n queue.append(pre_fact)\n return float(sum(self.rt.ops[i].cost for i in relaxed_plan))\n", "jupyddl/heuristics/lmcut.py": "\"\"\"The LM-cut heuristic (Helmert & Domshlak, 2009).\n\nLM-cut repeatedly computes h_max, finds a *landmark cut* of operators separating\nthe initial state from the goal in the justification graph, adds the cut's\ncheapest operator cost to the estimate, and discounts the cut operators. It is\none of the strongest admissible heuristics for optimal planning.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport heapq\nimport math\nfrom dataclasses import dataclass\n\nfrom .base import Heuristic\n\n\n@dataclass\nclass _AugOp:\n idx: int\n pre: frozenset\n add: frozenset\n base_cost: int\n\n\nclass LMCutHeuristic(Heuristic):\n name = \"lmcut\"\n admissible = True\n\n def __init__(self, task):\n super().__init__(task)\n nf = task.num_facts\n self.INIT = nf # artificial \"always true\" fact\n self.GOAL = nf + 1 # artificial goal fact\n self.num = nf + 2\n\n self.ops: list[_AugOp] = []\n for pre, add, cost in task.relaxed_operators():\n # Every operator must have at least one precondition for the\n # justification graph; use the artificial INIT fact if needed.\n pre = pre if pre else frozenset({self.INIT})\n self.ops.append(_AugOp(len(self.ops), pre, add, cost))\n # Artificial goal operator (cost 0) collecting the real goal.\n self.goal_op = len(self.ops)\n self.ops.append(_AugOp(self.goal_op, task.goals, frozenset({self.GOAL}), 0))\n\n self.consumers: list[list[int]] = [[] for _ in range(self.num)]\n for op in self.ops:\n for f in op.pre:\n self.consumers[f].append(op.idx)\n\n def __call__(self, state) -> float:\n base = set(state)\n base.add(self.INIT)\n source = frozenset(base)\n costs = [op.base_cost for op in self.ops]\n total = 0.0\n\n while True:\n hmax, pcf = self._hmax(source, costs)\n if hmax[self.GOAL] == 0:\n return total\n if math.isinf(hmax[self.GOAL]):\n return math.inf\n\n goal_zone = self._goal_zone(costs, pcf)\n cut = self._cut(source, goal_zone, pcf, hmax)\n if not cut: # safety; should not happen when hmax(goal) > 0\n return math.inf\n min_cost = min(costs[oi] for oi in cut)\n total += min_cost\n for oi in cut:\n costs[oi] -= min_cost\n\n def _hmax(self, source, costs):\n inf = math.inf\n cost = [inf] * self.num\n pcf = [-1] * len(self.ops)\n counter = [len(op.pre) for op in self.ops]\n pq: list = []\n for f in source:\n cost[f] = 0\n heapq.heappush(pq, (0, f))\n\n def relax(op: _AugOp):\n supporter = max(op.pre, key=lambda p: cost[p])\n value = cost[supporter] + costs[op.idx]\n pcf[op.idx] = supporter\n if math.isinf(cost[supporter]):\n return\n for f in op.add:\n if value < cost[f]:\n cost[f] = value\n heapq.heappush(pq, (value, f))\n\n for op in self.ops:\n if counter[op.idx] == 0:\n relax(op)\n while pq:\n c, f = heapq.heappop(pq)\n if c > cost[f]:\n continue\n for op_idx in self.consumers[f]:\n counter[op_idx] -= 1\n if counter[op_idx] == 0:\n relax(self.ops[op_idx])\n return cost, pcf\n\n def _goal_zone(self, costs, pcf):\n \"\"\"Facts that reach the goal through zero-cost justification edges.\"\"\"\n zone = {self.GOAL}\n changed = True\n while changed:\n changed = False\n for op in self.ops:\n supporter = pcf[op.idx]\n if (\n costs[op.idx] == 0\n and supporter >= 0\n and supporter not in zone\n and (op.add & zone)\n ):\n zone.add(supporter)\n changed = True\n return zone\n\n def _cut(self, source, goal_zone, pcf, hmax):\n \"\"\"Operators crossing from the init-reachable region into the goal zone.\"\"\"\n # Forward-reachable facts from the initial state that stay out of the\n # goal zone (the \"before\" region).\n before = {f for f in source if f not in goal_zone}\n changed = True\n while changed:\n changed = False\n for op in self.ops:\n supporter = pcf[op.idx]\n if supporter in before and not math.isinf(hmax[supporter]):\n for f in op.add:\n if f not in goal_zone and f not in before:\n before.add(f)\n changed = True\n cut = [\n op.idx for op in self.ops if pcf[op.idx] in before and (op.add & goal_zone)\n ]\n return cut\n", "jupyddl/heuristics/relaxation.py": "\"\"\"Delete-relaxation machinery shared by h_max, h_add, h_FF and LM-cut.\"\"\"\n\nfrom __future__ import annotations\n\nimport heapq\nimport math\nfrom dataclasses import dataclass\n\n\n@dataclass\nclass RelaxedOp:\n idx: int\n pre: frozenset\n add: frozenset\n cost: int\n\n\nclass RelaxedTask:\n \"\"\"Delete-relaxed view of a task: unary-ish operators + goal.\n\n Conditional effects are already expanded into separate relaxed operators by\n :meth:`jupyddl.task.Task.relaxed_operators`.\n \"\"\"\n\n def __init__(self, task):\n self.num_facts = task.num_facts\n self.goal = task.goals\n self.ops = [\n RelaxedOp(i, pre, add, cost)\n for i, (pre, add, cost) in enumerate(task.relaxed_operators())\n ]\n # Index: fact -> operators that have it as a precondition.\n self.consumers: list[list[int]] = [[] for _ in range(self.num_facts)]\n self.no_pre: list[int] = []\n for op in self.ops:\n if op.pre:\n for f in op.pre:\n self.consumers[f].append(op.idx)\n else:\n self.no_pre.append(op.idx)\n\n\ndef propagate_costs(rt: RelaxedTask, state, additive: bool):\n \"\"\"Generalised Dijkstra computing h_max (``additive=False``) or h_add costs.\n\n Returns ``(cost, supporter)`` where ``cost[f]`` is the estimated cost to\n achieve fact ``f`` and ``supporter[f]`` is the operator index that achieved\n it (used for FF's relaxed-plan extraction).\n \"\"\"\n inf = math.inf\n cost = [inf] * rt.num_facts\n supporter = [-1] * rt.num_facts\n counter = [len(op.pre) for op in rt.ops]\n pq: list = []\n\n for f in state:\n if cost[f] != 0:\n cost[f] = 0\n heapq.heappush(pq, (0, f))\n\n def op_value(op: RelaxedOp) -> float:\n if not op.pre:\n return op.cost\n pre_costs = [cost[p] for p in op.pre]\n agg = sum(pre_costs) if additive else max(pre_costs)\n return op.cost + agg\n\n def apply_op(op: RelaxedOp):\n value = op_value(op)\n if math.isinf(value):\n return\n for f in op.add:\n if value < cost[f]:\n cost[f] = value\n supporter[f] = op.idx\n heapq.heappush(pq, (value, f))\n\n for idx in rt.no_pre:\n apply_op(rt.ops[idx])\n\n while pq:\n c, f = heapq.heappop(pq)\n if c > cost[f]:\n continue\n for op_idx in rt.consumers[f]:\n counter[op_idx] -= 1\n if counter[op_idx] == 0:\n apply_op(rt.ops[op_idx])\n return cost, supporter\n\n\ndef goal_value(cost, goal, additive: bool) -> float:\n if not goal:\n return 0.0\n values = [cost[g] for g in goal]\n if any(math.isinf(v) for v in values):\n return math.inf\n return float(sum(values) if additive else max(values))\n", "jupyddl/heuristics/simple.py": "\"\"\"Cheap non-relaxation heuristics.\"\"\"\n\nfrom __future__ import annotations\n\nfrom ..task import facts_of\nfrom .base import Heuristic\n\n\nclass BlindHeuristic(Heuristic):\n \"\"\"0 in a goal state, otherwise the cheapest operator cost. Admissible.\"\"\"\n\n name = \"blind\"\n admissible = True\n\n def __init__(self, task):\n super().__init__(task)\n costs = [op.cost for op in task.operators if op.cost > 0]\n self.min_cost = min(costs) if costs else 1\n\n def __call__(self, state) -> float:\n return 0.0 if self.task.goal_reached(state) else float(self.min_cost)\n\n\nclass GoalCountHeuristic(Heuristic):\n \"\"\"Number of unsatisfied goal facts. Fast, informative, not admissible.\"\"\"\n\n name = \"goalcount\"\n\n def __call__(self, state) -> float:\n return float(len(self.task.goals - facts_of(state)))\n", "jupyddl/live.py": "\"\"\"A live search dashboard that runs in any terminal, with no dependencies.\n\n:class:`TerminalDashboard` is a :class:`~jupyddl.trace.SearchObserver` that\nrepaints a small block of the terminal while the planner works \u2014 sparklines for\nthe heuristic and the ``f`` frontier, a frontier gauge, live counters and a node\nrate. It is pure standard library (ANSI escapes and Unicode block characters),\nso watching a search never costs you a dependency::\n\n from jupyddl import build_task, solve_task\n from jupyddl.live import TerminalDashboard\n\n task = build_task(\"domain.pddl\", \"problem.pddl\")\n solve_task(task, \"astar\", \"lmcut\", observer=TerminalDashboard())\n\nOn a non-interactive stream it degrades to a periodic one-line progress report,\nso it is also safe to use in CI logs and notebooks.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport math\nimport shutil\nimport sys\nimport time\n\nfrom .trace import SearchObserver\n\nSPARKS = \"\u2581\u2582\u2583\u2584\u2585\u2586\u2587\u2588\"\nGAUGE_FULL = \"\u2588\"\nGAUGE_EMPTY = \"\u2591\"\n\n# 256-colour ANSI approximations of the jupyddl palette slots.\n_ANSI = {\n \"blue\": \"\\x1b[38;5;33m\",\n \"orange\": \"\\x1b[38;5;208m\",\n \"aqua\": \"\\x1b[38;5;36m\",\n \"yellow\": \"\\x1b[38;5;178m\",\n \"green\": \"\\x1b[38;5;34m\",\n \"muted\": \"\\x1b[38;5;245m\",\n \"bold\": \"\\x1b[1m\",\n \"reset\": \"\\x1b[0m\",\n}\n\n__all__ = [\"TerminalDashboard\", \"sparkline\"]\n\n\ndef sparkline(values, width: int = 40) -> str:\n \"\"\"Render ``values`` as a Unicode sparkline of at most ``width`` cells.\"\"\"\n values = [v for v in values if v is not None and not math.isinf(v)]\n if not values:\n return \"\"\n if len(values) > width: # keep the shape, drop the resolution\n bucket = len(values) / width\n values = [values[min(len(values) - 1, int(i * bucket))] for i in range(width)]\n low, high = min(values), max(values)\n if high == low:\n return SPARKS[3] * len(values)\n span = high - low\n return \"\".join(SPARKS[min(7, int((v - low) / span * 7.999))] for v in values)\n\n\ndef _gauge(fraction: float, width: int = 16) -> str:\n fraction = min(1.0, max(0.0, fraction))\n filled = int(round(fraction * width))\n return GAUGE_FULL * filled + GAUGE_EMPTY * (width - filled)\n\n\ndef _fmt(value) -> str:\n if value is None:\n return \"\u2013\"\n value = int(value)\n if value >= 1_000_000:\n return f\"{value / 1_000_000:.1f}M\"\n if value >= 10_000:\n return f\"{value / 1000:.1f}k\"\n return f\"{value:,}\"\n\n\nclass TerminalDashboard(SearchObserver):\n \"\"\"Repaint a live view of the search in the terminal.\n\n ``interval`` throttles repaints (seconds); ``history`` bounds the sparkline\n buffers so a multi-million-node search still costs constant memory.\n \"\"\"\n\n def __init__(\n self,\n stream=None,\n interval: float = 0.08,\n history: int = 400,\n color: bool = True,\n ):\n self.stream = stream if stream is not None else sys.stderr\n self.interval = interval\n self.history = history\n self.interactive = bool(getattr(self.stream, \"isatty\", lambda: False)())\n self.color = color and self.interactive\n\n self._h: list = []\n self._f: list = []\n self._open: list = []\n self._peak_open = 1\n self._best_h = math.inf\n self._start = time.perf_counter()\n self._last_paint = 0.0\n self._lines = 0\n self._title = \"jupyddl\"\n self._latest = None\n self._bounds = 0\n\n # ------------------------------------------------------------- painting\n def _c(self, key: str, text: str) -> str:\n if not self.color:\n return text\n return f\"{_ANSI[key]}{text}{_ANSI['reset']}\"\n\n def _width(self) -> int:\n try:\n return max(48, min(shutil.get_terminal_size().columns - 2, 100))\n except Exception: # pragma: no cover - very unusual terminals\n return 72\n\n def _clear(self) -> None:\n if self._lines and self.interactive:\n self.stream.write(f\"\\x1b[{self._lines}A\\x1b[0J\")\n\n def _render(self, final: bool = False) -> None:\n width = self._width()\n spark_width = max(16, width - 30)\n event = self._latest\n expanded = event.expanded if event else 0\n generated = event.generated if event else 0\n evaluated = event.evaluated if event else 0\n elapsed = max(1e-9, time.perf_counter() - self._start)\n rate = expanded / elapsed\n\n current_h = self._h[-1] if self._h else None\n current_f = self._f[-1] if self._f else None\n open_now = self._open[-1] if self._open else 0\n best_h = None if math.isinf(self._best_h) else self._best_h\n\n lines = [\n self._c(\"bold\", self._title),\n self._c(\"muted\", \"\u2500\" * width),\n \" h \"\n + self._c(\"aqua\", sparkline(self._h, spark_width))\n + self._c(\"muted\", f\" now {_num(current_h)} best {_num(best_h)}\"),\n \" f \"\n + self._c(\"blue\", sparkline(self._f, spark_width))\n + self._c(\"muted\", f\" now {_num(current_f)}\"),\n \" frontier \"\n + self._c(\"orange\", _gauge(open_now / max(1, self._peak_open), 18))\n + self._c(\"muted\", f\" {_fmt(open_now)} (peak {_fmt(self._peak_open)})\"),\n \" \"\n + self._c(\"bold\", _fmt(expanded))\n + self._c(\"muted\", \" expanded \")\n + self._c(\"bold\", _fmt(generated))\n + self._c(\"muted\", \" generated \")\n + self._c(\"bold\", _fmt(evaluated))\n + self._c(\"muted\", \" evaluated\"),\n self._c(\"muted\", f\" {_fmt(rate)} nodes/s \u00b7 {elapsed:.2f}s\")\n + (self._c(\"muted\", f\" \u00b7 {self._bounds} bounds\") if self._bounds else \"\"),\n ]\n self._clear()\n self.stream.write(\"\\n\".join(lines) + \"\\n\")\n self.stream.flush()\n self._lines = len(lines)\n\n def _maybe_paint(self) -> None:\n now = time.perf_counter()\n if now - self._last_paint < self.interval:\n return\n self._last_paint = now\n if self.interactive:\n self._render()\n else:\n # Non-tty: a single appended line, much less often.\n if now - self._start > 1 and int(now) % 2 == 0:\n event = self._latest\n if event is not None:\n self.stream.write(\n f\" ... {_fmt(event.expanded)} expanded, \"\n f\"{_fmt(event.generated)} generated, \"\n f\"{now - self._start:.1f}s\\n\"\n )\n self.stream.flush()\n\n # -------------------------------------------------------- observer hooks\n def on_start(self, task, planner: str, heuristic: str = \"\") -> None:\n self._start = time.perf_counter()\n label = f\"{planner}/{heuristic}\" if heuristic else planner\n name = getattr(task, \"name\", \"\") or \"task\"\n facts = getattr(task, \"num_facts\", 0)\n operators = len(getattr(task, \"operators\", ()))\n self._title = (\n f\"jupyddl \u00b7 {label} \u00b7 {name} \"\n f\"({facts} facts, {operators} ground actions)\"\n )\n if self.interactive:\n self._render()\n\n def on_expand(\n self,\n state,\n g: float = 0.0,\n h: float = 0.0,\n f: float = 0.0,\n depth: int = 0,\n open_size: int = 0,\n stats=None,\n parent=None,\n action: str = \"\",\n ) -> None:\n self._latest = stats\n if h is not None and not math.isinf(h):\n self._h.append(h)\n self._best_h = min(self._best_h, h)\n self._f.append(f if f else g + (h or 0))\n self._open.append(open_size)\n self._peak_open = max(self._peak_open, open_size)\n for buffer in (self._h, self._f, self._open):\n if len(buffer) > self.history:\n del buffer[: len(buffer) - self.history]\n self._maybe_paint()\n\n def on_bound(self, threshold: float, iteration: int, stats=None) -> None:\n self._bounds += 1\n self._latest = stats or self._latest\n\n def on_finish(self, result) -> None:\n self._latest = getattr(result, \"stats\", None) or self._latest\n if self.interactive:\n self._render(final=True)\n elapsed = time.perf_counter() - self._start\n if result is not None and getattr(result, \"solved\", False):\n verdict = self._c(\n \"green\",\n f\" \u2714 solved \u00b7 cost {result.cost} \u00b7 \"\n f\"{result.plan_length} actions \u00b7 {elapsed:.2f}s\",\n )\n else:\n verdict = self._c(\"orange\", f\" \u2718 no plan found \u00b7 {elapsed:.2f}s\")\n self.stream.write(verdict + \"\\n\")\n self.stream.flush()\n self._lines = 0\n\n\ndef _num(value) -> str:\n if value is None:\n return \"\u2013\"\n if isinstance(value, float):\n if math.isinf(value):\n return \"\u221e\"\n if value.is_integer():\n return str(int(value))\n return f\"{value:.1f}\"\n return str(value)\n", "jupyddl/parser/__init__.py": "\"\"\"PDDL parsing: tokenizer, AST and recursive-descent parser.\"\"\"\n\nfrom .ast import (\n Action,\n AddEffect,\n And,\n Arithmetic,\n Atom,\n Comparison,\n Conjunct,\n ConjunctiveEffect,\n DelEffect,\n DerivedPredicate,\n Domain,\n EqualityConstraint,\n Exists,\n Forall,\n ForallEffect,\n FluentRef,\n Function,\n IncreaseCostEffect,\n Literal,\n Number,\n NumericEffect,\n Or,\n PDDLError,\n Predicate,\n Problem,\n Truth,\n UnsupportedFeatureError,\n WhenEffect,\n)\nfrom .parser import (\n parse,\n parse_condition,\n parse_domain,\n parse_domain_file,\n parse_effect,\n parse_problem,\n parse_problem_file,\n)\nfrom .tokenizer import tokenize\n\n__all__ = [\n \"Action\",\n \"AddEffect\",\n \"And\",\n \"Arithmetic\",\n \"Atom\",\n \"Comparison\",\n \"Conjunct\",\n \"ConjunctiveEffect\",\n \"DelEffect\",\n \"DerivedPredicate\",\n \"Domain\",\n \"EqualityConstraint\",\n \"Exists\",\n \"Forall\",\n \"ForallEffect\",\n \"FluentRef\",\n \"Function\",\n \"IncreaseCostEffect\",\n \"Literal\",\n \"Number\",\n \"NumericEffect\",\n \"Or\",\n \"PDDLError\",\n \"Predicate\",\n \"Problem\",\n \"Truth\",\n \"UnsupportedFeatureError\",\n \"WhenEffect\",\n \"parse\",\n \"parse_condition\",\n \"parse_domain\",\n \"parse_effect\",\n \"parse_problem\",\n \"parse_domain_file\",\n \"parse_problem_file\",\n \"tokenize\",\n]\n", "jupyddl/parser/ast.py": "\"\"\"Structured AST for PDDL.\n\nConditions are stored as a **formula tree in negation normal form**: the parser\npushes every ``not`` down to the atoms and rewrites ``imply``, so the grounder\nonly ever sees negation applied to a literal. Quantifiers stay in the tree\nbecause expanding them needs the object pool, which only exists at grounding\ntime; the grounder then distributes the formula into DNF and emits one operator\nper disjunct.\n\nNumeric fluents and durative actions have their own small node types. See\n:mod:`jupyddl.requirements` for exactly which PDDL requirement flags are\nsupported and how.\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom dataclasses import dataclass, field\n\n\nclass PDDLError(Exception):\n \"\"\"Base class for all parsing / modelling errors.\"\"\"\n\n\nclass UnsupportedFeatureError(PDDLError):\n \"\"\"Raised when a PDDL construct outside the supported subset is used.\"\"\"\n\n\n@dataclass(frozen=True)\nclass Atom:\n \"\"\"A (possibly lifted) predicate application, e.g. ``(on ?x ?y)``.\n\n ``args`` holds terms as raw strings: variables keep their leading ``?``\n while constants/objects are stored verbatim.\n \"\"\"\n\n predicate: str\n args: tuple = ()\n\n def __str__(self) -> str:\n if not self.args:\n return f\"({self.predicate})\"\n return f\"({self.predicate} {' '.join(self.args)})\"\n\n\n@dataclass(frozen=True)\nclass Literal:\n \"\"\"A positive or negative atom used in preconditions and goals.\"\"\"\n\n atom: Atom\n positive: bool = True\n\n\n@dataclass(frozen=True)\nclass EqualityConstraint:\n \"\"\"An ``(= a b)`` (or its negation) constraint over terms.\"\"\"\n\n left: str\n right: str\n positive: bool = True\n\n\n# --- numeric expressions -----------------------------------------------------\n\n\n@dataclass(frozen=True)\nclass Number:\n \"\"\"A numeric literal.\"\"\"\n\n value: float\n\n\n@dataclass(frozen=True)\nclass FluentRef:\n \"\"\"A reference to a numeric fluent, e.g. ``(fuel ?truck)``.\"\"\"\n\n name: str\n args: tuple = ()\n\n def __str__(self) -> str:\n if not self.args:\n return f\"({self.name})\"\n return f\"({self.name} {' '.join(self.args)})\"\n\n\n@dataclass(frozen=True)\nclass Arithmetic:\n \"\"\"A binary arithmetic expression (``+``, ``-``, ``*``, ``/``).\n\n Unary minus is parsed as ``(- 0 x)``.\n \"\"\"\n\n op: str\n left: object\n right: object\n\n\n@dataclass(frozen=True)\nclass Comparison:\n \"\"\"A numeric comparison used in preconditions and goals.\"\"\"\n\n op: str # one of < <= = >= >\n left: object\n right: object\n\n\n# --- condition formulas (negation normal form) -------------------------------\n\n\n@dataclass(frozen=True)\nclass Truth:\n \"\"\"The constant ``true`` \u2014 what an empty ``(and)`` parses to.\"\"\"\n\n value: bool = True\n\n\n@dataclass(frozen=True)\nclass And:\n parts: tuple = ()\n\n\n@dataclass(frozen=True)\nclass Or:\n parts: tuple = ()\n\n\n@dataclass(frozen=True)\nclass Exists:\n params: tuple = () # ((var, type), ...)\n body: object = None\n\n\n@dataclass(frozen=True)\nclass Forall:\n params: tuple = () # ((var, type), ...)\n body: object = None\n\n\n@dataclass\nclass Conjunct:\n \"\"\"One ground DNF disjunct: a conjunction of literals and comparisons.\n\n Produced by the grounder, not by the parser. ``equalities`` are resolved\n during instantiation and never survive into a grounded operator.\n \"\"\"\n\n literals: list = field(default_factory=list)\n equalities: list = field(default_factory=list)\n comparisons: list = field(default_factory=list)\n\n\n# --- effects -----------------------------------------------------------------\n\n\n@dataclass\nclass AddEffect:\n atom: Atom\n\n\n@dataclass\nclass DelEffect:\n atom: Atom\n\n\n@dataclass\nclass IncreaseCostEffect:\n \"\"\"``(increase (total-cost) k)`` \u2014 the classical action-cost shorthand.\"\"\"\n\n amount: float\n\n\n@dataclass\nclass NumericEffect:\n \"\"\"An assignment to a numeric fluent.\n\n ``op`` is one of ``assign``, ``increase``, ``decrease``, ``scale-up`` or\n ``scale-down``.\n \"\"\"\n\n op: str\n target: FluentRef\n value: object # an arithmetic expression\n\n\n@dataclass\nclass ConjunctiveEffect:\n parts: list = field(default_factory=list)\n\n\n@dataclass\nclass ForallEffect:\n params: list # [(variable, type), ...]\n body: object\n\n\n@dataclass\nclass WhenEffect:\n condition: object # a condition formula\n body: object\n\n\n# --- domain / problem --------------------------------------------------------\n\n\n@dataclass\nclass Predicate:\n name: str\n params: list # [(variable, type), ...]\n\n\n@dataclass\nclass Function:\n \"\"\"A declared numeric function (``:functions``).\"\"\"\n\n name: str\n params: list = field(default_factory=list)\n\n\n@dataclass\nclass Action:\n name: str\n parameters: list # [(variable, type), ...]\n precondition: object # a condition formula\n effect: object # one of the *Effect nodes above\n # Set when the action came from a (:durative-action ...) block; the value is\n # its duration, and the compilation is documented in jupyddl.requirements.\n duration: object = None\n\n\n@dataclass\nclass DerivedPredicate:\n \"\"\"A ``(:derived (head ?x) body)`` axiom.\"\"\"\n\n head: Atom\n params: list # [(variable, type), ...]\n body: object # a condition formula\n\n\n@dataclass\nclass Domain:\n name: str\n requirements: list\n types: dict # child type -> parent type (\"object\" if none)\n constants: list # [(name, type), ...]\n predicates: list\n actions: list\n functions: list = field(default_factory=list)\n derived: list = field(default_factory=list)\n object_fluents: list = field(default_factory=list)\n constraints: list = field(default_factory=list)\n\n @property\n def has_durative_actions(self) -> bool:\n return any(action.duration is not None for action in self.actions)\n\n\n@dataclass\nclass Problem:\n name: str\n domain_name: str\n objects: list # [(name, type), ...]\n init: list # [Atom, ...]\n goal: object # a condition formula\n metric_minimize_cost: bool = False\n init_numeric: dict = field(default_factory=dict) # FluentRef -> float\n metric: object = None # (direction, expression) or None\n preferences: list = field(default_factory=list) # [Preference, ...]\n constraints: list = field(default_factory=list) # [Constraint | Preference]\n timed_initials: list = field(default_factory=list) # [TimedInitial, ...]\n init_objects: dict = field(default_factory=dict) # object-fluent assignments\n violation_weights: dict = field(default_factory=dict) # preference -> weight\n\n\n# --- PDDL 3: preferences and trajectory constraints --------------------------\n\n\n@dataclass(frozen=True)\nclass Preference:\n \"\"\"A named soft goal or soft constraint.\n\n ``body`` is either a condition formula (a goal preference) or a\n :class:`Constraint` (a soft trajectory constraint). Violating it is legal;\n the ``:metric`` says what it costs.\n \"\"\"\n\n name: str\n body: object\n\n\n@dataclass(frozen=True)\nclass Constraint:\n \"\"\"One state-trajectory constraint from a ``(:constraints ...)`` block.\n\n ``kind`` is the PDDL modal operator (``always``, ``sometime``, ``at-end``,\n ``at-most-once``, ``sometime-before``, ``sometime-after``) and ``args`` holds\n its operand formulas, already in negation normal form.\n \"\"\"\n\n kind: str\n args: tuple = ()\n\n\n@dataclass(frozen=True)\nclass TimedInitial:\n \"\"\"A ``(at