⚡ Use compact scalar RL observations

CHANGELOG.md

@@ -12,6 +12,8 @@ releases may include breaking changes.
 
 ### Added
 
+- ✨ Encode the RL qubit-count and depth observations as normalized one-element
+  `float32` arrays ([#784]) ([**@flowerthrower**])
 - ✨ Expand the RL observation with normalized OpenQASM operation frequencies
   and include measurements in the shared ML feature schema ([#758])
   ([**@flowerthrower**])
@@ -92,6 +94,7 @@ for previous changelogs._
 
 [#773]: https://github.com/munich-quantum-toolkit/predictor/pull/771
 [#769]: https://github.com/munich-quantum-toolkit/predictor/pull/769
+[#784]: https://github.com/munich-quantum-toolkit/predictor/pull/784
 [#758]: https://github.com/munich-quantum-toolkit/predictor/pull/758
 [#755]: https://github.com/munich-quantum-toolkit/predictor/pull/755
 [#731]: https://github.com/munich-quantum-toolkit/predictor/pull/731

UPGRADING.md

@@ -6,6 +6,14 @@ of changes including minor and patch releases, please refer to the
 
 ## [Unreleased]
 
+### Compact scalar RL observations
+
+The `num_qubits` and `depth` entries in RL observations are now one-element
+`float32` arrays in `[0, 1]` instead of discrete integers. The qubit count is
+linearly scaled by the target device's qubit count; the depth is `log1p`-scaled
+and capped at 999,999. Existing RL models must be retrained, and code that
+consumes `PredictorEnv` observations directly must handle the new array values.
+
 ### RL operation features
 
 The RL observation now includes normalized frequencies for supported OpenQASM

src/mqt/predictor/rl/helper.py

@@ -11,6 +11,7 @@
 from __future__ import annotations
 
 import logging
+from math import log1p
 from pathlib import Path
 from typing import TYPE_CHECKING
 
@@ -28,6 +29,7 @@
 
 logger = logging.getLogger("mqt-predictor")
 
+MAX_CIRCUIT_DEPTH = 999_999
 OBSERVATION_OPERATIONS = tuple(get_openqasm_gates())
 
 
@@ -72,20 +74,22 @@ def get_state_sample(max_qubits: int, path_training_circuits: Path, rng: Generat
     return qc, str(file_list[random_index])
 
 
-def create_feature_dict(qc: QuantumCircuit) -> dict[str, int | NDArray[np.float32]]:
+def create_feature_dict(qc: QuantumCircuit, max_num_qubits: int) -> dict[str, NDArray[np.float32]]:
     """Create a normalized feature dictionary for a quantum circuit."""
     operation_counts = dict(qc.count_ops())
     total_operations = sum(value for gate, value in operation_counts.items() if gate != "barrier")
-    feature_dict: dict[str, int | NDArray[np.float32]] = {
+    normalized_num_qubits = min(qc.num_qubits, max_num_qubits) / max_num_qubits
+    normalized_depth = log1p(min(qc.depth(), MAX_CIRCUIT_DEPTH)) / log1p(MAX_CIRCUIT_DEPTH)
+    feature_dict = {
         **{
             operation: np.array(
                 [operation_counts.get(operation, 0) / total_operations if total_operations else 0.0],
                 dtype=np.float32,
             )
             for operation in OBSERVATION_OPERATIONS
         },
-        "num_qubits": qc.num_qubits,
-        "depth": qc.depth(),
+        "num_qubits": np.array([normalized_num_qubits], dtype=np.float32),
+        "depth": np.array([normalized_depth], dtype=np.float32),
     }
 
     supermarq_features = calc_supermarq_features(qc)

src/mqt/predictor/rl/predictorenv.py

@@ -183,8 +183,8 @@ def __init__(
             operation: Box(low=0, high=1, shape=(1,), dtype=np.float32) for operation in OBSERVATION_OPERATIONS
         }
         spaces: dict[str, Space] = {
-            "num_qubits": Discrete(128),
-            "depth": Discrete(1000000),
+            "num_qubits": Box(low=0, high=1, shape=(1,), dtype=np.float32),
+            "depth": Box(low=0, high=1, shape=(1,), dtype=np.float32),
             **operation_spaces,
             "program_communication": Box(low=0, high=1, shape=(1,), dtype=np.float32),
             "critical_depth": Box(low=0, high=1, shape=(1,), dtype=np.float32),
@@ -297,7 +297,7 @@ def step(self, action: int) -> tuple[dict[str, Any], float, bool, bool, dict[Any
             action_duration = time.perf_counter() - start_time
             # Different passes may fail for various reasons (e.g., found no routing solution).
             self.error_occurred = True
-            obs = create_feature_dict(self.state)
+            obs = create_feature_dict(self.state, self.device.num_qubits)
 
             # Trace the error before aborting
             self._collect_tracer_data(
@@ -337,7 +337,7 @@ def step(self, action: int) -> tuple[dict[str, Any], float, bool, bool, dict[Any
             reward_val = 0
             done = False
 
-        obs = create_feature_dict(self.state)
+        obs = create_feature_dict(self.state, self.device.num_qubits)
 
         # Trace+truncate if step limit is reached
         if not done and self.max_steps is not None and self.num_steps >= self.max_steps:
@@ -447,7 +447,7 @@ def reset(
         self.num_qubits_uncompiled_circuit = self.state.num_qubits
         self.has_parameterized_gates = len(self.state.parameters) > 0
 
-        obs = create_feature_dict(self.state)
+        obs = create_feature_dict(self.state, self.device.num_qubits)
 
         # Setup tracer for the new episode
         if self.tracer_output_path is not None:

src/mqt/predictor/rl/tracer.py

@@ -209,7 +209,7 @@ def record_step(
         reward: float,
         current_qc: QuantumCircuit,
         figures_of_merit: FigureOfMeritMetrics,
-        features: dict[str, int | NDArray[np.float32]],
+        features: dict[str, NDArray[np.float32]],
         synthesized: bool,
         laid_out: bool,
         routed: bool,
@@ -337,8 +337,6 @@ def _extract_device_metadata(device: Target) -> DeviceMetadata:
         )
 
     @staticmethod
-    def _extract_float(val: int | NDArray[np.float32]) -> float:
-        """Extract a float from a scalar or one-element feature array."""
-        if isinstance(val, int):
-            return float(val)
+    def _extract_float(val: NDArray[np.float32]) -> float:
+        """Extract a float from a one-element feature array."""
         return float(val.item())

tests/compilation/test_helper_rl.py

@@ -45,22 +45,20 @@ def test_create_feature_dict() -> None:
     qc.ccx(0, 1, 2)
     qc.measure_all()
 
-    features = create_feature_dict(qc)
+    max_num_qubits = 20
+    features = create_feature_dict(qc, max_num_qubits)
 
-    for name, feature in features.items():
-        if name in {"num_qubits", "depth"}:
-            assert isinstance(feature, int)
-        else:
-            assert isinstance(feature, np.ndarray)
-            assert feature.dtype == np.float32
+    for feature in features.values():
+        assert isinstance(feature, np.ndarray)
+        assert feature.dtype == np.float32
 
     np.testing.assert_allclose(features["h"], [1 / 6])
     np.testing.assert_allclose(features["cx"], [1 / 6])
     np.testing.assert_allclose(features["ccx"], [1 / 6])
     np.testing.assert_allclose(features["measure"], [0.5])
     np.testing.assert_allclose(features["x"], [0])
-    assert features["num_qubits"] == qc.num_qubits
-    assert features["depth"] == qc.depth()
+    np.testing.assert_allclose(features["num_qubits"], [3 / max_num_qubits])
+    np.testing.assert_allclose(features["depth"], [np.log1p(qc.depth()) / np.log1p(999_999)])
 
 
 def test_get_path_trained_model() -> None:

tests/compilation/test_predictor_rl.py

@@ -50,8 +50,8 @@ def test_predictor_env_reset_from_string() -> None:
         dump(qc, f)
     observation, _ = predictor.env.reset(qc=qasm_path)
 
-    assert observation == create_feature_dict(qc)
-    assert observation["num_qubits"] == qc.num_qubits
+    assert observation == create_feature_dict(qc, device.num_qubits)
+    assert observation["num_qubits"][0] == pytest.approx(qc.num_qubits / device.num_qubits)
 
 
 def test_predictor_env_esp_error() -> None:

tests/compilation/test_tracer.py

@@ -77,8 +77,8 @@ def test_compilation_tracer_generates_valid_json(tmp_path: Path) -> None:
     qc_terminal.rz(0.5, 2)
 
     # 2. Mock the feature dictionaries that the RL environment normally passes
-    features_baseline = create_feature_dict(qc_baseline)
-    features_terminal = create_feature_dict(qc_terminal)
+    features_baseline = create_feature_dict(qc_baseline, device.num_qubits)
+    features_terminal = create_feature_dict(qc_terminal, device.num_qubits)
 
     # 3. Use the actual `record_step` API instead of manually building CompilationSteps
     tracer.record_step(