Update reward wrappers to be more consistent (#519)

* rewardoverridden -> rewardtransformed

* minor updates

* Update src/ReinforcementLearningEnvironments/src/environments/wrappers/RewardOverriddenEnv.jl

Co-authored-by: Jun Tian <[email protected]>

Author: albheim

docs/experiments/experiments/DQN/Dopamine_DQN_Atari.jl

@@ -80,7 +80,7 @@ function atari_env_factory(
     n_replica = nothing,
 )
     init(seed) =
-        RewardOverriddenEnv(
+        RewardTransformedEnv(
             StateCachedEnv(
                 StateTransformedEnv(
                     AtariEnv(;
@@ -101,8 +101,8 @@ function atari_env_factory(
                     ),
                     state_space_mapping= _ -> Space(fill(0..256, state_size..., n_frames))
                 )
-            ),
-            r -> clamp(r, -1, 1)
+            );
+            reward_mapping = r -> clamp(r, -1, 1)
         )
 
     if isnothing(n_replica)
@@ -130,7 +130,7 @@ end
 function (hook::TotalOriginalRewardPerEpisode)(
     ::PostActStage,
     agent,
-    env::RewardOverriddenEnv,
+    env::RewardTransformedEnv,
 )
     hook.reward += reward(env.env)
 end
@@ -153,7 +153,7 @@ end
 function (hook::TotalBatchOriginalRewardPerEpisode)(
     ::PostActStage,
     agent,
-    env::MultiThreadEnv{<:RewardOverriddenEnv},
+    env::MultiThreadEnv{<:RewardTransformedEnv},
 )
     for (i, e) in enumerate(env.envs)
         hook.reward[i] += reward(e.env)

docs/experiments/experiments/DQN/Dopamine_IQN_Atari.jl

@@ -85,7 +85,7 @@ function atari_env_factory(
     n_replica = nothing,
 )
     init(seed) =
-        RewardOverriddenEnv(
+        RewardTransformedEnv(
             StateCachedEnv(
                 StateTransformedEnv(
                     AtariEnv(;
@@ -106,8 +106,8 @@ function atari_env_factory(
                     ),
                     state_space_mapping= _ -> Space(fill(0..256, state_size..., n_frames))
                 )
-            ),
-            r -> clamp(r, -1, 1)
+            );
+            reward_mapping = r -> clamp(r, -1, 1)
         )
 
     if isnothing(n_replica)
@@ -135,7 +135,7 @@ end
 function (hook::TotalOriginalRewardPerEpisode)(
     ::PostActStage,
     agent,
-    env::RewardOverriddenEnv,
+    env::RewardTransformedEnv,
 )
     hook.reward += reward(env.env)
 end
@@ -158,7 +158,7 @@ end
 function (hook::TotalBatchOriginalRewardPerEpisode)(
     ::PostActStage,
     agent,
-    env::MultiThreadEnv{<:RewardOverriddenEnv},
+    env::MultiThreadEnv{<:RewardTransformedEnv},
 )
     for (i, e) in enumerate(env.envs)
         hook.reward[i] += reward(e.env)

docs/experiments/experiments/DQN/Dopamine_Rainbow_Atari.jl

@@ -84,7 +84,7 @@ function atari_env_factory(
     n_replica = nothing,
 )
     init(seed) =
-        RewardOverriddenEnv(
+        RewardTransformedEnv(
             StateCachedEnv(
                 StateTransformedEnv(
                     AtariEnv(;
@@ -105,8 +105,8 @@ function atari_env_factory(
                     ),
                     state_space_mapping= _ -> Space(fill(0..256, state_size..., n_frames))
                 )
-            ),
-            r -> clamp(r, -1, 1)
+            );
+            reward_mapping = r -> clamp(r, -1, 1)
         )
 
     if isnothing(n_replica)
@@ -134,7 +134,7 @@ end
 function (hook::TotalOriginalRewardPerEpisode)(
     ::PostActStage,
     agent,
-    env::RewardOverriddenEnv,
+    env::RewardTransformedEnv,
 )
     hook.reward += reward(env.env)
 end
@@ -157,7 +157,7 @@ end
 function (hook::TotalBatchOriginalRewardPerEpisode)(
     ::PostActStage,
     agent,
-    env::MultiThreadEnv{<:RewardOverriddenEnv},
+    env::MultiThreadEnv{<:RewardTransformedEnv},
 )
     for (i, e) in enumerate(env.envs)
         hook.reward[i] += reward(e.env)

docs/experiments/experiments/DQN/JuliaRL_BasicDQN_SingleRoomUndirected.jl

@@ -24,8 +24,8 @@ function RL.Experiment(
 
     env = GridWorlds.SingleRoomUndirectedModule.SingleRoomUndirected(rng=rng)
     env = GridWorlds.RLBaseEnv(env)
-    env = RLEnvs.StateTransformedEnv(env;state_mapping=x -> vec(Float32.(x)))
-    env = RewardOverriddenEnv(env, x -> x - convert(typeof(x), 0.01))
+    env = RLEnvs.StateTransformedEnv(env; state_mapping=x -> vec(Float32.(x)))
+    env = RewardTransformedEnv(env; reward_mapping = x -> x - convert(typeof(x), 0.01))
     env = MaxTimeoutEnv(env, 240)
 
     ns, na = length(state(env)), length(action_space(env))

src/ReinforcementLearningEnvironments/src/environments/wrappers/RewardOverriddenEnv.jl

@@ -3,12 +3,12 @@ export RewardOverriddenEnv
 """
     RewardOverriddenEnv(env, f)
 
-Apply `f` on `reward(env)`.
+Apply `f` on `env` to generate a custom reward.
 """
 struct RewardOverriddenEnv{F,E<:AbstractEnv} <: AbstractEnvWrapper
     env::E
     f::F
 end
 
 RLBase.reward(env::RewardOverriddenEnv, args...; kwargs...) =
-    env.f(reward(env.env, args...; kwargs...))
+    env.f(env.env, args...; kwargs...)

src/ReinforcementLearningEnvironments/src/environments/wrappers/RewardTransformedEnv.jl

@@ -0,0 +1,17 @@
+export RewardTransformedEnv
+
+"""
+    RewardTransformedEnv(env, f)
+
+Apply `f` on `reward(env)`.
+"""
+struct RewardTransformedEnv{F,E<:AbstractEnv} <: AbstractEnvWrapper
+    env::E
+    reward_mapping::F
+end
+
+RewardTransformedEnv(env; reward_mapping=identity) = 
+    RewardTransformedEnv(env, reward_mapping)
+
+RLBase.reward(env::RewardTransformedEnv, args...; kwargs...) =
+    env.reward_mapping(reward(env.env, args...; kwargs...))

src/ReinforcementLearningEnvironments/src/environments/wrappers/wrappers.jl

@@ -30,6 +30,7 @@ include("ActionTransformedEnv.jl")
 include("DefaultStateStyle.jl")
 include("MaxTimeoutEnv.jl")
 include("RewardOverriddenEnv.jl")
+include("RewardTransformedEnv.jl")
 include("StateCachedEnv.jl")
 include("StateTransformedEnv.jl")
 include("StochasticEnv.jl")

src/ReinforcementLearningEnvironments/test/environments/wrappers/wrappers.jl

@@ -53,10 +53,24 @@
         @test is_terminated(env′) == false
     end
 
+    @testset "RewardTransformedEnv" begin
+        rng = StableRNG(123)
+        env = TigerProblemEnv(; rng=rng)
+        env′ = RewardTransformedEnv(env; reward_mapping = x -> sign(x))
+
+        RLBase.test_interfaces!(env′)
+        RLBase.test_runnable!(env′)
+
+        while !is_terminated(env′)
+            env′(rand(rng, legal_action_space(env′)))
+            @test reward(env′) ∈ (-1, 0, 1)
+        end
+    end
+
     @testset "RewardOverriddenEnv" begin
         rng = StableRNG(123)
         env = TigerProblemEnv(; rng=rng)
-        env′ = RewardOverriddenEnv(env, x -> sign(x))
+        env′ = RewardOverriddenEnv(env, e -> sign(reward(e)))
 
         RLBase.test_interfaces!(env′)
         RLBase.test_runnable!(env′)