Change to support string fields

PiperOrigin-RevId: 452597155 Change-Id: Ife5d5fe0d35229344f13cfd4a3e3d2e48875c418
tensorflow · Jun 2, 2022 · d6f84de · d6f84de
1 parent 243d4a2
commit d6f84de
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tf_agents/environments/gym_wrapper.py b/tf_agents/environments/gym_wrapper.py
@@ -111,6 +111,8 @@ def nested_spec(spec, child_name):
       dtype = space.dtype
     else:
       dtype = dtype_map.get(gym.spaces.Box, np.float32)
+    if dtype == tf.string:
+      return specs.ArraySpec(shape=space.shape, dtype=dtype, name=name)
     minimum = np.asarray(space.low, dtype=dtype)
     maximum = np.asarray(space.high, dtype=dtype)
     if simplify_box_bounds: