thrust_rl/env/games/
signaling.rs

1//! Referential signaling game — reference `CommunicatingEnvironment` (issue
2//! #274).
3//!
4//! The smallest environment that exercises the fixed-vocab comms surface from
5//! [`crate::multi_agent::comms`] end-to-end. It is the Phase 1 reference impl
6//! called out by [`docs/COMMS_DESIGN.md`](../../../../docs/COMMS_DESIGN.md)
7//! section 5: "one small reference env implementing `CommunicatingEnvironment`
8//! (e.g. a 2-agent referential signaling game)".
9//!
10//! # The game
11//!
12//! Two agents, a fixed vocabulary of `V` tokens:
13//!
14//! - **Agent 0 (speaker)** observes a hidden token `t ∈ 0..V` (one-hot) and has
15//!   *no task action* — its only action dim is a **message token** it emits.
16//! - **Agent 1 (listener)** sends nothing. Its observation is the message token
17//!   it received, and its *task action* is a guess `g ∈ 0..V`.
18//!
19//! Both agents earn `+1` when the listener's guess matches the hidden token and
20//! `0` otherwise, so the speaker is incentivized to transmit `t` faithfully.
21//! This is the canonical Lewis signaling game.
22//!
23//! # Comms wiring (the point of this env)
24//!
25//! | agent | `agent_action_space` | `message_vocab` | full action | `message_obs_size` |
26//! |-------|----------------------|-----------------|-------------|--------------------|
27//! | 0 speaker | `[]` (no task action) | `[V]` | `[token]` | `0` |
28//! | 1 listener | `[V]` (the guess) | `[]` | `[guess]` | `1` (received token) |
29//!
30//! Inside [`SignalingGame::step_multi`] the speaker's message dim is sliced off
31//! with [`crate::multi_agent::comms::split_action`], routed by broadcast, and
32//! written into the listener's observation with
33//! [`crate::multi_agent::comms::place_message`]. No new `MultiAgentResult`
34//! machinery is involved — the message rides inline in the action / observation
35//! vectors, exactly as Phase 1 prescribes.
36//!
37//! # Determinism
38//!
39//! The env consumes no internal RNG: the hidden token is fixed at construction
40//! (or by [`SignalingGame::set_hidden`]) and unchanged by stepping, so
41//! [`Environment::clone_state`] / [`Environment::restore_state`] reproduce
42//! every subsequent step exactly.
43
44use crate::{
45    env::{Environment, SpaceInfo, SpaceType, StepInfo, StepResult},
46    multi_agent::{
47        comms::{CommunicatingEnvironment, Delivery, place_message, split_action},
48        environment::{MultiAgentEnvironment, MultiAgentResult},
49        joint::{JointEnv, JointStepResult},
50    },
51};
52
53/// Agent id of the speaker (observes the hidden token, emits a message).
54pub const SPEAKER: usize = 0;
55/// Agent id of the listener (receives the message, guesses the token).
56pub const LISTENER: usize = 1;
57
58/// 2-agent referential signaling game over a fixed vocabulary.
59///
60/// See the [module docs](self) for the full description and the comms layout
61/// table.
62#[derive(Debug, Clone)]
63pub struct SignalingGame {
64    /// Vocabulary size `V`: both the hidden-token cardinality and the message /
65    /// guess cardinality.
66    vocab: usize,
67    /// Hidden token the speaker must transmit (`0..vocab`).
68    hidden: i64,
69    /// Message token the listener received on the most recent step (`-1` before
70    /// the first step of an episode).
71    received: i64,
72}
73
74impl SignalingGame {
75    /// Number of agents (always 2: speaker + listener).
76    pub const NUM_AGENTS: usize = 2;
77
78    /// Construct a signaling game over a `vocab`-token vocabulary with hidden
79    /// token `0`.
80    ///
81    /// # Panics
82    ///
83    /// Panics if `vocab == 0` — a signaling game needs at least one token.
84    pub fn new(vocab: usize) -> Self {
85        assert!(vocab > 0, "SignalingGame requires vocab >= 1");
86        Self { vocab, hidden: 0, received: -1 }
87    }
88
89    /// Construct a signaling game with an explicit hidden token (for tests /
90    /// deterministic rollouts).
91    ///
92    /// # Panics
93    ///
94    /// Panics if `vocab == 0` or `hidden` is outside `0..vocab`.
95    pub fn with_hidden(vocab: usize, hidden: i64) -> Self {
96        assert!(vocab > 0, "SignalingGame requires vocab >= 1");
97        assert!(
98            hidden >= 0 && (hidden as usize) < vocab,
99            "hidden token {hidden} out of range 0..{vocab}"
100        );
101        Self { vocab, hidden, received: -1 }
102    }
103
104    /// Vocabulary size `V`.
105    pub fn vocab(&self) -> usize {
106        self.vocab
107    }
108
109    /// The hidden token the speaker must transmit this episode.
110    pub fn hidden(&self) -> i64 {
111        self.hidden
112    }
113
114    /// The message token the listener received on the most recent step, or `-1`
115    /// if no step has occurred since the last reset.
116    pub fn received(&self) -> i64 {
117        self.received
118    }
119
120    /// Set the hidden token (`0..vocab`) and clear the received message.
121    ///
122    /// # Panics
123    ///
124    /// Panics if `hidden` is outside `0..vocab`.
125    pub fn set_hidden(&mut self, hidden: i64) {
126        assert!(
127            hidden >= 0 && (hidden as usize) < self.vocab,
128            "hidden token {hidden} out of range 0..{}",
129            self.vocab
130        );
131        self.hidden = hidden;
132        self.received = -1;
133    }
134
135    /// One-hot encoding of the hidden token — the speaker's observation.
136    fn speaker_obs(&self) -> Vec<f32> {
137        let mut obs = vec![0.0_f32; self.vocab];
138        obs[self.hidden as usize] = 1.0;
139        obs
140    }
141
142    /// The listener's observation: a single dim carrying the received token
143    /// (`-1` before the first step), laid out via
144    /// [`crate::multi_agent::comms::place_message`].
145    fn listener_obs(&self) -> Vec<f32> {
146        let mut obs = vec![-1.0_f32; self.message_obs_size(LISTENER)];
147        place_message(&mut obs, 0, &[self.received]);
148        obs
149    }
150}
151
152impl Default for SignalingGame {
153    /// A 2-token signaling game (the smallest non-trivial vocabulary).
154    fn default() -> Self {
155        Self::new(2)
156    }
157}
158
159impl Environment for SignalingGame {
160    type Action = i64;
161    /// Full snapshot: the env consumes no RNG, so restoring reproduces every
162    /// subsequent step exactly.
163    type State = (i64, i64);
164
165    fn reset(&mut self) {
166        self.received = -1;
167    }
168
169    /// Single-agent surface returns the speaker's observation. Multi-agent
170    /// consumers should use [`MultiAgentEnvironment::get_agent_observation`].
171    fn get_observation(&self) -> Vec<f32> {
172        self.speaker_obs()
173    }
174
175    /// Degenerate single-agent step: interpret `action` as the speaker's
176    /// emitted message and broadcast it. Multi-agent consumers should use
177    /// [`MultiAgentEnvironment::step_multi`].
178    fn step(&mut self, action: i64) -> StepResult {
179        let r = self.step_multi(&[vec![action], vec![action]]);
180        StepResult {
181            observation: r.observations.into_iter().next().unwrap_or_default(),
182            reward: r.rewards.first().copied().unwrap_or(0.0),
183            terminated: r.terminated.first().copied().unwrap_or(false),
184            truncated: false,
185            info: StepInfo::default(),
186        }
187    }
188
189    fn observation_space(&self) -> SpaceInfo {
190        SpaceInfo { shape: vec![self.vocab], space_type: SpaceType::Box }
191    }
192
193    fn action_space(&self) -> SpaceInfo {
194        SpaceInfo { shape: vec![], space_type: SpaceType::Discrete(self.vocab) }
195    }
196
197    fn render(&self) -> Vec<u8> {
198        Vec::new()
199    }
200
201    fn close(&mut self) {}
202
203    fn clone_state(&self) -> (i64, i64) {
204        (self.hidden, self.received)
205    }
206
207    fn restore_state(&mut self, state: &(i64, i64)) {
208        self.hidden = state.0;
209        self.received = state.1;
210    }
211}
212
213impl MultiAgentEnvironment for SignalingGame {
214    fn num_agents(&self) -> usize {
215        Self::NUM_AGENTS
216    }
217
218    fn get_agent_observation(&self, agent_id: usize) -> Vec<f32> {
219        match agent_id {
220            SPEAKER => self.speaker_obs(),
221            LISTENER => self.listener_obs(),
222            other => panic!("SignalingGame has 2 agents; no agent {other}"),
223        }
224    }
225
226    fn agent_action_space(&self, agent_id: usize) -> Vec<usize> {
227        match agent_id {
228            // Speaker has no task action — only the message dim (see `message_vocab`).
229            SPEAKER => Vec::new(),
230            // Listener's task action is its guess over the vocabulary.
231            LISTENER => vec![self.vocab],
232            other => panic!("SignalingGame has 2 agents; no agent {other}"),
233        }
234    }
235
236    fn step_multi(&mut self, actions: &[Vec<i64>]) -> MultiAgentResult {
237        assert_eq!(actions.len(), Self::NUM_AGENTS, "signaling game requires 2 action vectors");
238
239        // Peel the message dims off each agent's action using the shared helper.
240        // Speaker: task_len 0 -> the whole action is its message token.
241        let (_speaker_task, speaker_msg) =
242            split_action(&actions[SPEAKER], self.agent_action_space(SPEAKER).len());
243        // Listener: task_len == vocab dims -> its guess; it sends no message.
244        let (listener_task, _listener_msg) =
245            split_action(&actions[LISTENER], self.agent_action_space(LISTENER).len());
246
247        // Broadcast delivery: the speaker's message reaches the listener.
248        debug_assert!(self.delivery().reaches(SPEAKER, LISTENER));
249        self.received = speaker_msg.first().copied().unwrap_or(-1);
250
251        // Reward: both agents win iff the listener's guess matches the hidden token.
252        let guess = listener_task.first().copied().unwrap_or(-1);
253        let correct = guess == self.hidden;
254        let reward = if correct { 1.0_f32 } else { 0.0 };
255
256        MultiAgentResult::new(
257            vec![self.speaker_obs(), self.listener_obs()],
258            vec![reward, reward],
259            // Single-shot game: the round terminates every step.
260            vec![true, true],
261            vec![false, false],
262        )
263    }
264
265    fn active_agents(&self) -> Vec<bool> {
266        vec![true; Self::NUM_AGENTS]
267    }
268}
269
270/// Thin adapter so the joint PPO trainer can drive the signaling game
271/// directly (issue #275, Phase 2). `reset_joint` / `step_joint` delegate to
272/// the existing [`MultiAgentEnvironment`] surface — all message routing is
273/// already handled inside [`SignalingGame::step_multi`], so no extra plumbing
274/// lives here.
275impl JointEnv for SignalingGame {
276    fn reset_joint(&mut self, _seed: Option<u64>) -> Vec<Vec<f32>> {
277        // The env consumes no RNG (the hidden token is fixed at construction),
278        // so the seed is intentionally ignored.
279        self.reset();
280        (0..self.num_agents()).map(|i| self.get_agent_observation(i)).collect()
281    }
282
283    fn step_joint(&mut self, actions: &[Vec<i64>]) -> JointStepResult {
284        let result = self.step_multi(actions);
285        JointStepResult {
286            rewards: result.rewards,
287            // Single-shot game: every step terminates the round.
288            done: result.terminated.iter().all(|&t| t),
289            observations: result.observations,
290        }
291    }
292}
293
294impl CommunicatingEnvironment for SignalingGame {
295    fn message_vocab(&self, agent_id: usize) -> Vec<usize> {
296        match agent_id {
297            // Speaker emits one token over the full vocabulary.
298            SPEAKER => vec![self.vocab],
299            // Listener sends nothing.
300            _ => Vec::new(),
301        }
302    }
303
304    fn message_obs_size(&self, agent_id: usize) -> usize {
305        match agent_id {
306            // Listener receives the speaker's single message token.
307            LISTENER => 1,
308            // Speaker receives nothing.
309            _ => 0,
310        }
311    }
312
313    fn delivery(&self) -> Delivery {
314        Delivery::Broadcast
315    }
316}
317
318#[cfg(test)]
319mod tests {
320    use super::*;
321
322    #[test]
323    fn comms_layout_matches_design() {
324        let env = SignalingGame::new(4);
325        // Speaker: no task action, one 4-way message token, receives nothing.
326        assert_eq!(env.agent_action_space(SPEAKER), Vec::<usize>::new());
327        assert_eq!(env.message_vocab(SPEAKER), vec![4]);
328        assert_eq!(env.message_obs_size(SPEAKER), 0);
329        // Listener: a 4-way guess, sends nothing, receives one token dim.
330        assert_eq!(env.agent_action_space(LISTENER), vec![4]);
331        assert_eq!(env.message_vocab(LISTENER), Vec::<usize>::new());
332        assert_eq!(env.message_obs_size(LISTENER), 1);
333        assert_eq!(env.delivery(), Delivery::Broadcast);
334    }
335
336    #[test]
337    fn speaker_observes_hidden_token_one_hot() {
338        let env = SignalingGame::with_hidden(5, 3);
339        let obs = env.get_agent_observation(SPEAKER);
340        assert_eq!(obs, vec![0.0, 0.0, 0.0, 1.0, 0.0]);
341    }
342
343    #[test]
344    fn message_round_trips_speaker_to_listener() {
345        // Core acceptance criterion: agent 0 emits token T; after `step_multi`
346        // agent 1 receives T in its observation vector.
347        let mut env = SignalingGame::with_hidden(4, 2);
348        let token = 2_i64;
349        // Speaker action = [token]; listener guesses (value irrelevant here).
350        let result = env.step_multi(&[vec![token], vec![0]]);
351
352        // Listener's post-step observation carries the emitted token.
353        let listener_obs = &result.observations[LISTENER];
354        assert_eq!(listener_obs.len(), env.message_obs_size(LISTENER));
355        assert_eq!(listener_obs[0], token as f32);
356        // And it is reflected by the standalone observation accessor too.
357        assert_eq!(env.get_agent_observation(LISTENER)[0], token as f32);
358        assert_eq!(env.received(), token);
359    }
360
361    #[test]
362    fn reward_is_one_when_listener_guesses_hidden_token() {
363        let mut env = SignalingGame::with_hidden(3, 1);
364        // Listener guesses correctly (1 == hidden).
365        let correct = env.step_multi(&[vec![1], vec![1]]);
366        assert_eq!(correct.rewards, vec![1.0, 1.0]);
367
368        // Listener guesses incorrectly.
369        let mut env = SignalingGame::with_hidden(3, 1);
370        let wrong = env.step_multi(&[vec![1], vec![0]]);
371        assert_eq!(wrong.rewards, vec![0.0, 0.0]);
372    }
373
374    #[test]
375    fn listener_obs_before_step_is_sentinel() {
376        let env = SignalingGame::new(3);
377        // No step yet -> received == -1 sentinel.
378        assert_eq!(env.get_agent_observation(LISTENER), vec![-1.0]);
379    }
380
381    #[test]
382    fn reset_clears_received_message() {
383        let mut env = SignalingGame::with_hidden(4, 0);
384        env.step_multi(&[vec![3], vec![0]]);
385        assert_eq!(env.received(), 3);
386        env.reset();
387        assert_eq!(env.received(), -1);
388        assert_eq!(env.get_agent_observation(LISTENER), vec![-1.0]);
389    }
390
391    #[test]
392    fn snapshot_restore_round_trips() {
393        let mut env = SignalingGame::with_hidden(4, 2);
394        env.step_multi(&[vec![3], vec![0]]);
395        let snap = env.clone_state();
396        env.step_multi(&[vec![1], vec![0]]);
397        assert_eq!(env.received(), 1);
398        env.restore_state(&snap);
399        assert_eq!(env.hidden(), 2);
400        assert_eq!(env.received(), 3);
401    }
402
403    #[test]
404    fn num_agents_and_active() {
405        let env = SignalingGame::default();
406        assert_eq!(env.num_agents(), 2);
407        assert_eq!(env.active_agents(), vec![true, true]);
408    }
409}
thrust_rl/env/games/signaling.rs

thrust_rl/env/games/
signaling.rs
