diff --git a/tensorneat/examples/brax/half_cheetah.py b/tensorneat/examples/brax/half_cheetah.py
index da4823c..4d31efe 100644
--- a/tensorneat/examples/brax/half_cheetah.py
+++ b/tensorneat/examples/brax/half_cheetah.py
@@ -23,7 +23,7 @@ if __name__ == '__main__':
             ),
         ),
         problem=BraxEnv(
-            env_name='halhcheetah',
+            env_name='halfcheetah',
         ),
         generation_limit=10000,
         fitness_target=5000
diff --git a/tensorneat/problem/rl_env/brax_env.py b/tensorneat/problem/rl_env/brax_env.py
index 9c34501..4e0b505 100644
--- a/tensorneat/problem/rl_env/brax_env.py
+++ b/tensorneat/problem/rl_env/brax_env.py
@@ -39,7 +39,7 @@ class BraxEnv(RLEnv):
 
         def step(key, env_state, obs):
             key, _ = jax.random.split(key)
-            action = act_func(state, obs, params)
+            action = act_func(obs, params)
             obs, env_state, r, done, _ = self.step(randkey, env_state, action)
             return key, env_state, obs, r, done