[vizdoom] expose combined dm action as discrete (#350)

## Summary
- add explicit discrete metadata to C++ specs and plumb it into Python
array specs
- use that metadata only for Vizdoom combined actions so dm_env exposes
a discrete action spec even though the transport dtype is float
- add a regression test for `make_spec("D1Basic-v1",
use_combined_action=True).action_spec()`

## Testing
- python3 -m py_compile envpool/python/data.py
envpool/python/protocol.py envpool/vizdoom/vizdoom_test.py
- devbox validation in progress: `make lint`, targeted Bazel tests, then
`make bazel-test`

Author: Trinkle23897

envpool/core/py_envpool.h

@@ -104,7 +104,8 @@ template <typename Spec>
 struct SpecTupleHelper {
   static decltype(auto) Make(const Spec& spec) {
     return std::make_tuple(py::dtype::of<typename Spec::dtype>(), spec.shape,
-                           spec.bounds, spec.elementwise_bounds);
+                           spec.bounds, spec.elementwise_bounds,
+                           spec.is_discrete);
   }
 };
 
@@ -121,7 +122,8 @@ struct SpecTupleHelper<Spec<Container<dtype>>> {
     return std::make_tuple(py::dtype::of<dtype>(),
                            std::make_tuple(spec.shape, spec.inner_spec.shape),
                            spec.inner_spec.bounds,
-                           spec.inner_spec.elementwise_bounds);
+                           spec.inner_spec.elementwise_bounds,
+                           spec.inner_spec.is_discrete);
   }
 };
 

envpool/core/spec.h

@@ -59,6 +59,7 @@ template <typename D>
 class Spec : public ShapeSpec {
  public:
   using dtype = D;  // NOLINT
+  bool is_discrete{false};
   std::tuple<dtype, dtype> bounds = {std::numeric_limits<dtype>::min(),
                                      std::numeric_limits<dtype>::max()};
   std::tuple<std::vector<dtype>, std::vector<dtype>> elementwise_bounds;
@@ -109,6 +110,12 @@ class Spec : public ShapeSpec {
   }
 };
 
+template <typename D>
+Spec<D> MarkDiscrete(Spec<D> spec) {
+  spec.is_discrete = true;
+  return spec;
+}
+
 template <typename dtype>
 class TArray;
 

envpool/python/data.py

@@ -28,6 +28,34 @@
 ACTION_THRESHOLD = 2**20
 
 
+def _maybe_scalar_int(value: Any) -> int | None:
+    arr = np.asarray(value)
+    if arr.size != 1:
+        return None
+    scalar = arr.item()
+    integer = int(scalar)
+    if not np.isclose(scalar, integer):
+        return None
+    return integer
+
+
+def _maybe_discrete_range(
+    spec: ArraySpec, spec_type: str
+) -> tuple[int, int] | None:
+    if np.prod(np.abs(spec.shape)) != 1:
+        return None
+    minimum = _maybe_scalar_int(spec.minimum)
+    maximum = _maybe_scalar_int(spec.maximum)
+    if minimum is None or maximum is None or maximum >= ACTION_THRESHOLD:
+        return None
+    if spec_type == "act":
+        if not (spec.is_discrete or np.issubdtype(spec.dtype, np.integer)):
+            return None
+    elif not np.issubdtype(spec.dtype, np.integer):
+        return None
+    return minimum, maximum - minimum + 1
+
+
 def to_nested_dict(
     flatten_dict: dict[str, Any], generator: type = dict
 ) -> dict[str, Any]:
@@ -70,16 +98,15 @@ def dm_spec_transform(
     name: str, spec: ArraySpec, spec_type: str
 ) -> dm_env.specs.Array:
     """Transform ArraySpec to dm_env compatible specs."""
-    if (
-        np.prod(np.abs(spec.shape)) == 1
-        and np.isclose(spec.minimum, 0)
-        and spec.maximum < ACTION_THRESHOLD
-    ):
-        # special treatment for discrete action space
+    discrete_range = _maybe_discrete_range(spec, spec_type)
+    if discrete_range is not None and discrete_range[0] == 0:
+        # dm_env only supports zero-based discrete arrays.
         return dm_env.specs.DiscreteArray(
             name=name,
-            dtype=spec.dtype,
-            num_values=int(spec.maximum - spec.minimum + 1),
+            dtype=spec.dtype
+            if np.issubdtype(spec.dtype, np.integer)
+            else np.int32,
+            num_values=discrete_range[1],
         )
     return dm_env.specs.BoundedArray(
         name=name,
@@ -92,19 +119,13 @@ def dm_spec_transform(
 
 def gym_spec_transform(name: str, spec: ArraySpec, spec_type: str) -> gym.Space:
     """Transform ArraySpec to gym.Env compatible spaces."""
-    if (
-        np.prod(np.abs(spec.shape)) == 1
-        and np.isclose(spec.minimum, 0)
-        and spec.maximum < ACTION_THRESHOLD
-    ):
-        # special treatment for discrete action space
-        discrete_range = int(spec.maximum - spec.minimum + 1)
+    discrete_range = _maybe_discrete_range(spec, spec_type)
+    if discrete_range is not None:
+        start, num_values = discrete_range
         try:
-            return gym.spaces.Discrete(
-                n=discrete_range, start=int(spec.minimum)
-            )
+            return gym.spaces.Discrete(n=num_values, start=start)
         except TypeError:  # old gym version doesn't have `start`
-            return gym.spaces.Discrete(n=discrete_range)
+            return gym.spaces.Discrete(n=num_values)
     return gym.spaces.Box(
         shape=[s for s in spec.shape if s != -1],
         dtype=spec.dtype,
@@ -117,16 +138,10 @@ def gymnasium_spec_transform(
     name: str, spec: ArraySpec, spec_type: str
 ) -> gymnasium.Space:
     """Transform ArraySpec to gymnasium.Env compatible spaces."""
-    if (
-        np.prod(np.abs(spec.shape)) == 1
-        and np.isclose(spec.minimum, 0)
-        and spec.maximum < ACTION_THRESHOLD
-    ):
-        # special treatment for discrete action space
-        discrete_range = int(spec.maximum - spec.minimum + 1)
-        return gymnasium.spaces.Discrete(
-            n=discrete_range, start=int(spec.minimum)
-        )
+    discrete_range = _maybe_discrete_range(spec, spec_type)
+    if discrete_range is not None:
+        start, num_values = discrete_range
+        return gymnasium.spaces.Discrete(n=num_values, start=start)
     return gymnasium.spaces.Box(
         shape=[s for s in spec.shape if s != -1],
         dtype=spec.dtype,

envpool/python/protocol.py

@@ -96,10 +96,12 @@ def __init__(
         shape: list[int],
         bounds: tuple[Any, Any],
         element_wise_bounds: tuple[Any, Any],
+        is_discrete: bool = False,
     ):
         """Constructor of ArraySpec."""
         self.dtype = dtype
         self.shape = shape
+        self.is_discrete = is_discrete
         if element_wise_bounds[0]:
             self.minimum = np.array(element_wise_bounds[0])
         else:

envpool/vizdoom/vizdoom_env.h

@@ -123,8 +123,8 @@ class VizdoomEnvFns {
     }
     auto action_set =
         BuildActionSet(button_list, conf["force_speed"_], delta_config);
-    return MakeDict(
-        "action"_.Bind(Spec<double>({-1}, {0.0, action_set.size() - 1.0})));
+    return MakeDict("action"_.Bind(
+        MarkDiscrete(Spec<double>({-1}, {0.0, action_set.size() - 1.0}))));
   }
 };
 

envpool/vizdoom/vizdoom_test.py

@@ -20,7 +20,7 @@
 from absl.testing import absltest
 
 import envpool.vizdoom.registration  # noqa: F401
-from envpool.registration import make_dm, make_gym
+from envpool.registration import make_dm, make_gym, make_spec
 
 
 class _VizdoomEnvPoolBasicTest(absltest.TestCase):
@@ -155,6 +155,12 @@ def test_obs_space(self) -> None:
             == 1 * 4
         )
 
+    def test_action_spec(self) -> None:
+        spec = make_spec("D1Basic-v1", use_combined_action=True)
+        action_spec = spec.action_spec()
+        assert action_spec.num_values == 6
+        assert np.issubdtype(action_spec.dtype, np.integer)
+
     def test_explicit_reset_with_episodic_life_gymnasium(self) -> None:
         env = make_gym(
             "D1Basic-v1",