Fix PyType errors in input normalization wrapper.

PiperOrigin-RevId: 389885508 Change-Id: I4b6d5984e933a61a64ce92732c9a81b3087f2058
google-deepmind · Aug 10, 2021 · eee201a · eee201a
1 parent 1bb5a0b
commit eee201a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/acme/agents/jax/normalization.py b/acme/agents/jax/normalization.py
@@ -258,7 +258,7 @@ def make_actor(
         backend='cpu')
 
 
-# Have to disable pytype invalid-annotation error, as it fails on Kokorox with
+# Have to disable pytype invalid-annotation error, as it fails on Kokoro with
 # Invalid type annotation 'TrainingState': Appears only once in the signature.
 def wrap_learner_core(
     learner_core: learner_core_lib.LearnerCore[reverb.ReplaySample,
@@ -273,7 +273,7 @@ def wrap_learner_core(
 
   State = NormalizationLearnerWrapperState[TrainingState]  # pylint: disable=invalid-name
 
-  def init(random_key: PRNGKey) -> State:
+  def init(random_key: PRNGKey):  # Returns a State.
     return NormalizationLearnerWrapperState(
         learner_core.init(random_key),
         running_statistics.init_state(environment_spec.observations))