[Feature] Replace RewardClipping with SignTransform in Atari examples (…

…pytorch#1870)
ai4co · Feb 12, 2024 · 69d44f5 · 69d44f5
1 parent 6f6c896
commit 69d44f5
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 8 deletions.
diff --git a/examples/a2c/utils_atari.py b/examples/a2c/utils_atari.py
@@ -20,8 +20,8 @@
     NoopResetEnv,
     ParallelEnv,
     Resize,
-    RewardClipping,
     RewardSum,
+    SignTransform,
     StepCounter,
     ToTensorImage,
     TransformedEnv,
@@ -73,7 +73,7 @@ def make_parallel_env(env_name, num_envs, device, is_test=False):
     env.append_transform(RewardSum())
     env.append_transform(StepCounter(max_steps=4500))
     if not is_test:
-        env.append_transform(RewardClipping(-1, 1))
+        env.append_transform(SignTransform(in_keys=["reward"]))
     env.append_transform(DoubleToFloat())
     env.append_transform(VecNorm(in_keys=["pixels"]))
     return env

diff --git a/examples/dqn/utils_atari.py b/examples/dqn/utils_atari.py
@@ -14,8 +14,8 @@
     GymEnv,
     NoopResetEnv,
     Resize,
-    RewardClipping,
     RewardSum,
+    SignTransform,
     StepCounter,
     ToTensorImage,
     TransformedEnv,
@@ -42,7 +42,7 @@ def make_env(env_name, frame_skip, device, is_test=False):
     env.append_transform(NoopResetEnv(noops=30, random=True))
     if not is_test:
         env.append_transform(EndOfLifeTransform())
-        env.append_transform(RewardClipping(-1, 1))
+        env.append_transform(SignTransform(in_keys=["reward"]))
     env.append_transform(ToTensorImage())
     env.append_transform(GrayScale())
     env.append_transform(Resize(84, 84))

diff --git a/examples/impala/utils.py b/examples/impala/utils.py
@@ -16,8 +16,8 @@
     GymEnv,
     NoopResetEnv,
     Resize,
-    RewardClipping,
     RewardSum,
+    SignTransform,
     StepCounter,
     ToTensorImage,
     TransformedEnv,
@@ -46,7 +46,7 @@ def make_env(env_name, device, is_test=False):
     env.append_transform(NoopResetEnv(noops=30, random=True))
     if not is_test:
         env.append_transform(EndOfLifeTransform())
-        env.append_transform(RewardClipping(-1, 1))
+        env.append_transform(SignTransform(in_keys=["reward"]))
     env.append_transform(ToTensorImage(from_int=False))
     env.append_transform(GrayScale())
     env.append_transform(Resize(84, 84))

diff --git a/examples/ppo/utils_atari.py b/examples/ppo/utils_atari.py
@@ -19,8 +19,8 @@
     NoopResetEnv,
     ParallelEnv,
     Resize,
-    RewardClipping,
     RewardSum,
+    SignTransform,
     StepCounter,
     ToTensorImage,
     TransformedEnv,
@@ -71,7 +71,7 @@ def make_parallel_env(env_name, num_envs, device, is_test=False):
     env.append_transform(RewardSum())
     env.append_transform(StepCounter(max_steps=4500))
     if not is_test:
-        env.append_transform(RewardClipping(-1, 1))
+        env.append_transform(SignTransform(in_keys=["reward"]))
     env.append_transform(DoubleToFloat())
     env.append_transform(VecNorm(in_keys=["pixels"]))
     return env