struct PolicyRewardTest <: TrainingTest
    idx::Integer
    state::AbstractState
    image::AbstractImage
    actions::Vector{<:Integer}
    action::Integer
    reward::Float32
end

mutable struct PolicyRewardTests <: TrainingTests
    tests::Vector{PolicyRewardTest}
    images::AbstractArray{Float32, 4}
    actual_reward::AbstractArray{Float32, 2}
    period::Int
    warmup::Int
    function PolicyRewardTests(tests; period=1, warmup=0)
        new(
            tests, 
            tensorize(get_image.(tests)),
            Flux.unsqueeze(get_reward.(tests), 1),
            period,
            warmup,
        )
    end
end

function gettestfn(tests::PolicyRewardTests)
    function fn(n::Network)
        initial_output = initialinferencefn(tests)(n)
        acc, fails = getpolicyacc(tests, initial_output)
        recurrent_output = recurrentinferencefn(tests, initial_output)(n)
        acc, fails, getrewardmse(tests, recurrent_output)
    end
end

struct PolicyRewardTestsSVG
    tests::PolicyRewardTests
end
svg(x::PolicyRewardTests) = PolicyRewardTestsSVG(x)

function Base.show(io::IO, m::MIME"image/svg+xml", x::PolicyRewardTestsSVG)
    num_tests = length(x.tests)
    range = 1:(min(10, num_tests))
    subset = 1:num_tests
    subset_indices = subset[range]
    svgbatch = [gridsvg(x.tests[i]) for i in subset_indices]
    Base.show(io, m, GridSVG(svgbatch, main_title="PolicyReward Tests"))
end