Add configurable observation encoder and tests

examples/grid_world/src/lib.rs

@@ -5,6 +5,7 @@ use rand::thread_rng;
 use std::cmp;
 
 use twisterl::rl::env::Env;
+use twisterl::rl::codec::{ObservationCodec, build_codec};
 use twisterl::python_interface::env::{PyBaseEnv, get_env_ref, get_env_mut};
 
 #[derive(Clone)]
@@ -17,6 +18,7 @@ pub struct GridWorld {
     agent: (usize, usize),
     goal: (usize, usize),
     trap: (usize, usize),
+    codec: std::sync::Arc<dyn ObservationCodec>,
 }
 
 impl GridWorld {
@@ -34,7 +36,9 @@ impl GridWorld {
             steps_left: max_steps,
             agent: (0, 0),
             goal: (0, 0),
-            trap: (0,0)
+            trap: (0,0),
+            codec: build_codec("multi_hot", width * height, width * height)
+                .expect("Failed to create grid world observation codec"),
         };
         // env.reset();
         env
@@ -155,7 +159,8 @@ impl Env for GridWorld {
     }
 
     fn observe(&self) -> Vec<usize> {
-        self.get_state().iter().enumerate().map(|(i, v)| i * self.height * self.width + v).collect()  
+        let state = self.get_state();
+        self.codec.encode_state(&state)
     }
 }
 
@@ -204,4 +209,4 @@ impl PyGridWorldEnv {
 fn grid_world(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<PyGridWorldEnv>()?;
     Ok(())
-}
+}

rust/src/envs/puzzle.rs

@@ -13,6 +13,8 @@ that they have been altered from the originals.
 
 use rand::distributions::{Distribution, Uniform};
 use crate::rl::env::Env;
+use crate::rl::codec::{ObservationCodec, build_codec};
+use std::sync::Arc;
 
 
 // This is the Env definition
@@ -27,6 +29,7 @@ pub struct Puzzle {
     pub difficulty: usize,
     pub depth_slope: usize,
     pub max_depth: usize,
+    codec: Arc<dyn ObservationCodec>,
 }
 
 
@@ -38,7 +41,20 @@ impl Puzzle {
         depth_slope: usize,
         max_depth: usize,
     ) -> Self {
-        Puzzle {state: (0..(width*height)).collect(), zero_location: (0,0), depth:1, width, height, difficulty, depth_slope, max_depth}
+        let num_slots = width * height;
+        let codec = build_codec("multi_hot", num_slots, num_slots)
+            .expect("Failed to create puzzle observation codec");
+        Puzzle {
+            state: (0..(width*height)).collect(),
+            zero_location: (0,0),
+            depth:1,
+            width,
+            height,
+            difficulty,
+            depth_slope,
+            max_depth,
+            codec,
+        }
     }
 
     pub fn solved(&self) -> bool {
@@ -176,8 +192,8 @@ impl Env for Puzzle {
         }
     }
 
-fn observe(&self,) -> Vec<usize> {
-        self.state.iter().enumerate().map(|(i, v)| i * self.height * self.width + v).collect()
+    fn observe(&self,) -> Vec<usize> {
+        self.codec.encode_state(&self.state)
     }
 
 }

rust/src/python_interface/codec.rs

@@ -0,0 +1,58 @@
+// -*- coding: utf-8 -*-
+/*
+(C) Copyright 2025 IBM. All Rights Reserved.
+
+This code is licensed under the Apache License, Version 2.0. You may
+obtain a copy of this license in the LICENSE.txt file in the root directory
+of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+
+Any modifications or derivative works of this code must retain this
+copyright notice, and modified files need to carry a notice indicating
+that they have been altered from the originals.
+*/
+
+use std::sync::Arc;
+
+use pyo3::prelude::*;
+
+use crate::rl::codec::{build_codec, ObservationCodec};
+use crate::python_interface::error_mapping::MyError;
+
+#[pyclass(name = "ObservationCodec")]
+pub struct PyObservationCodec {
+    codec: Arc<dyn ObservationCodec>,
+}
+
+impl Clone for PyObservationCodec {
+    fn clone(&self) -> Self {
+        Self {
+            codec: Arc::clone(&self.codec),
+        }
+    }
+}
+
+#[pymethods]
+impl PyObservationCodec {
+    pub fn encode(&self, obs: Vec<Vec<usize>>) -> Vec<Vec<f32>> {
+        self.codec.encode_indices(&obs)
+    }
+
+    #[getter]
+    pub fn obs_size(&self) -> usize {
+        self.codec.obs_size()
+    }
+
+    pub fn clone(&self) -> Self {
+        self.clone()
+    }
+}
+
+#[pyfunction]
+pub fn make_observation_codec(
+    kind: &str,
+    num_slots: usize,
+    domain_size: usize,
+) -> PyResult<PyObservationCodec> {
+    let codec = build_codec(kind, num_slots, domain_size).map_err(MyError::from)?;
+    Ok(PyObservationCodec { codec })
+}

rust/src/python_interface/mod.rs

@@ -19,3 +19,4 @@ pub mod modules;
 pub mod policy;
 pub mod pyenv;
 pub mod python_bindings;
+pub mod codec;

rust/src/python_interface/python_bindings.rs

@@ -17,6 +17,7 @@ use pyo3::prelude::*;
 use crate::python_interface::modules::PySequential;
 use crate::python_interface::layers::{PyEmbeddingBag, PyLinear};
 use crate::python_interface::policy::PyPolicy;
+use crate::python_interface::codec::{PyObservationCodec, make_observation_codec};
 
 fn init_nn_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<PyEmbeddingBag>()?;
@@ -52,6 +53,12 @@ fn init_collector_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     Ok(())
 }
 
+fn init_codec_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_class::<PyObservationCodec>()?;
+    m.add_function(wrap_pyfunction!(make_observation_codec, m)?)?;
+    Ok(())
+}
+
 fn add_twisterl_functionality(py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
     let nn_submod = PyModule::new(py, "nn")?;
     init_nn_module(&nn_submod)?;
@@ -65,6 +72,10 @@ fn add_twisterl_functionality(py: Python, m: &Bound<'_, PyModule>) -> PyResult<(
     init_collector_module(&collector_submod)?;
     m.add_submodule(&collector_submod)?;
 
+    let codec_submod = PyModule::new(py, "codec")?;
+    init_codec_module(&codec_submod)?;
+    m.add_submodule(&codec_submod)?;
+
     Ok(())
 }
 

rust/src/rl/codec.rs

@@ -0,0 +1,85 @@
+// -*- coding: utf-8 -*-
+/*
+(C) Copyright 2025 IBM. All Rights Reserved.
+
+This code is licensed under the Apache License, Version 2.0. You may
+obtain a copy of this license in the LICENSE.txt file in the root directory
+of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+
+Any modifications or derivative works of this code must retain this
+copyright notice, and modified files need to carry a notice indicating
+that they have been altered from the originals.
+*/
+
+use std::sync::Arc;
+
+/// Trait implemented by observation codecs that transform raw environment
+/// representations into sparse indices and dense vectors.
+pub trait ObservationCodec: Send + Sync {
+    fn obs_size(&self) -> usize;
+    fn encode_state(&self, state: &[usize]) -> Vec<usize>;
+    fn encode_indices(&self, obs: &[Vec<usize>]) -> Vec<Vec<f32>>;
+}
+
+/// Multi-hot codec that expands categorical slots into multi-hot vectors.
+pub struct MultiHotCodec {
+    num_slots: usize,
+    domain_size: usize,
+    obs_size: usize,
+}
+
+impl MultiHotCodec {
+    pub fn new(num_slots: usize, domain_size: usize) -> Self {
+        Self {
+            num_slots,
+            domain_size,
+            obs_size: num_slots * domain_size,
+        }
+    }
+}
+
+impl ObservationCodec for MultiHotCodec {
+    fn obs_size(&self) -> usize {
+        self.obs_size
+    }
+
+    fn encode_state(&self, state: &[usize]) -> Vec<usize> {
+        state
+            .iter()
+            .enumerate()
+            .map(|(slot, &value)| {
+                let v = if value < self.domain_size {
+                    value
+                } else {
+                    value % self.domain_size
+                };
+                slot * self.domain_size + v
+            })
+            .collect()
+    }
+
+    fn encode_indices(&self, obs: &[Vec<usize>]) -> Vec<Vec<f32>> {
+        let mut result = Vec::with_capacity(obs.len());
+        for sample in obs {
+            let mut dense = vec![0.0f32; self.obs_size];
+            for &idx in sample {
+                if idx < self.obs_size {
+                    dense[idx] = 1.0;
+                }
+            }
+            result.push(dense);
+        }
+        result
+    }
+}
+
+pub fn build_codec(
+    kind: &str,
+    num_slots: usize,
+    domain_size: usize,
+) -> anyhow::Result<Arc<dyn ObservationCodec>> {
+    match kind {
+        "multi_hot" => Ok(Arc::new(MultiHotCodec::new(num_slots, domain_size))),
+        _ => Err(anyhow::anyhow!("Unknown observation codec type: {}", kind)),
+    }
+}

rust/src/rl/mod.rs

@@ -15,4 +15,5 @@ pub mod search;
 pub mod evaluate;
 pub mod tree;
 pub mod env;
-pub mod solve;
+pub mod solve;
+pub mod codec;

src/twisterl/defaults.py

@@ -120,6 +120,7 @@ def make_config(algo_name, input_config):
         "evals": copy.deepcopy(EVALS_CONFIG),
         "learning": copy.deepcopy(LEARNING_CONFIG),
         "logging": copy.deepcopy(LOGGING_CONFIG),
+        "observation_encoder": {"type": "multi_hot"},
         **copy.deepcopy(ALGO_CONFIG[algo_name]),
     }
     conf.update(input_config)

src/twisterl/rl/algorithm.py

@@ -20,6 +20,7 @@
 from torch.utils.tensorboard import SummaryWriter
 
 from twisterl.defaults import make_config
+from twisterl.rl.observation import make_observation_encoder
 
 from twisterl import twisterl
 
@@ -60,11 +61,23 @@ def __init__(self, env, policy, config, run_path=None):
         self.policy.device = self.config["device"]
         self.rs_pol = self.policy.to_rust()
 
+        # Observation encoder handles sparse/dense conversions
+        self.obs_shape = getattr(self.policy, "obs_shape", None)
+        if self.obs_shape is None:
+            raise ValueError("Policy must expose obs_shape for observation encoding.")
+        self.obs_encoder = make_observation_encoder(
+            self.obs_shape, self.config.get("observation_encoder")
+        )
+
         # Make optimizer
         self.optimizer = torch.optim.Adam(
             self.policy.parameters(), **self.config["optimizer"]
         )
 
+    def encode_obs(self, obs):
+        """Returns numpy array representation for the collected observations."""
+        return self.obs_encoder(obs)
+
     @timed
     @abstractmethod
     def data_to_torch(self, data):

src/twisterl/rl/az.py

@@ -11,7 +11,6 @@
 # that they have been altered from the originals.
 
 import torch
-import numpy as np
 
 from twisterl.rl.algorithm import Algorithm, timed
 from twisterl import twisterl
@@ -30,11 +29,11 @@ def data_to_torch(self, data):
             data.additional_data["remaining_values"],
         )
 
-        np_obs = np.zeros((len(obs), self.obs_size), dtype=float)
-        for i, obs_i in enumerate(obs):
-            np_obs[i, obs_i] = 1.0
+        encoded_obs = self.encode_obs(obs)
 
-        pt_obs = torch.tensor(np_obs, dtype=torch.float, device=self.config["device"])
+        pt_obs = torch.tensor(
+            encoded_obs, dtype=torch.float, device=self.config["device"]
+        )
         pt_probs = torch.tensor(probs, dtype=torch.float, device=self.config["device"])
         pt_vals = torch.tensor(
             vals, dtype=torch.float, device=self.config["device"]