Squeeze converted specs and observations if necessary.

PiperOrigin-RevId: 416275759

pysc2/env/BUILD

@@ -23,10 +23,12 @@ pytype_library(
     name = "converted_env",
     srcs = ["converted_env.py"],
     deps = [
+        "@dm_env_archive//:dm_env",
+        requirement("numpy"),
         "//pysc2/env/converter",
         "//pysc2/env/converter/proto:converter_py_pb2",
         "//pysc2/lib:actions",
-        "@dm_env_archive//:dm_env",
+        requirement("dm-tree"),
         requirement("typing_extensions"),
         "@s2client_proto//s2clientprotocol:common_py_pb2",
         "@s2client_proto//s2clientprotocol:raw_py_pb2",

pysc2/env/converted_env.py

@@ -19,9 +19,11 @@
 from typing import Any, Mapping, NamedTuple, Sequence
 
 import dm_env
+import numpy as np
 from pysc2.env.converter import converter as converter_lib
 from pysc2.env.converter.proto import converter_pb2
 from pysc2.lib import actions as sc2_actions
+import tree
 import typing_extensions
 
 from s2clientprotocol import common_pb2
@@ -31,6 +33,22 @@
 _BARRIER_TIMEOUT = 30.0
 
 
+def squeeze_if_necessary(x: np.ndarray) -> np.ndarray:
+  """Remove the trailing 1 of inputs."""
+  if x.shape and x.shape[-1] == 1:
+    return np.squeeze(x, axis=-1)
+  else:
+    return x
+
+
+def squeeze_spec_if_necessary(x):
+  """Remove the trailing 1 of specs inputs."""
+  if x.shape and x.shape[-1] == 1:
+    return x.replace(shape=x.shape[:-1])
+  else:
+    return x
+
+
 class ConverterFactory(typing_extensions.Protocol):
 
   def __call__(self, game_info: sc2api_pb2.ResponseGameInfo) -> Any:
@@ -121,7 +139,7 @@ def step(self, actions) -> dm_env.TimeStep:
     return timestep
 
   def observation_spec(self):
-    return self._obs_specs
+    return tree.map_structure(squeeze_spec_if_necessary, self._obs_specs)
 
   def action_spec(self):
     return self._action_specs
@@ -155,7 +173,8 @@ def _convert_obs(obs, converter):
       if not isinstance(obs, sc2api_pb2.ResponseObservation):
         obs = obs['_response_observation']()
       env_obs = converter_pb2.Observation(player=obs)
-      return converter.convert_observation(observation=env_obs)
+      env_obs = converter.convert_observation(observation=env_obs)
+      return tree.map_structure(squeeze_if_necessary, env_obs)
 
     # Merge the timesteps from a sequence to a single timestep
     return dm_env.TimeStep(
@@ -322,7 +341,9 @@ def get_environment_spec(
   """
   env_info = converter_pb2.EnvironmentInfo(game_info=_dummy_game_info())
   cvr = converter_lib.Converter(converter_settings, env_info)
-  return EnvironmentSpec(cvr.observation_spec(), cvr.action_spec())
+  obs_spec = tree.map_structure(squeeze_spec_if_necessary,
+                                cvr.observation_spec())
+  return EnvironmentSpec(obs_spec, cvr.action_spec())
 
 
 def make_converter_factories(