try to fix bugs of ActionTransformedEnv (#447)

Author: peterchen96

src/ReinforcementLearningCore/src/policies/agents/agent.jl

@@ -137,11 +137,11 @@ function RLBase.update!(
     # the global rng. In theory it shouldn't affect the performance of specific
     # algorithm.
     # TODO: how to inject a local rng here to avoid polluting the global rng
-    action = rand(action_space(env))
 
     s = policy isa NamedPolicy ? state(env, nameof(policy)) : state(env)
+    a = policy isa NamedPolicy ? rand(action_space(env, nameof(policy))) : rand(action_space(env))
     push!(trajectory[:state], s)
-    push!(trajectory[:action], action)
+    push!(trajectory[:action], a)
     if haskey(trajectory, :legal_actions_mask)
         lasm =
             policy isa NamedPolicy ? legal_action_space_mask(env, nameof(policy)) :

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

@@ -17,10 +17,10 @@ ActionTransformedEnv(env; action_mapping = identity, action_space_mapping = iden
     ActionTransformedEnv(env, action_mapping, action_space_mapping)
 
 RLBase.action_space(env::ActionTransformedEnv, args...) =
-    env.action_space_mapping(action_space(env.env), args...)
+    env.action_space_mapping(action_space(env.env, args...))
 
 RLBase.legal_action_space(env::ActionTransformedEnv, args...) =
-    env.action_space_mapping(legal_action_space(env.env), args...)
+    env.action_space_mapping(legal_action_space(env.env, args...))
 
 (env::ActionTransformedEnv)(action, args...; kwargs...) =
     env.env(env.action_mapping(action), args...; kwargs...)