Skip to content

ENH: Mixed discrete-continuous action spaces (ProductActions) #107

Description

@oyamad

Goal

Add ProductActions, the product of a discrete and a continuous action space, completing the action-space design of #104/#105 (cf. the request for mixed problems in #5). A solver for x = (x_d, x_c) where x_d ranges over a finite set and x_c over an M-dimensional box: the inner problem at each state is solved by enumeration over the discrete branches, with one continuous sub-maximization per branch reusing the existing scalar/multivariate inner solvers unchanged.

This issue is written as an implementation specification: follow it closely, including the MUST/MUST NOT items, which encode bugs actually encountered while building #93-#105. Read .github/copilot-instructions.md first; all conventions there (PR format, benchmark tables, test discipline) apply.

Type and user contract

struct ProductActions{TD<:DiscreteActions,TC<:ContinuousActions} <: ActionSpace
    discrete::TD
    continuous::TC
end
  • Actions are passed to the user's f(s, x) and g(s, x, e) as 2-tuples x = (x_d, x_c): x[1] is a discrete value from discrete.vals, x[2] is the continuous part (a Float64 for M == 1, a length-M indexable collection for M > 1).
  • _policy_eltype(::ProductActions{DiscreteActions{TA},ContinuousActions{1}}) = Tuple{TA,Float64}; for M > 1 the second element is NTuple{M,Float64}. res.X is a Vector of these tuples; res.X_ind holds the discrete-part indices (same convention as DiscreteActions).
  • Infeasibility convention unchanged: f returns -Inf.
  • LQA throws ArgumentError. The legacy s_wise_max!/compute_greedy! helpers are not extended (their docstrings already scope them to one-dimensional continuous spaces).

Core algorithm: branch decomposition via curried models

For each discrete index k, the continuous sub-problem is the existing problem with x_d fixed. Implement this by constructing, once per solve, K lightweight "curried" ContinuousDPs that close over the branch value and forward to the user's functions:

function _branch_cdp(cdp::ContinuousDP, k::Int)
    a = cdp.actions        # ::ProductActions
    dval = a.discrete.vals[k]
    f_k(s, xc) = cdp.f(s, (dval, xc))
    g_k(s, xc, e) = cdp.g(s, (dval, xc), e)
    # MUST: construct via the inner struct constructor, REUSING cdp.interp.
    # MUST NOT: call the basis-taking constructor, which re-runs Interp()
    # and re-factorizes Phi (O(n^3)) once per branch.
    return ContinuousDP(f_k, g_k, cdp.discount, cdp.shocks, cdp.weights,
                        a.continuous, cdp.interp)
end

The per-state inner solve is then:

# xc_warm[k] is the warm-start container row for branch k at this state
best_v, best_k = -Inf, 1
for k in 1:K
    v_k = <continuous inner solve on branch_cdps[k] at s, warm-started
           from branch k's own warm slot, writing its maximizer back>
    if v_k > best_v
        best_v, best_k = v_k, k
    end
end
# store best_v in Tv[i], best_k in X_ind[i]
  • For M == 1 the continuous solve is _s_wise_max_foc! (with ws.dfecs) or _s_wise_max! (Brent), selected exactly as in bellman_operator! today; for M > 1 it is _s_wise_max_multi! with use_foc = (inner_solver == :foc && dfecs !== nothing).
  • MUST: warm starts are per (state, branch) — every branch's sub-solve updates its own slot each sweep, including non-chosen branches (the sub-solve does this naturally when run). A single shared warm slot destroys warm-start validity when the chosen branch switches.
  • The fec and dfecs caches are state-side and branch-independent: share one set across all branches. MUST: set_coefs!(dfec, C) once per sweep before any funeval_point! on a DerivFunEvalCache — evaluating with unset coefficients reads uninitialized memory (this exact bug occurred in ENH: Add multi-dimensional continuous actions #105's development).

Workspace and containers

Extend CDPWorkspace construction for ProductActions:

  • fec as usual; dfecs built iff inner_solver == :foc and every basis dimension satisfies _foc_suitable (same rule as continuous spaces; the discrete part is irrelevant to it).
  • X_ind::Vector{Int} sized n (chosen branch per node).
  • Continuous warm-start storage for all branches: Vector{typeof(policy buffer)} of length K — for M == 1 a Vector{Vector{Float64}} (each fill(NaN, n)), for M > 1 a Vector{Matrix{Float64}} (each fill(NaN, n, M)). NaN is the cold-start marker, as elsewhere.
  • The curried branch_cdps vector (element type is concrete: the closures differ only in the captured value).

Downstream dispatch points (all of them)

Mirror the DiscreteActions branches added in #104 at every one of these sites; do not skip any:

  1. bellman_operator!(cdp, C, ws) and policy_iteration_operator!(cdp, C, ws): product sweep; the policy-iteration variant passes action values to evaluate_policy! as a lazily-built value vector [(vals[X_ind[i]], xc_i) ...] or an equivalent non-copying wrapper.
  2. The legacy bellman_operator!(cdp, C, Tv) overload: supported, allocating its own temporaries (its Tv-only interface is action-type-agnostic).
  3. evaluate!(res, fec) and the callable res(s_nodes): product enumeration at the evaluation nodes. MUST: respect res.inner_solver — build dfecs only when it is :foc (a :brent solve must never be re-evaluated through the derivative path; this is the ENH: Add multi-dimensional continuous actions #105 P0-1 review fix, do not regress it).
  4. simulate!: recompute the greedy action exactly at each visited state (as for DiscreteActions); never interpolate the discrete part.
  5. CDPSolveResult construction (_empty_policy/_policy_buffer/_policy_size_ok methods for ProductActions).
  6. The copy constructor: x_lb/x_ub keywords must throw ArgumentError for ProductActions (bounds live inside the continuous component; replacing them would require rebuilding the product explicitly).
  7. solve docstring: inner_solver applies to the continuous sub-solves; enumeration handles the discrete part.

Test plan (new file test/test_product_actions.jl, included from runtests.jl)

MUST: every test model keeps g(s, x, e) within the interpolation domain for all feasible actions and shocks (clamp in g or bound the continuous box by the domain) — out-of-domain next states make Chebyshev-basis solves silently diverge (see PR #97).

  1. Dominance test (exact): a growth model with a binary discrete part where branch 2's reward carries a -1e3 penalty, so branch 1 strictly dominates. The ProductActions solve must give all(res.X_ind .== 1) and V/continuous policy equal to the plain ContinuousActions solve of branch 1's model to rtol = 1e-10 (same fixed point, same inner solver).
  2. Genuine trade-off test (agreement with full discretization): a model where the discrete choice matters (e.g. two technologies with different productivity and a switching cost in f); compare against a DiscreteActions solve on the product of the discrete values and a 501-point grid of the continuous part — V within O(h) of each other (tolerance ~20h as in test_discrete_actions.jl).
  3. inner_solver = :brent respected end-to-end: res.inner_solver == :brent and res.X continuous parts equal a direct derivative-free recomputation at a few states (mirror the ENH: Add multi-dimensional continuous actions #105 regression test).
  4. Containers: res.X isa Vector{Tuple{...}}, res.X == [(vals[k], xc) ...] consistency with res.X_ind, workspace warm slots NaN-initialized, per-branch slots independent.
  5. set_eval_nodes! and callable consistency (== between callable output and res fields at the same nodes), simulate domain-safety, copy-constructor and ArgumentError cases (LQA, x_lb kwarg, invalid inner_solver still validated).
  6. All pre-existing tests pass unchanged (backward compatibility).

Benchmarks

Add one case to benchmark/benchmarks.jl (e.g. 1d_cheb_product: the trade-off model with 2 discrete branches, scalar continuous part, Chebyshev basis): solve_PFI, solve_VFI_50iter, bellman_operator (workspace variant). Run full PkgBenchmark.benchmarkpkg on main and the PR branch; all pre-existing benchmarks must be within noise; report the table in the PR body per the repository convention. Expected cost shape: product sweep ~ K x the continuous sweep; state that in the PR body rather than implying it is free.

Known traps (encountered while building #93-#105; each cost real debugging time)

  • MUST hoist all object construction out of try blocks; catch only around the numerical optimize call, and rethrow InterruptException. A swallowed UndefVarError inside a broad catch silently disabled an entire solver path in ENH: Add multi-dimensional continuous actions #105's development.
  • MUST NOT use Fminbox(NelderMead()) anywhere — it can terminate at its starting point. The derivative-free multivariate path is cyclic coordinate-wise Brent (already implemented; reuse it).
  • Optim.only_fg! does not exist: use NLSolversBase: only_fg! (already a direct dependency).
  • Warm-start containers are the same arrays as the policy output; read the warm value before overwriting (per state, per branch).
  • The whole-sweep JIT warm-up dominates first-call timing; benchmark and time only after a warm-up pass.

Definition of done

  • All dispatch points above implemented; no MethodError reachable from public API with a ProductActions model.
  • Full test suite passes, including the new file; pre-existing tests untouched.
  • Benchmark table (both sides) in the PR body; pre-existing benchmarks within noise.
  • Docstrings for ProductActions and every touched public function; .github/copilot-instructions.md action-space section updated (and, per the synchronization rule recorded there, README.md/docs/src/index.md get a one-sentence mention, the example notebook none).
  • PR message follows the repository conventions (title prefix ENH:, design/validation/benchmark sections, generated-by attribution).

🤖 Drafted with Claude Code (Claude Fable 5)

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions