diff --git a/rllib/examples/pettingzoo_env.py b/rllib/examples/pettingzoo_env.py
index 1f38338bf..45effd16f 100644
--- a/rllib/examples/pettingzoo_env.py
+++ b/rllib/examples/pettingzoo_env.py
@@ -6,7 +6,7 @@ except ImportError:
     from ray.rllib.agents.registry import get_agent_class
 from ray.tune.registry import register_env
 from ray.rllib.env import PettingZooEnv
-from pettingzoo.gamma import prison_v0
+from pettingzoo.butterfly import pistonball_v0
 from supersuit.aec_wrappers import normalize_obs, dtype, color_reduction
 
 from numpy import float32
@@ -24,10 +24,10 @@ if __name__ == "__main__":
 
     # function that outputs the environment you wish to register.
     def env_creator(config):
-        env = prison_v0.env(num_floors=config.get("num_floors", 4))
+        env = pistonball_v0.env(local_ratio=config.get("local_ratio", 0.2))
         env = dtype(env, dtype=float32)
-        env = color_reduction(env, dtype=float32)
-        env = normalize_obs(env, mode="R")
+        env = color_reduction(env, mode="R")
+        env = normalize_obs(env)
         return env
 
     num_cpus = 1
@@ -38,7 +38,7 @@ if __name__ == "__main__":
 
     # 2. Set environment config. This will be passed to
     # the env_creator function via the register env lambda below
-    config["env_config"] = {"num_floors": 5}
+    config["env_config"] = {"local_ratio": 0.5}
 
     # 3. Register env
     register_env("prison", lambda config: PettingZooEnv(env_creator(config)))