code: def test_normalize_return():
      env = DummyRewardEnv(return_reward_idx=0)
      env = NormalizeReward(env)
      env.reset()
      env.step(env.action_space.sample())
      assert_almost_equal(
          env.return_rms.mean,
          np.mean([1]),  # [first return]
          decimal=4,
      )
      env.step(env.action_space.sample())
      assert_almost_equal(
          env.return_rms.mean,
          np.mean(
              [2 + env.gamma * 1, 1]
          ),  # [second return, first return]
          decimal=4,
      )
  
file_path: gym/tests/wrappers/test_normalize.py
imports: from typing import Optional
  import numpy
  from numpy.testing import assert_almost_equal
  import gym
  from gym.wrappers.normalize import NormalizeObservation, NormalizeReward
label_desc: Use image node_id 4 for calling a global function with example usage: test_normalize_return() without return types