core_rust/core/utils/
matrices.rs

1use crate::core::game::*;
2use game_logic::Game;
3use state::State;
4use utils::actions::Action;
5
6/// Build a dense 3D indicator tensor T[s][a][s'] == 1 for legal transitions.
7#[allow(dead_code)]
8fn build_state_action_transition_matrix(
9    game: &Game,
10    all_states: &[State],
11    all_actions: &[Action],
12) -> Vec<Vec<Vec<usize>>> {
13    let num_states = all_states.len();
14    let num_actions = all_actions.len();
15
16    // precompute a fast lookup from unit_position → state_idx
17    let mut pos2idx = std::collections::HashMap::with_capacity(num_states);
18    for (i, st) in all_states.iter().enumerate() {
19        pos2idx.insert(st.unit_position, i);
20    }
21
22    // allocate: [S][A][S] zeroed
23    let mut t = vec![vec![vec![0; num_states]; num_actions]; num_states];
24
25    for (s_idx, state) in all_states.iter().enumerate() {
26        for (a_idx, &action) in all_actions.iter().enumerate() {
27            let (next_state, _) = game.simulate(state, &action).expect("simulation failed");
28
29            if let Some(&ns_idx) = pos2idx.get(&next_state.unit_position) {
30                t[s_idx][a_idx][ns_idx] = 1;
31            }
32        }
33    }
34
35    t
36}
37
38/// Build a dense reward matrix R[s][a] for one-step rewards.
39#[allow(dead_code)]
40fn build_state_action_reward_matrix(
41    game: &Game,
42    all_states: &[State],
43    all_actions: &[Action],
44) -> Vec<Vec<f32>> {
45    let num_states = all_states.len();
46    let num_actions = all_actions.len();
47
48    // allocate: [S][A] zeroed
49    let mut r = vec![vec![0.0; num_actions]; num_states];
50
51    for (s_idx, state) in all_states.iter().enumerate() {
52        for (a_idx, &action) in all_actions.iter().enumerate() {
53            let (_next_state, game_vars) =
54                game.simulate(state, &action).expect("simulation failed");
55            r[s_idx][a_idx] = game_vars.score;
56        }
57    }
58
59    r
60}
61
62/// Build transition and reward matrices for the provided states/actions.
63#[allow(dead_code)]
64pub fn build_matrices(
65    game: &Game,
66    all_states: &[State],
67    all_actions: &[Action],
68) -> (Vec<Vec<Vec<usize>>>, Vec<Vec<f32>>) {
69    let transition_matrix = build_state_action_transition_matrix(game, all_states, all_actions);
70    let reward_matrix = build_state_action_reward_matrix(game, all_states, all_actions);
71
72    (transition_matrix, reward_matrix)
73}
core_rust/core/utils/matrices.rs

core_rust/core/utils/
matrices.rs
