Task: robust_value_iteration_hypercube
Category: Operations research
Difficulty: 3
Tags: Dynamic Programming,Robust Optimization,Markov Decision Process,Value Iteration
Source: Grand-Cl\'ement, Petrik \
Description: Vieille 2026, "Beyond Discounted Returns: Robust Markov Decision Processes with Average and Blackwell Optimality," Operations Research: https://doi.org/10.1287/opre.2023.0694 & Robust value iteration for MDPs with hypercube transition-uncertainty sets.

Prompt template (YAML):
task_description: |
  A robust Markov Decision Process (MDP) value iteration algorithm where transition probabilities are uncertain and bounded within a hypercube (L-infinity norm).
  At each iteration, the algorithm computes the worst-case expected value by sorting the next-state values and greedily assigning maximum possible transition probabilities to the lowest values until the probability distribution sums to 1.
  
  Given `rew` (S x A), `P_up` (S x A x S), `P_down` (S x A x S), `gamma`, and `T`:
  1. Initialize `V` as an array of zeros of size `S`.
  2. For `t` from 0 to `T - 1`:
     a. Sort `V` in ascending order to get an array of indices `permutation`.
     b. Initialize `V_next` as an array of zeros of size `S`.
     c. For each state `s`:
        i. Initialize `max_Q = -infinity`.
        ii. For each action `a`:
            - Calculate `total_P_down = sum_{j=0}^{S-1} P_down[s, a, j]`.
            - Initialize `run_sum = total_P_down - P_down[s, a, permutation[0]] + P_up[s, a, permutation[0]]`.
            - Initialize `index = 0`.
            - While `run_sum < 1.0` and `index < S - 1`:
                `index += 1`
                `run_sum += P_up[s, a, permutation[index]] - P_down[s, a, permutation[index]]`
            - Compute `expected_V = 0.0`.
            - For `j` from 0 to `index - 1`:
                `expected_V += P_up[s, a, permutation[j]] * V[permutation[j]]`
            - `expected_V += (1.0 - (run_sum - P_up[s, a, permutation[index]])) * V[permutation[index]]`
            - For `j` from `index + 1` to `S - 1`:
                `expected_V += P_down[s, a, permutation[j]] * V[permutation[j]]`
            - `Q_sa = rew[s, a] + gamma * expected_V`
            - `max_Q = max(max_Q, Q_sa)`
        iii. `V_next[s] = max_Q`
     d. `V = V_next`
  3. Return `V`.

interface: |
  void solution_compute(int S, int A, int T, float gamma, const float* rew, const float* P_up, const float* P_down, float* V);

input_size_notes: |
  - S is the number of states, up to 2000
  - A is the number of actions, up to 20
  - T is the number of iterations, up to 500
  - rew is of size S * A
  - P_up and P_down are of size S * A * S

output_contract: |
  - The output array `V` must be filled with the computed values after `T` iterations.
  - The output array `V` is pre-allocated on the host.
  - `solution_compute` should allocate device memory, copy the inputs, perform the computation, and copy the result back to the host array `V`.

algorithm_background: |
  Robust Value Iteration finds the optimal policy for an MDP when the transition probabilities are not known exactly but lie in an uncertainty set.
  For L-infinity norm (hypercube) uncertainty, the worst-case transition probabilities can be found efficiently by sorting the states by their current value and greedily assigning the maximum possible probability (P_up) to the states with the lowest values, and the minimum possible probability (P_down) to the rest, such that the total probability sums to 1.

hints_l1: |
  - The sorting step can be done using Thrust or CUB on the GPU.
  - Since the array to sort is of size S, it can be sorted quickly.
  - The inner loops over `s` and `a` can be parallelized across GPU threads.

hints_l2: |
  - Precompute `total_P_down` for each `(s, a)` during initialization to avoid an O(S) reduction inside the hot loop.
  - Memory coalescing is crucial. Threads for different `(s, a)` access `P_up` and `P_down` using the same sorted `permutation` of indices.
  - If `P_up` and `P_down` are stored in `[s_prime, s, a]` layout instead of `[s, a, s_prime]`, memory accesses will be perfectly coalesced. Consider transposing the arrays at the start of `solution_compute`.


CPU reference implementation (C):
#include <stdlib.h>
#include <stdio.h>
#include <math.h>

typedef struct {
    float val;
    int idx;
} ValueIndex;

int compare_vi(const void* a, const void* b) {
    float va = ((ValueIndex*)a)->val;
    float vb = ((ValueIndex*)b)->val;
    if (va < vb) return -1;
    if (va > vb) return 1;
    return 0;
}

void solution_compute(int S, int A, int T, float gamma, const float* rew, const float* P_up, const float* P_down, float* V) {
    for (int s = 0; s < S; ++s) {
        V[s] = 0.0f;
    }
    
    float* V_next = (float*)malloc(S * sizeof(float));
    ValueIndex* vi = (ValueIndex*)malloc(S * sizeof(ValueIndex));
    
    float* total_P_down = (float*)malloc(S * A * sizeof(float));
    for (int s = 0; s < S; ++s) {
        for (int a = 0; a < A; ++a) {
            float sum = 0.0f;
            for (int j = 0; j < S; ++j) {
                sum += P_down[(s * A + a) * S + j];
            }
            total_P_down[s * A + a] = sum;
        }
    }
    
    for (int t = 0; t < T; ++t) {
        for (int s = 0; s < S; ++s) {
            vi[s].val = V[s];
            vi[s].idx = s;
        }
        qsort(vi, S, sizeof(ValueIndex), compare_vi);
        
        for (int s = 0; s < S; ++s) {
            float max_Q = -1e30f;
            for (int a = 0; a < A; ++a) {
                int sa_offset = (s * A + a) * S;
                float run_sum = total_P_down[s * A + a];
                
                int index = 0;
                int p_idx = vi[0].idx;
                run_sum += P_up[sa_offset + p_idx] - P_down[sa_offset + p_idx];
                
                while (run_sum < 1.0f && index < S - 1) {
                    index++;
                    p_idx = vi[index].idx;
                    run_sum += P_up[sa_offset + p_idx] - P_down[sa_offset + p_idx];
                }
                
                float expected_V = 0.0f;
                for (int j = 0; j < index; ++j) {
                    int p_j = vi[j].idx;
                    expected_V += P_up[sa_offset + p_j] * V[p_j];
                }
                
                int p_index = vi[index].idx;
                expected_V += (1.0f - (run_sum - P_up[sa_offset + p_index])) * V[p_index];
                
                for (int j = index + 1; j < S; ++j) {
                    int p_j = vi[j].idx;
                    expected_V += P_down[sa_offset + p_j] * V[p_j];
                }
                
                float Q_sa = rew[s * A + a] + gamma * expected_V;
                if (Q_sa > max_Q) {
                    max_Q = Q_sa;
                }
            }
            V_next[s] = max_Q;
        }
        
        for (int s = 0; s < S; ++s) {
            V[s] = V_next[s];
        }
    }
    
    free(V_next);
    free(vi);
    free(total_P_down);
}

void solution_free(void) {
    // No persistent state to free
}