From 1df382300a3c0a530463c062ae9c97ec0caaec08 Mon Sep 17 00:00:00 2001 From: Daniel Toyama Date: Tue, 7 Jul 2026 10:22:03 -0700 Subject: [PATCH] Internal change. PiperOrigin-RevId: 943965316 --- envlogger/environment_logger_test.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/envlogger/environment_logger_test.py b/envlogger/environment_logger_test.py index 2c6e9d2..50aa737 100644 --- a/envlogger/environment_logger_test.py +++ b/envlogger/environment_logger_test.py @@ -21,7 +21,6 @@ import pickle import tempfile import threading -from typing import Optional from unittest import mock import uuid @@ -309,12 +308,12 @@ def test_episode_fn(self): """Checks that `episode_fn` produces expected custom data.""" v = 100 - def increment_fn(timestep, unused_action, unused_env) -> Optional[int]: + def increment_fn(timestep, unused_action, unused_env) -> int | None: """Increments `v` on the last timestep and returns it in even episodes.""" nonlocal v if timestep.first(): v += 1 - return np.int32(v) if v % 2 == 0 else None + return v if v % 2 == 0 else None env = catch_env.Catch() env = environment_logger.EnvLogger(