From 9cb21fd2095ac815472e61035a9b5e9d67f02715 Mon Sep 17 00:00:00 2001
From: zombie-einstein <13398815+zombie-einstein@users.noreply.github.com>
Date: Tue, 11 Nov 2025 10:39:26 +0000
Subject: [PATCH] Use correct action specification

---
 src/thants/specs.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/src/thants/specs.py b/src/thants/specs.py
index d41fce8..145f6c2 100644
--- a/src/thants/specs.py
+++ b/src/thants/specs.py
@@ -95,7 +95,7 @@ def get_observation_spec(
     )
 
 
-def get_action_spec(num_agents: int, num_signals: int) -> specs.BoundedArray:
+def get_action_spec(num_agents: int, num_signals: int) -> specs.MultiDiscreteArray:
     """
     Get action specification for a colony
 
@@ -111,10 +111,9 @@ def get_action_spec(num_agents: int, num_signals: int) -> specs.BoundedArray:
     Spec
         Action specification
     """
-    return specs.BoundedArray(
-        shape=(num_agents,),
-        minimum=0,
-        maximum=7 + num_signals,
+    n_actions = 7 + num_signals
+    return specs.MultiDiscreteArray(
+        num_values=jnp.full((num_agents,), n_actions, dtype=int),
         dtype=int,
     )