# Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. """ Quantum Circuit Optimization Environment Implementation. A noise-aware, hardware-constrained reinforcement learning environment for quantum circuit design and optimization. Physics backend: Qiskit Statevector (stateless — no Aer, no noise models). RL logic: Pure Python / NumPy (unchanged from original design). Design principle: Environment = decision + RL logic Qiskit = stateless physics engine """ # --------------------------------------------------------------------------- # sys.path bootstrap -- makes `python server/my_env_environment.py` work # when run directly from the my_env/ root directory. # --------------------------------------------------------------------------- import os import sys _HERE = os.path.dirname(os.path.abspath(__file__)) # .../my_env/server/ _ROOT = os.path.dirname(_HERE) # .../my_env/ for _p in (_ROOT, _HERE): if _p not in sys.path: sys.path.insert(0, _p) import logging import math from typing import Any, Dict, List, Optional, Tuple from uuid import uuid4 logger = logging.getLogger(__name__) import numpy as np # --------------------------------------------------------------------------- # STEP 1: Qiskit lightweight imports (statevector only — NO Aer) # --------------------------------------------------------------------------- from qiskit import QuantumCircuit from qiskit.quantum_info import Statevector from openenv.core.env_server.interfaces import Environment # Try relative imports first (package mode -- same class identity as app.py). # Fall back to absolute imports (direct execution mode). try: from ..models import ActionType, GateType, QuantumAction, QuantumObservation, QuantumState from ..graders import ( AggregateGrader, ConstraintsGrader, EfficiencyGrader, FidelityGrader, NoiseGrader, UnitaryGrader, ) except ImportError: from models import ActionType, GateType, QuantumAction, QuantumObservation, QuantumState from graders import ( AggregateGrader, ConstraintsGrader, EfficiencyGrader, FidelityGrader, NoiseGrader, UnitaryGrader, ) try: from .tasks import TASK_REGISTRY except ImportError: from tasks import TASK_REGISTRY # --------------------------------------------------------------------------- # STEP 8: Performance constraints (critical) # --------------------------------------------------------------------------- MAX_QUBITS: int = 4 # Maximum qubits supported MAX_STEPS: int = 15 # Soft cap per spec; per-task configs retain their limits MAX_DEPTH: int = 20 # Maximum circuit depth before penalty # --------------------------------------------------------------------------- # STEP 2: Circuit converter (Qiskit — stateless) # --------------------------------------------------------------------------- def build_qiskit_circuit(circuit_gates: List[Dict[str, Any]], num_qubits: int) -> QuantumCircuit: """ Convert a list of gate dicts into a Qiskit QuantumCircuit. Supports: H, X, CNOT, SWAP, RX, RZ. All other gate names are silently skipped. Args: circuit_gates: List of dicts with keys ``gate``, ``qubits``, ``parameter``. num_qubits: Number of qubits in the circuit. Returns: A Qiskit :class:`QuantumCircuit` with the gates applied in order. """ qc = QuantumCircuit(num_qubits) for g in circuit_gates: name = g["gate"] q = g["qubits"] p = g.get("parameter") if name == "H": qc.h(q[0]) elif name == "X": qc.x(q[0]) elif name == "CNOT": qc.cx(q[0], q[1]) elif name == "SWAP": # STEP 4: Real physics SWAP — used for qubit routing on # limited-connectivity hardware (e.g. q0—q1—q2 topology). qc.swap(q[0], q[1]) elif name == "RX": qc.rx(float(p) if p is not None else 0.0, q[0]) elif name == "RZ": qc.rz(float(p) if p is not None else 0.0, q[0]) # Unknown gates are intentionally skipped (graceful degradation) return qc # --------------------------------------------------------------------------- # STEP 3: Statevector computation (replaces all old NumPy simulator logic) # --------------------------------------------------------------------------- def compute_statevector(circuit_gates: List[Dict[str, Any]], num_qubits: int) -> np.ndarray: """ Compute the statevector of a circuit using Qiskit. Starts from the |00...0> ground state (Qiskit default). Does NOT store any Qiskit objects — returns a plain NumPy complex128 array. Args: circuit_gates: Gate list as produced by the environment's _gates field. num_qubits: Number of qubits. Returns: Complex NumPy array of shape (2**num_qubits,). """ if not circuit_gates: # Empty circuit → |00...0> sv = np.zeros(2 ** num_qubits, dtype=np.complex128) sv[0] = 1.0 return sv qc = build_qiskit_circuit(circuit_gates, num_qubits) sv = np.asarray(Statevector.from_instruction(qc).data, dtype=np.complex128) # Normalize explicitly to handle Qiskit precision quirks norm = np.linalg.norm(sv) if norm > 0: sv = sv / norm # STEP 9: Validation assert len(sv) == 2 ** num_qubits, ( f"Statevector length {len(sv)} != 2^{num_qubits}={2**num_qubits}" ) return sv # --------------------------------------------------------------------------- # Circuit Depth Calculator (pure Python — no Qiskit dependency) # --------------------------------------------------------------------------- def compute_circuit_depth(gates: List[Dict[str, Any]], num_qubits: int) -> int: """ Compute circuit depth (longest path through any qubit). Each gate occupies one time-step on each of its qubits. Depth = max over all qubits of the number of gate layers touching that qubit. """ if not gates: return 0 qubit_depths = [0] * num_qubits for gate_info in gates: qubits = gate_info["qubits"] if len(qubits) == 1: qubit_depths[qubits[0]] += 1 elif len(qubits) >= 2: # Multi-qubit gate: all involved qubits advance to the max + 1 max_d = max(qubit_depths[q] for q in qubits) for q in qubits: qubit_depths[q] = max_d + 1 return max(qubit_depths) if qubit_depths else 0 # --------------------------------------------------------------------------- # Environment # --------------------------------------------------------------------------- class QuantumCircuitEnvironment(Environment): """ RL environment for noise-aware, hardware-constrained quantum circuit optimization. The agent builds a quantum circuit step by step to match a target quantum state while minimising depth, gate count, noise impact, and SWAP overhead. Physics backend: Qiskit Statevector (stateless — called per step, not stored). Supports concurrent WebSocket sessions. """ SUPPORTS_CONCURRENT_SESSIONS: bool = True def __init__(self, seed: Optional[int] = None): """ Initialise the environment. Args: seed: Optional RNG seed for reproducibility. """ self._seed = seed or 42 self._rng = np.random.RandomState(self._seed) # Task config -- set on reset() self._task_id: str = "easy" self._task_config: Dict[str, Any] = {} self._target_sv: Optional[np.ndarray] = None self._target_unitary: Optional[np.ndarray] = None # Circuit state (plain Python dicts — NO Qiskit objects stored) self._gates: List[Dict[str, Any]] = [] self._num_qubits: int = 2 # Episode tracking self._step_count: int = 0 self._max_steps: int = 20 self._done: bool = False self._prev_score: float = 0.0 self._prev_fidelity: float = 0.0 self._prev_depth: int = 0 self._prev_action: Optional[ActionType] = None self._last_reward: float = 0.0 self._episode_id: str = str(uuid4()) self._swap_count: int = 0 self._prev_penalty: float = 0.0 # Graders self._aggregate_grader = AggregateGrader() # ------------------------------------------------------------------ # Core API # ------------------------------------------------------------------ def reset(self, config: Optional[Dict] = None) -> QuantumObservation: """ Reset the environment for a new episode. Args: config: Optional configuration dictionary containing task_id. Returns: Initial QuantumObservation. """ # Resolve task_id from config task_id = config.get("task_id", "easy") if config else "easy" # Load task task_cls = TASK_REGISTRY.get(task_id, TASK_REGISTRY["easy"]) self._task_config = task_cls.config() self._task_id = self._task_config["task_id"] self._num_qubits = min(self._task_config["num_qubits"], MAX_QUBITS) self._max_steps = self._task_config["max_steps"] # Enforce MAX_QUBITS constraint if self._task_config["num_qubits"] > MAX_QUBITS: logger.warning( "Task '%s' requests %d qubits; clamped to MAX_QUBITS=%d.", self._task_id, self._task_config["num_qubits"], MAX_QUBITS, ) # Check if task uses statevector or unitary targets if hasattr(task_cls, "target_unitary"): self._target_unitary = task_cls.target_unitary() self._target_sv = None logger.info("Loaded task '%s' (%d qubits) with UNITARY target.", self._task_id, self._num_qubits) else: self._target_sv = task_cls.target_statevector() self._target_unitary = None logger.info("Loaded task '%s' (%d qubits) with STATEVECTOR target.", self._task_id, self._num_qubits) # Set grader weights for this task weights = self._task_config.get("grader_weights") if weights: self._aggregate_grader = AggregateGrader(weights=weights) # Reset circuit and episode state self._gates = [] self._step_count = 0 self._done = False self._prev_action = None self._last_reward = 0.0 self._episode_id = str(uuid4()) self._swap_count = 0 # Compute initial baseline score so the first step evaluates relatively initial_scores = self._compute_all_scores() self._prev_score = initial_scores["aggregate"] self._prev_fidelity = initial_scores["fidelity"] self._prev_depth = initial_scores["depth"] self._prev_penalty = ( 0.5 * (1.0 - initial_scores["efficiency"]) + 0.3 * (1.0 - initial_scores["noise"]) + 0.2 * (1.0 - initial_scores["constraints"]) ) return self._build_observation(reward=0.0) def step(self, action: QuantumAction) -> QuantumObservation: # type: ignore[override] """ Execute one step in the environment. STEP 6: Each step: 1. Applies action to circuit_gates 2. Computes statevector via Qiskit 3. Computes reward from state 4. Returns updated observation Args: action: QuantumAction specifying the gate operation. Returns: QuantumObservation with updated metrics and shaped reward. """ if self._done: return self._build_observation(reward=0.0) self._step_count += 1 # Handle STOP action if action.action_type == ActionType.STOP: self._done = True reward = self._compute_terminal_reward() if self._prev_score > 0.8: reward += 0.2 else: reward -= 0.2 # Clamp reward reward = max(-1.0, min(1.0, reward)) logger.info( "[STOP] task=%s qubits=%d step=%d reward=%.4f", self._task_id, self._num_qubits, self._step_count, reward, ) return self._build_observation(reward=reward) # Validate action valid, penalty_msg = self._validate_action(action) if not valid: self._check_done() logger.info("[INVALID] task=%s step=%d error=%s", self._task_id, self._step_count, penalty_msg) return self._build_observation(reward=-0.1, error=penalty_msg) # Capture pre-action state for REMOVE bonus/penalty logic pre_action_fidelity = self._prev_fidelity pre_action_depth = self._prev_depth # Apply action to circuit_gates # STEP 5 (REMOVE): pops last gate; state is recomputed from scratch by Qiskit self._apply_action(action) # Compute scores (triggers compute_statevector → Qiskit) scores = self._compute_all_scores() current_score = scores["aggregate"] current_fidelity = scores["fidelity"] current_depth = scores["depth"] # ------------------------------------------------------------------ # STEP 7: Reward function (physically meaningful) # ------------------------------------------------------------------ # --- Penalty computation --- current_penalty = ( 0.5 * (1.0 - scores["efficiency"]) + 0.3 * (1.0 - scores["noise"]) + 0.2 * (1.0 - scores["constraints"]) ) # Redundancy penalty if len(self._gates) >= 2: prev_gate = self._gates[-2] curr_gate = self._gates[-1] if prev_gate.get("gate") == curr_gate.get("gate") and prev_gate.get("qubits") == curr_gate.get("qubits"): current_penalty += 0.005 # PARAM penalty if action.action_type == ActionType.PARAM: current_penalty += 0.003 # SWAP penalty if action.action_type == ActionType.SWAP: current_penalty += 0.02 # --- Reward core --- fidelity_delta = current_fidelity - self._prev_fidelity penalty_delta = current_penalty - getattr(self, '_prev_penalty', 0.0) reward = fidelity_delta * 3.0 - penalty_delta * 0.5 # --- ADD vs REMOVE balance: reward constructive actions --- if action.action_type == ActionType.ADD and fidelity_delta < 0: reward += 0.15 # Increased from 0.1 to reduce ADD penalty # --- REMOVE logic: NOT profitable --- if action.action_type == ActionType.REMOVE: # Cap REMOVE reward - it can never profit more than half of last reward if self._last_reward != 0: reward = min(reward, abs(self._last_reward) * 0.5) # Small optimization bonus only if depth improved without fidelity loss if current_depth < pre_action_depth and abs(current_fidelity - pre_action_fidelity) < 1e-6: reward = min(reward + 0.005, 0.01) # Tiny bonus, capped at 0.01 elif current_fidelity < pre_action_fidelity: reward -= 0.05 # Stronger penalty for fidelity loss # --- Penalize ADD→REMOVE loops (direct exploit killer) --- if action.action_type == ActionType.REMOVE and self._prev_action == ActionType.ADD: reward -= 0.05 # --- Prevent REMOVE spam loops --- if self._prev_action == ActionType.REMOVE: reward -= 0.02 # --- Exploration bias --- reward += 0.02 # --- Overfitting penalty --- optimal_depth = self._task_config.get("optimal_depth", 5) if current_fidelity > 0.95 and current_depth > optimal_depth: reward -= 0.1 # --- Smooth --- reward = np.tanh(reward) # Update trackers self._prev_action = action.action_type self._last_reward = reward # Episode limit if self._step_count >= self._max_steps: self._done = True reward -= 0.1 reward = max(-1.0, min(1.0, reward)) # Update trackers self._prev_score = current_score self._prev_fidelity = current_fidelity self._prev_depth = current_depth self._prev_penalty = current_penalty logger.info( "[STEP] task=%s qubits=%d step=%d | " "fid=%.4f eff=%.4f noise=%.4f cstr=%.4f agg=%.4f | " "depth=%d gates=%d swaps=%d reward=%.4f", self._task_id, self._num_qubits, self._step_count, scores["fidelity"], scores["efficiency"], scores["noise"], scores["constraints"], scores["aggregate"], current_depth, scores["gate_count"], self._swap_count, reward, ) self._check_done() return self._build_observation(reward=reward) @property def state(self) -> QuantumState: """Return the full internal state.""" scores = self._compute_all_scores() return QuantumState( episode_id=self._episode_id, step_count=self._step_count, circuit_gates=list(self._gates), target_description=self._task_config.get("description", ""), task_id=self._task_id, max_steps=self._max_steps, noise_model_name=self._task_config.get("noise_model", "none"), current_fidelity=scores.get("fidelity", 0.0), current_score=scores.get("aggregate", 0.0), ) # ------------------------------------------------------------------ # Action validation # ------------------------------------------------------------------ def _validate_action(self, action: QuantumAction) -> Tuple[bool, Optional[str]]: """Validate that the action is legal.""" at = action.action_type if at == ActionType.ADD: if action.gate is None: return False, "ADD action requires a gate type" if not action.qubits: return False, "ADD action requires target qubits" for q in action.qubits: if q < 0 or q >= self._num_qubits: return False, f"Qubit index {q} out of range [0, {self._num_qubits - 1}]" if action.gate == GateType.CNOT and len(action.qubits) != 2: return False, "CNOT requires exactly 2 qubits" if action.gate in (GateType.H, GateType.X, GateType.RX, GateType.RZ): if len(action.qubits) != 1: return False, f"{action.gate.value} requires exactly 1 qubit" if action.gate in (GateType.RX, GateType.RZ) and action.parameter is None: return False, f"{action.gate.value} requires a parameter (angle)" elif at == ActionType.REMOVE: if not self._gates: return False, "Cannot REMOVE from an empty circuit" elif at == ActionType.SWAP: if len(action.qubits) != 2: return False, "SWAP requires exactly 2 qubits" for q in action.qubits: if q < 0 or q >= self._num_qubits: return False, f"Qubit index {q} out of range [0, {self._num_qubits - 1}]" if action.qubits[0] == action.qubits[1]: return False, "SWAP qubit indices must be distinct" elif at == ActionType.PARAM: if not self._gates: return False, "Cannot PARAM-tune with empty circuit" if action.parameter is None: return False, "PARAM action requires a parameter value" return True, None # ------------------------------------------------------------------ # Action application # ------------------------------------------------------------------ def _apply_action(self, action: QuantumAction) -> None: """ Apply a validated action to the circuit gate list. IMPORTANT: No Qiskit objects are stored — only plain Python dicts. State recomputation happens lazily in _compute_all_scores(). """ at = action.action_type if at == ActionType.ADD: gate_dict: Dict[str, Any] = { "gate": action.gate.value if action.gate else "H", "qubits": list(action.qubits), } if action.parameter is not None: gate_dict["parameter"] = action.parameter self._gates.append(gate_dict) elif at == ActionType.REMOVE: # STEP 5: Remove the LAST gate; state recomputed via Qiskit next step if self._gates: self._gates.pop() elif at == ActionType.SWAP: # STEP 4: SWAP appended as a gate dict; # build_qiskit_circuit translates this to qc.swap(q0, q1) self._gates.append({ "gate": "SWAP", "qubits": list(action.qubits), }) self._swap_count += 1 elif at == ActionType.PARAM: # Tune the parameter of the last parametric gate for i in range(len(self._gates) - 1, -1, -1): if self._gates[i]["gate"] in ("RX", "RZ"): self._gates[i]["parameter"] = action.parameter break # ------------------------------------------------------------------ # Scoring # ------------------------------------------------------------------ def _compute_all_scores(self) -> Dict[str, float]: """ Compute all grader scores for the current circuit. Calls compute_statevector() → Qiskit for the statevector path. UnitaryGrader uses its own matrix math (unchanged). """ # --- Fidelity --- if self._target_unitary is not None: fidelity_score = UnitaryGrader.grade(self._gates, self._target_unitary, self._num_qubits) # Use ground state as placeholder for metadata current_sv = np.zeros(2 ** self._num_qubits, dtype=np.complex128) current_sv[0] = 1.0 else: # STEP 3: Qiskit statevector computation current_sv = compute_statevector(self._gates, self._num_qubits) # Dimension assertion (also inside compute_statevector, but belt-and-suspenders) expected_dim = 2 ** self._num_qubits assert current_sv.shape[0] == expected_dim, ( f"Statevector dim {current_sv.shape[0]} != 2^{self._num_qubits}={expected_dim}" ) if self._target_sv is not None: assert self._target_sv.shape[0] == expected_dim, ( f"Target dim {self._target_sv.shape[0]} != 2^{self._num_qubits}={expected_dim}" ) fidelity_score = FidelityGrader.grade(current_sv, self._target_sv) # --- Efficiency --- depth = compute_circuit_depth(self._gates, self._num_qubits) gate_count = len(self._gates) efficiency_score = EfficiencyGrader.grade( depth, gate_count, self._task_config.get("max_depth", 10), self._task_config.get("max_gate_count", 15), ) # --- Noise --- noise_model = self._task_config.get("noise_model", "none") noise_score, noise_estimate = NoiseGrader.grade(self._gates, noise_model) # --- Constraints --- connectivity = self._task_config.get("connectivity", []) conn_tuples = [tuple(c) for c in connectivity] if connectivity else [] constraints_score = ConstraintsGrader.grade(self._gates, conn_tuples, self._num_qubits) # --- Aggregate --- if self._target_unitary is not None: aggregate = 0.7 * fidelity_score + 0.3 * efficiency_score else: aggregate = self._aggregate_grader.grade( fidelity_score, efficiency_score, noise_score, constraints_score ) # Almost-correct zone (smooth transition band 0.7–0.95 fidelity) if 0.7 < fidelity_score < 0.95: scale = 0.8 + 0.2 * (fidelity_score - 0.7) / 0.25 aggregate *= scale # Approximate task tolerance scaling target_tolerance = self._task_config.get("target_tolerance") if target_tolerance is not None and fidelity_score > target_tolerance: aggregate += 0.2 * (fidelity_score - target_tolerance) / (1.0 - target_tolerance) # Noise impact (realism) noise_penalty = (1.0 - noise_score) * 2.0 aggregate -= noise_penalty # STEP 7: Depth / time cost aggregate -= depth * 0.005 # STEP 7: SWAP count penalty in aggregate score aggregate -= 0.02 * self._swap_count # MAX_DEPTH soft penalty if depth > MAX_DEPTH: aggregate -= 0.1 * (depth - MAX_DEPTH) # Budget overflow penalty (smooth) strict_budget = self._task_config.get("strict_gate_budget") if strict_budget is not None and gate_count > strict_budget: overflow = gate_count - strict_budget penalty = overflow / strict_budget aggregate *= max(0.3, 1.0 - penalty) # Final clamp aggregate = max(0.0001, min(0.9999, aggregate)) assert 0.0 <= aggregate <= 1.0, f"Aggregate score out of range: {aggregate}" return { "fidelity": fidelity_score, "efficiency": efficiency_score, "noise": noise_score, "noise_estimate": noise_estimate, "constraints": constraints_score, "aggregate": aggregate, "depth": depth, "gate_count": gate_count, } def _compute_terminal_reward(self) -> float: """Compute bonus/penalty at episode end based on aggregate quality.""" scores = self._compute_all_scores() aggregate = scores["aggregate"] return 0.2 if aggregate > 0.8 else -0.2 # ------------------------------------------------------------------ # Termination check # ------------------------------------------------------------------ def _check_done(self) -> None: """Check if episode should terminate.""" if self._step_count >= self._max_steps: self._done = True # ------------------------------------------------------------------ # Observation builder # ------------------------------------------------------------------ def _build_observation( self, reward: float, error: Optional[str] = None ) -> QuantumObservation: """Build and return a QuantumObservation.""" scores = self._compute_all_scores() depth = int(scores.get("depth", 0)) gate_count = int(scores.get("gate_count", 0)) # Update trackers from latest scores self._prev_fidelity = scores["fidelity"] self._prev_depth = depth return QuantumObservation( circuit_gates=list(self._gates), fidelity=scores["fidelity"], depth=depth, gate_count=gate_count, noise_estimate=scores.get("noise_estimate", 0.0), valid_actions=self._get_valid_actions(), score=scores["aggregate"], task_id=self._task_id, num_qubits=self._num_qubits, max_steps=self._max_steps, steps_taken=self._step_count, target_description=self._task_config.get("description", ""), done=self._done, reward=reward, metadata={ "fidelity_score": scores["fidelity"], "efficiency_score": scores["efficiency"], "noise_score": scores["noise"], "constraints_score": scores["constraints"], "aggregate_score": scores["aggregate"], "error": error, }, ) def _get_valid_actions(self) -> List[str]: """Return list of valid action type strings.""" actions = ["ADD", "STOP"] if self._gates: actions.append("REMOVE") has_param_gates = any(g["gate"] in ("RX", "RZ") for g in self._gates) if has_param_gates: actions.append("PARAM") if self._num_qubits >= 2: actions.append("SWAP") return actions # --------------------------------------------------------------------------- # STEP 10: Direct smoke tests # --------------------------------------------------------------------------- if __name__ == "__main__": logging.basicConfig(level=logging.INFO, format="%(name)s | %(message)s") print("=== Quantum Circuit Optimization Environment -- Qiskit Smoke Test ===") print() env = QuantumCircuitEnvironment(seed=42) # --- Test 1: Bell State (H + CNOT) → fidelity ≈ 1 ------------------ print("--- Test 1: Bell State (H + CNOT) ---") obs = env.reset(config={"task_id": "easy"}) print(f" Initial fidelity: {obs.fidelity:.4f}, score: {obs.score:.4f}") obs = env.step(QuantumAction(action_type=ActionType.ADD, gate=GateType.H, qubits=[0])) print(f" After H(0): fid={obs.fidelity:.4f} score={obs.score:.4f} reward={obs.reward:.4f}") obs = env.step(QuantumAction(action_type=ActionType.ADD, gate=GateType.CNOT, qubits=[0, 1])) print(f" After CNOT(0,1): fid={obs.fidelity:.4f} score={obs.score:.4f} reward={obs.reward:.4f}") obs = env.step(QuantumAction(action_type=ActionType.STOP)) bell_score = obs.score assert bell_score > 0.5, f"FAIL: Bell correct circuit scored {bell_score:.4f}, expected > 0.5" print(f" PASS: Bell score = {bell_score:.4f}") print() # --- Test 2: GHZ Incomplete ------------------------------------------- print("--- Test 2: GHZ State (INCOMPLETE: H+CNOT, missing CNOT(1,2)) ---") obs = env.reset(config={"task_id": "medium"}) obs = env.step(QuantumAction(action_type=ActionType.ADD, gate=GateType.H, qubits=[0])) obs = env.step(QuantumAction(action_type=ActionType.ADD, gate=GateType.CNOT, qubits=[0, 1])) obs = env.step(QuantumAction(action_type=ActionType.STOP)) ghz_incomplete = obs.score print(f" STOP: score={ghz_incomplete:.4f} reward={obs.reward:.4f}") assert ghz_incomplete < 0.5, f"FAIL: GHZ incomplete scored {ghz_incomplete:.4f}, expected < 0.5" print(f" PASS: GHZ incomplete score = {ghz_incomplete:.4f}") print() # --- Test 3: GHZ Correct ---------------------------------------------- print("--- Test 3: GHZ State (CORRECT: H+CNOT(0,1)+CNOT(1,2)) ---") obs = env.reset(config={"task_id": "medium"}) obs = env.step(QuantumAction(action_type=ActionType.ADD, gate=GateType.H, qubits=[0])) obs = env.step(QuantumAction(action_type=ActionType.ADD, gate=GateType.CNOT, qubits=[0, 1])) obs = env.step(QuantumAction(action_type=ActionType.ADD, gate=GateType.CNOT, qubits=[1, 2])) obs = env.step(QuantumAction(action_type=ActionType.STOP)) ghz_correct = obs.score print(f" STOP: score={ghz_correct:.4f} reward={obs.reward:.4f}") assert ghz_correct > 0.5, f"FAIL: GHZ correct scored {ghz_correct:.4f}, expected > 0.5" print(f" PASS: GHZ correct score = {ghz_correct:.4f}") print() # --- Test 4: SWAP physically swaps qubit states ----------------------- print("--- Test 4: SWAP — qubit routing effect ---") # Qiskit uses LITTLE-ENDIAN qubit ordering: # qubit 0 = LSB, qubit 1 = MSB # X on qubit 0 → |10⟩ in big-endian = index 1 (binary ...01 in little-endian) # After SWAP(0,1) → qubit states exchange → different basis index gates_before_swap = [{"gate": "X", "qubits": [0]}] sv_before = compute_statevector(gates_before_swap, 2) gates_after_swap = [{"gate": "X", "qubits": [0]}, {"gate": "SWAP", "qubits": [0, 1]}] sv_after = compute_statevector(gates_after_swap, 2) # Find the dominant (non-zero) index for each state idx_before = int(np.argmax(np.abs(sv_before))) idx_after = int(np.argmax(np.abs(sv_after))) print(f" Before SWAP: dominant index = {idx_before} (amplitude={abs(sv_before[idx_before]):.4f})") print(f" After SWAP: dominant index = {idx_after} (amplitude={abs(sv_after[idx_after]):.4f})") # Both must be pure basis states assert abs(sv_before[idx_before]) > 0.99, f"Before SWAP not a pure state: {sv_before}" assert abs(sv_after[idx_after]) > 0.99, f"After SWAP not a pure state: {sv_after}" # SWAP must have changed which basis state is occupied assert idx_before != idx_after, ( f"SWAP did not change basis state index: before={idx_before}, after={idx_after}" ) print(" PASS: SWAP physically swapped qubit states") print() # --- Test 5: REMOVE reverts circuit state ----------------------------- print("--- Test 5: REMOVE — state reverts ---") obs = env.reset(config={"task_id": "easy"}) fid_empty = obs.fidelity obs_h = env.step(QuantumAction(action_type=ActionType.ADD, gate=GateType.H, qubits=[0])) fid_after_h = obs_h.fidelity obs_rm = env.step(QuantumAction(action_type=ActionType.REMOVE)) fid_after_remove = obs_rm.fidelity print(f" Fidelity (empty): {fid_empty:.6f}") print(f" Fidelity (H(0)): {fid_after_h:.6f}") print(f" Fidelity (REMOVE→ reverted): {fid_after_remove:.6f}") assert abs(fid_after_remove - fid_empty) < 1e-6, ( f"FAIL: state did not revert after REMOVE " f"(empty={fid_empty:.6f}, after_remove={fid_after_remove:.6f})" ) print(" PASS: REMOVE reverted circuit state correctly") print() # --- Discrimination summary ------------------------------------------- print("=" * 50) print("SCORE DISCRIMINATION SUMMARY:") print(f" Bell correct: {bell_score:.4f} (should be > 0.5)") print(f" GHZ incomplete: {ghz_incomplete:.4f} (should be < 0.5)") print(f" GHZ correct: {ghz_correct:.4f} (should be > 0.5)") assert bell_score > ghz_incomplete, "FAIL: Bell should score higher than incomplete GHZ" assert ghz_correct > ghz_incomplete, "FAIL: Correct GHZ should score higher than incomplete" print("\n ALL SMOKE TESTS PASSED!")