Skip to content

Commit 0dfe312

Browse files
Change generator type to iterator
This is a bit cleaner. As reccomended by an internal (to Google) static python analysis tool.
1 parent dc93b3d commit 0dfe312

File tree

1 file changed

+2
-3
lines changed

1 file changed

+2
-3
lines changed

compiler_opt/rl/imitation_learning/generate_bc_trajectories_lib.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
import contextlib
1818
import gin
1919
from typing import Any
20-
from collections.abc import Callable, Generator
20+
from collections.abc import Callable, Iterator
2121
import json
2222

2323
from absl import flags
@@ -518,8 +518,7 @@ def explore_at_state_generator(
518518
policy: Callable[[time_step.TimeStep | None], np.ndarray],
519519
explore_policy: Callable[[time_step.TimeStep], policy_step.PolicyStep],
520520
num_samples: int = 1,
521-
) -> Generator[tuple[tf.train.SequenceExample, ExplorationWithPolicy], None,
522-
None]:
521+
) -> Iterator[tuple[tf.train.SequenceExample, ExplorationWithPolicy]]:
523522
"""Generate sequence examples and next exploration policy while exploring.
524523
525524
Generator that defines how to explore at the given explore_step. This

0 commit comments

Comments
 (0)