add tests on GPU

.github/workflows/CI.yml

@@ -19,7 +19,7 @@ jobs:
       matrix:
         version:
           - '1.6'
-          - '1.7'
+          - '1'
           - 'nightly'
         os:
           - ubuntu-latest

Project.toml

@@ -3,6 +3,7 @@ uuid = "6486599b-a3cd-4e92-a99a-2cea90cc8c3c"
 version = "0.1.5"
 
 [deps]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 CircularArrayBuffers = "9de3a189-e0c0-4e15-ba3b-b14b9fb0aec1"
 ElasticArrays = "fdbdab4c-e67f-52f5-8c3f-e7b388dad3d4"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
@@ -14,12 +15,13 @@ StackViews = "cae243ae-269e-4f55-b966-ac2d0dc13c15"
 CircularArrayBuffers = "0.1"
 ElasticArrays = "1"
 MacroTools = "0.5"
+OnlineStats = "1"
 StackViews = "0.1"
 julia = "1.6"
-OnlineStats = "1"
 
 [extras]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test"]
+test = ["Test", "CUDA"]

src/traces.jl

@@ -3,6 +3,7 @@ export Trace, Traces, MultiplexTraces, Episode, Episodes
 import MacroTools: @forward
 
 import CircularArrayBuffers
+import Adapt
 
 #####
 
@@ -13,11 +14,23 @@ Base.convert(::Type{AbstractTrace}, x::AbstractTrace) = x
 Base.summary(io::IO, t::AbstractTrace) = print(io, "$(length(t))-element $(nameof(typeof(t)))")
 
 #####
+
+"""
+    Trace(A::AbstractArray)
+
+Similar to
+[`Slices`](https://github.com/JuliaLang/julia/blob/master/base/slicearray.jl)
+which will be introduced in `[email protected]`. The main difference is that, the
+`axes` info in the `Slices` is static, while it may be dynamic with `Trace`.
+
+We only support slices along the last dimension since it's the most common usage
+in RL.
+"""
 struct Trace{T,E} <: AbstractTrace{E}
     parent::T
 end
 
-Base.summary(io::IO, t::Trace{T}) where {T} = print(io, "$(length(t))-element $(nameof(typeof(t))){$T}")
+Base.summary(io::IO, t::Trace{T}) where {T} = print(io, "$(length(t))-element$(length(t) > 0 ? 's' : "") $(nameof(typeof(t))){$T}")
 
 function Trace(x::T) where {T<:AbstractArray}
     E = eltype(x)
@@ -27,6 +40,8 @@ function Trace(x::T) where {T<:AbstractArray}
     Trace{T,SubArray{E,N,P,I,true}}(x)
 end
 
+Adapt.adapt_structure(to, t::Trace) = Trace(Adapt.adapt_structure(to, t.parent))
+
 Base.convert(::Type{AbstractTrace}, x::AbstractArray) = Trace(x)
 
 Base.size(x::Trace) = (size(x.parent, ndims(x.parent)),)
@@ -59,6 +74,21 @@ Base.haskey(t::AbstractTraces{names}, k::Symbol) where {names} = k in names
 
 #####
 
+"""
+Dedicated for `MultiplexTraces` to avoid scalar indexing when `view(view(t::MultiplexTrace, 1:end-1), I)`.
+"""
+struct RelativeTrace{left,right,T,E} <: AbstractTrace{E}
+    trace::Trace{T,E}
+end
+RelativeTrace{left,right}(t::Trace{T,E}) where {left,right,T,E} = RelativeTrace{left,right,T,E}(t)
+
+Base.size(x::RelativeTrace{0,-1}) = (max(0, length(x.trace) - 1),)
+Base.size(x::RelativeTrace{1,0}) = (max(0, length(x.trace) - 1),)
+Base.getindex(s::RelativeTrace{0,-1}, I) = getindex(s.trace, I)
+Base.getindex(s::RelativeTrace{1,0}, I) = getindex(s.trace, I .+ 1)
+Base.setindex!(s::RelativeTrace{0,-1}, v, I) = setindex!(s.trace, v, I)
+Base.setindex!(s::RelativeTrace{1,0}, v, I) = setindex!(s.trace, v, I .+ 1)
+
 """
     MultiplexTraces{names}(trace)
 
@@ -89,12 +119,14 @@ function MultiplexTraces{names}(t) where {names}
     MultiplexTraces{names,typeof(trace),eltype(trace)}(trace)
 end
 
+Adapt.adapt_structure(to, t::MultiplexTraces{names}) where {names} = MultiplexTraces{names}(Adapt.adapt_structure(to, t.trace))
+
 function Base.getindex(t::MultiplexTraces{names}, k::Symbol) where {names}
     a, b = names
     if k == a
-        convert(AbstractTrace, t.trace[1:end-1])
+        RelativeTrace{0,-1}(convert(AbstractTrace, t.trace))
     elseif k == b
-        convert(AbstractTrace, t.trace[2:end])
+        RelativeTrace{1,0}(convert(AbstractTrace, t.trace))
     else
         throw(ArgumentError("unknown trace name: $k"))
     end
@@ -133,6 +165,8 @@ end
 
 Episode(t::AbstractTraces{names,T}) where {names,T} = Episode{typeof(t),names,T}(t, Ref(false))
 
+Adapt.adapt_structure(to, t::Episode{T,names,E}) where {T,names,E} = Episode{T,names,E}(Adapt.adapt_structure(to, t.traces), t.is_terminated)
+
 @forward Episode.traces Base.getindex, Base.setindex!, Base.size
 
 Base.getindex(e::Episode) = getindex(e.is_terminated)
@@ -175,6 +209,11 @@ struct Episodes{names,E,T} <: AbstractTraces{names,E}
     inds::Vector{Tuple{Int,Int}}
 end
 
+Adapt.adapt_structure(to, t::Episodes) =
+    Episodes() do
+        Adapt.adapt_structure(to, t.init())
+    end
+
 function Episodes(init)
     x = init()
     T = typeof(x)
@@ -249,6 +288,11 @@ struct Traces{names,T,N,E} <: AbstractTraces{names,E}
     inds::NamedTuple{names,NTuple{N,Int}}
 end
 
+function Adapt.adapt_structure(to, t::Traces{names,T,N,E}) where {names,T,N,E}
+    data = Adapt.adapt_structure(to, t.traces)
+    # FIXME: `E` is not adapted here
+    Traces{names,typeof(data),length(names),E}(data, t.inds)
+end
 
 function Traces(; kw...)
     data = map(x -> convert(AbstractTrace, x), values(kw))

test/common.jl

@@ -31,20 +31,21 @@ end
         action=Float32 => (2,),
         reward=Float32 => (),
         terminal=Bool => ()
-    )
+    ) |> gpu
 
     @test t isa CircularArraySARTTraces
 
-    push!(t, (state=ones(Float32, 2, 3), action=ones(Float32, 2)))
+    push!(t, (state=ones(Float32, 2, 3), action=ones(Float32, 2)) |> gpu)
     @test length(t) == 0
 
-    push!(t, (reward=1.0f0, terminal=false))
+    push!(t, (reward=1.0f0, terminal=false) |> gpu)
     @test length(t) == 0 # next_state and next_action is still missing
 
-    push!(t, (next_state=ones(Float32, 2, 3) * 2, next_action=ones(Float32, 2) * 2))
+    push!(t, (next_state=ones(Float32, 2, 3) * 2, next_action=ones(Float32, 2) * 2) |> gpu)
     @test length(t) == 1
 
-    @test t[1] == (
+    # this will trigger the scalar indexing of CuArray
+    CUDA.@allowscalar @test t[1] == (
         state=ones(Float32, 2, 3),
         next_state=ones(Float32, 2, 3) * 2,
         action=ones(Float32, 2),
@@ -54,28 +55,30 @@ end
     )
 
     push!(t, (reward=2.0f0, terminal=false))
-    push!(t, (state=ones(Float32, 2, 3) * 3, action=ones(Float32, 2) * 3))
+    push!(t, (state=ones(Float32, 2, 3) * 3, action=ones(Float32, 2) * 3) |> gpu)
 
     @test length(t) == 2
 
     push!(t, (reward=3.0f0, terminal=false))
-    push!(t, (state=ones(Float32, 2, 3) * 4, action=ones(Float32, 2) * 4))
+    push!(t, (state=ones(Float32, 2, 3) * 4, action=ones(Float32, 2) * 4) |> gpu)
 
     @test length(t) == 3
 
     push!(t, (reward=4.0f0, terminal=false))
-    push!(t, (state=ones(Float32, 2, 3) * 5, action=ones(Float32, 2) * 5))
+    push!(t, (state=ones(Float32, 2, 3) * 5, action=ones(Float32, 2) * 5) |> gpu)
 
     @test length(t) == 3
-    @test t[1] == (
+
+    # this will trigger the scalar indexing of CuArray
+    CUDA.@allowscalar @test t[1] == (
         state=ones(Float32, 2, 3) * 2,
         next_state=ones(Float32, 2, 3) * 3,
         action=ones(Float32, 2) * 2,
         next_action=ones(Float32, 2) * 3,
         reward=2.0f0,
         terminal=false,
     )
-    @test t[end] == (
+    CUDA.@allowscalar @test t[end] == (
         state=ones(Float32, 2, 3) * 4,
         next_state=ones(Float32, 2, 3) * 5,
         action=ones(Float32, 2) * 4,
@@ -87,8 +90,8 @@ end
     batch = t[1:3]
     @test size(batch.state) == (2, 3, 3)
     @test size(batch.action) == (2, 3)
-    @test batch.reward == [2.0, 3.0, 4.0]
-    @test batch.terminal == Bool[0, 0, 0]
+    @test batch.reward == [2.0, 3.0, 4.0] |> gpu
+    @test batch.terminal == Bool[0, 0, 0] |> gpu
 end
 
 @testset "ElasticArraySARTTraces" begin

test/runtests.jl

@@ -1,6 +1,14 @@
 using ReinforcementLearningTrajectories
 using CircularArrayBuffers
 using Test
+using CUDA
+using Adapt
+
+struct TestAdaptor end
+
+gpu(x) = Adapt.adapt(TestAdaptor(), x)
+
+Adapt.adapt_storage(to::TestAdaptor, x) = CUDA.functional() ? CUDA.cu(x) : x
 
 @testset "ReinforcementLearningTrajectories.jl" begin
     include("traces.jl")