struct TrainingBatchTensor
    images::AbstractArray{Float32, 4}
    image_targets::AbstractArray{Float32, 2}
    actions::AbstractArray{Float32, 4}
    action_targets::AbstractArray{Float32, 2}
    function TrainingBatchTensor(
        images::AbstractArray, 
        image_targets::AbstractArray, 
        actions::AbstractArray, 
        action_targets::AbstractArray)
        new(gpu(images), gpu(image_targets), gpu(actions), gpu(action_targets))
    end
end

images(x::TrainingBatchTensor) = x.images
actions(x::TrainingBatchTensor) = x.actions
image_targets(x::TrainingBatchTensor) = view(x.image_targets, :, :)
action_indices(x::TrainingBatchTensor) = 1:maximum(view(x.action_targets, 2, :))
calc_num_initial_inference_targets(x::TrainingBatchTensor) = size(x.image_targets, 2)
calc_num_recurrent_inference_targets(x::TrainingBatchTensor) = size(x.action_targets, 2)

function TrainingBatchTensor(game_type::Type{<:AbstractEnv}, action_space_size::Integer, samples::Vector)
    images = tensorize(map(image, samples))
    image_targets = reduce(hcat, map(s->begin
        target = s.targets[1]
        [
            value(target), 
            reward(target), 
            policy(target)...
        ] 
    end, samples))
    action_idx_to_targets = Dict()
    for (game_idx, sample) in enumerate(samples)
        gradient_scale = Float32(1) / length(sample.actions)
        for (action_idx, action) in enumerate(sample.actions)
            target = sample.targets[action_idx+1]
            if !haskey(action_idx_to_targets, action_idx)
                action_idx_to_targets[action_idx] = []
            end
            push!(action_idx_to_targets[action_idx], 
                [
                    game_idx, 
                    action_idx, 
                    action,
                    value(target), 
                    reward(target), 
                    policy(target)...,
                    gradient_scale,
                ]
            )
        end
    end
    ordered_action_indices = sort(collect(keys(action_idx_to_targets)))
    action_targets = reduce(hcat, 
        map(action_idx->begin
            reduce(hcat, action_idx_to_targets[action_idx])
        end, ordered_action_indices))

    # @todo, extract this logic out into the game logic
    #reflect = [0 0 1; 0 1 0; 1 0 0]
    ordered_action_planes = map(action_idx->begin
        planes = map(
            image->image.channels,
            get_action_image.(
                game_type, 
                map(x->Int(x[3]), action_idx_to_targets[action_idx])
            )
        )
        # ohb = Flux.onehotbatch(actions, 1:action_space_size)
        # action_mat = convert(Array{Float32, 4}, reshape(ohb, 3, 3, 1, :))
        # rotated = map(i->reflect * (action_mat[:,:,1,i]'), 1:(length(actions)))
        reduce((a,b)->cat(a, b; dims=4), planes)
    end, ordered_action_indices)
    actions = reduce((a,b)->cat(a, b; dims=4), ordered_action_planes)
    TrainingBatchTensor(images, image_targets, actions, action_targets)
end