Fix ZeroDiscountonLifeLoss wrapper and add test covereage.

PiperOrigin-RevId: 318249207
Change-Id: I01c2f0b76d1b15fc42bf83538fbb1bc8c0757776

acme/wrappers/atari_wrapper.py

@@ -353,7 +353,7 @@ def __init__(self, environment: dm_env.Environment):
     self._last_num_lives = None
 
   def reset(self) -> dm_env.TimeStep:
-    timestep = self._env.reset()
+    timestep = self._environment.reset()
     self._reset_next_step = False
     self._last_num_lives = timestep.observation[LIVES_INDEX]
     return timestep

acme/wrappers/atari_wrapper_test.py

@@ -17,6 +17,7 @@
 
 import unittest
 from absl.testing import absltest
+from absl.testing import parameterized
 from acme.wrappers import atari_wrapper
 from dm_env import specs
 import numpy as np
@@ -33,12 +34,14 @@
 
 
 @unittest.skipIf(SKIP_GYM_TESTS, SKIP_GYM_MESSAGE)
-class AtariWrapperTest(absltest.TestCase):
+class AtariWrapperTest(parameterized.TestCase):
 
-  def test_pong(self):
+  @parameterized.parameters(True, False)
+  def test_pong(self, zero_discount_on_life_loss: bool):
     env = gym.make('PongNoFrameskip-v4', full_action_space=True)
     env = gym_wrapper.GymAtariAdapter(env)
-    env = atari_wrapper.AtariWrapper(env)
+    env = atari_wrapper.AtariWrapper(
+        env, zero_discount_on_life_loss=zero_discount_on_life_loss)
 
     # Test converted observation spec.
     observation_spec = env.observation_spec()