diff --git a/donkey_gym/__init__.py b/donkey_gym/__init__.py
index 191cf90..7c9c533 100644
--- a/donkey_gym/__init__.py
+++ b/donkey_gym/__init__.py
@@ -3,7 +3,7 @@
 register(
     id='DonkeyVae-v0',
     entry_point='donkey_gym.vae_env.vae_env:DonkeyVAEEnv',
-    timestep_limit=None,
+    max_episode_steps=None,
 )
 #
 # register(