From 9bd09fcace8bf6cca5b5a741faca5d95643489cd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Florent=20Altch=C3=A9?= <altche@google.com>
Date: Wed, 26 Aug 2020 10:05:31 +0100
Subject: [PATCH] Minor fix to type annotations.

PiperOrigin-RevId: 328491975
---
 rl_unplugged/dm_control_suite.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/rl_unplugged/dm_control_suite.py b/rl_unplugged/dm_control_suite.py
index b59939e..5360e03 100644
--- a/rl_unplugged/dm_control_suite.py
+++ b/rl_unplugged/dm_control_suite.py
@@ -30,7 +30,7 @@ observations.
 import collections
 import functools
 import os
-from typing import Dict, Tuple, Set
+from typing import Dict, Optional, Tuple, Set
 
 from acme import wrappers
 from dm_control import composer
@@ -778,7 +778,7 @@ def dataset(root_path: str,
             shapes: Dict[str, Tuple[int]],
             num_threads: int,
             batch_size: int,
-            uint8_features: Set[str] = None,
+            uint8_features: Optional[Set[str]] = None,
             num_shards: int = 100,
             shuffle_buffer_size: int = 100000,
             sarsa: bool = True) -> tf.data.Dataset: