Fix multiplex

Author: HenriDeh

src/normalization.jl

@@ -122,7 +122,11 @@ When sampling a normalized trace, it will first normalize the samples to zero me
 variance. Traces that do not have a normalizer are sample as usual.
 
 Note that when used in combination with [`Episodes`](@ref), `NormalizedTraces` must wrap 
-the `Episodes` struct, not the inner `AbstractTraces` contained in an `Episode`, otherwise
+the `Episodes` struct, not the inner `AbstractTraces` contained in an `Episode`, otherwise 
+the running estimate will reset after each episode.
+
+When used with a MultiplexTraces, the normalizer used with for one symbol (e.g. :state) will
+be the same used for the other one (e.g. :next_state).
 
 Preconfigured normalizers are provided for scalar (see [`scalar_normalizer`](@ref)) and 
 arrays (see [`array_normalizer`](@ref)).
@@ -131,6 +135,7 @@ arrays (see [`array_normalizer`](@ref)).
 ```
 t = CircularArraySARTTraces(capacity = 10, state = Float64 => (5,))
 nt = NormalizedTraces(t, reward = scalar_normalizer(), state = array_normalizer((5,)))
+# :next_state will also be normalized. 
 traj = Trajectory(
     container = nt,
     sampler = BatchSampler(10)
@@ -147,6 +152,18 @@ function NormalizedTraces(traces::AbstractTraces{names, TT}; trace_normalizer_pa
         @assert key in keys(traces) "Traces do not have key $key, valid keys are $(keys(traces))."
     end
     nt = (; trace_normalizer_pairs...)
+    for trace in traces.traces
+        #check if all traces of MultiplexTraces are in pairs
+        if trace isa MultiplexTraces
+            if length(intersect(keys(trace), keys(trace_normalizer_pairs))) in [0, length(keys(trace))] #check if none or all keys are in normalizers
+                continue
+            else #if not then one is missing
+                present_key = only(intersect(keys(trace), keys(trace_normalizer_pairs)))
+                absent_key = only(setdiff(keys(trace), keys(trace_normalizer_pairs)))
+                nt = merge(nt, (;(absent_key => nt[present_key],)...)) #assign the same normalizer
+            end
+        end
+    end
     NormalizedTraces{names, TT, typeof(traces), keys(nt), typeof(values(nt))}(traces, nt)
 end
 

test/normalization.jl

@@ -3,21 +3,21 @@ using ReinforcementLearningTrajectories
 import ReinforcementLearningTrajectories: sample
 import OnlineStats: mean, std
 
-@testset "normalization.jl" begin
+#@testset "normalization.jl" begin
     t = CircularArraySARTTraces(capacity = 10, state = Float64 => (5,))
     nt = NormalizedTraces(t, reward = scalar_normalizer(), state = array_normalizer((5,)))
     m = mean(0:4)
     s = std(0:4)
-
+    ss = std([0,1,2,2,3,4])
     for i in 0:4
         r = ((1.0:5.0) .+ i) .% 5
         push!(nt, (state = [r;], action = 1, reward = Float32(i), terminal = false))
     end
-    push!(nt, (next_state = fill(m, 5), next_action = 1)) #does not update because next_state is not in keys of normalizers. Is this desirable or not ?
+    push!(nt, (next_state = fill(m, 5), next_action = 1)) #this also updates state moments
 
     @test mean(nt.normalizers[:reward].os) == m && std(nt.normalizers[:reward].os) == s
     @test all(nt.normalizers[:state].os) do moments
-        mean(moments) == m && std(moments) == s
+        mean(moments) == m && std(moments) == ss
     end
 
     unnormalized_batch = t[[1:5;]]
@@ -31,7 +31,8 @@ import OnlineStats: mean, std
         controller = InsertSampleRatioController(ratio = Inf, threshold = 0)
     )
     normalized_batch = sample(traj)
-    @test all(extrema(normalized_batch[:state]) .≈ ((0, 4) .- m)./s)
+    @test all(extrema(normalized_batch[:state]) .≈ ((0, 4) .- m)./ss)
+    @test all(extrema(normalized_batch[:next_state]) .≈ ((0, 4) .- m)./ss)
     @test all(extrema(normalized_batch[:reward]) .≈ ((0, 4) .- m)./s)
     #check for no mutation
     unnormalized_batch = t[[1:5;]]