improvement

README.md

@@ -74,7 +74,7 @@ Current LLM benchmarks rely on static toy puzzles. This environment bridges the
 | `1` | **Cancel Identical Gates** | Removes self-inverse gate pairs (X·X=I, H·H=I, CNOT·CNOT=I) on the same qubits, not blocked by overlapping intermediate gates. | `+1.0` |
 | `2` | **Swap Commuting Gates** | Swaps the target gate with the next adjacent gate **only if** their qubit sets do not intersect. Enables bringing distant cancellable pairs together. | `-0.05` |
 | `3` | **H-X-H Identity Collapse** | Replaces a `H → X → H` sequence on the same qubit with a single `Z` gate (net: 2 gates removed). | `+2.0` |
-| `4` | **Entanglement Compression** | Replaces an adjacent `CNOT → SWAP` on the same qubits with a single `CZ` gate (net: 1 gate removed). | `+1.0` |
+| `4` | **Entanglement Compression** | Replaces a `CNOT(a,b) → CNOT(b,a) → CNOT(a,b)` sequence with a single `SWAP` gate — a standard compiler identity (net: 2 gates removed). | `+2.0` |
 
 > **Invalid actions** (out-of-bounds index, illegal non-commuting swap, pattern not present) incur a `-0.10` penalty. Circuit state remains unchanged.
 

inference.py

@@ -79,7 +79,7 @@ SYSTEM_PROMPT = textwrap.dedent(
               the same qubits, not blocked by intermediate gates sharing those qubits.
     Action 2: Swap adjacent commuting gates (gates on entirely non-overlapping qubits).
     Action 3: Replace an H-X-H sequence on the same qubit with a Z gate.
-    Action 4: Replace a CNOT-SWAP sequence on the same qubits with a CZ gate.
+    Action 4: Replace a CNOT(a,b)→CNOT(b,a)→CNOT(a,b) sequence with a single SWAP gate (3 alternating CNOTs collapse to 1 SWAP).
 
     You MUST output ONLY a valid JSON object with exactly two keys:
       "target_index" (integer) and "action_type" (integer 1-4).

models.py

@@ -41,7 +41,7 @@ class QuantumAction(Action):
         default=1, 
         ge=1, # Minimum action type is 1
         le=4, # Maximum action type is 4
-        description="1: Cancel identical gates, 2: Swap commuting gates, 3: Replace H-X-H with Z, 4: Replace CNOT-SWAP with CZ"
+        description="1: Cancel identical gates, 2: Swap commuting gates, 3: Replace H-X-H with Z, 4: Replace 3-CNOT sequence with SWAP"
     )
 
 

server/quantum_openenv_env_environment.py

@@ -71,6 +71,16 @@ class TaskConfig:
             insert_idx_2 = rng.randint(insert_idx_1, len(circuit))
             circuit.insert(insert_idx_2, gate2)
 
+        if self.use_entangling and self.num_qubits > 1:
+            num_patterns = 1 if self.name == "medium" else 2  # hard gets 2
+            for _ in range(num_patterns):
+                if rng.random() > 0.3:  # 70% chance per pattern, keeps it non-deterministic
+                    q1, q2 = rng.sample(range(self.num_qubits), 2)
+                    insert_at = rng.randint(0, len(circuit))
+                    circuit.insert(insert_at,     QuantumGate(name="CNOT", target_qubits=[q1, q2]))
+                    circuit.insert(insert_at + 1, QuantumGate(name="CNOT", target_qubits=[q2, q1]))
+                    circuit.insert(insert_at + 2, QuantumGate(name="CNOT", target_qubits=[q1, q2]))
+
         return circuit
 
 
@@ -247,20 +257,26 @@ class QuantumCircuitOptimizationEnvironment(Environment):
                     action_result = "identity_hxh_to_z"
                     self._used_advanced_actions = True
 
-        # ACTION 4: Replace CNOT-SWAP with CZ  (advanced identity)
+        # ACTION 4: Replace CNOT(a,b)→CNOT(b,a)→CNOT(a,b) with SWAP  (advanced identity)
         elif action_type == 4:
-            if target_index + 1 < len(self._circuit):
+            if target_index + 2 < len(self._circuit):
                 g1 = self._circuit[target_index]
                 g2 = self._circuit[target_index + 1]
+                g3 = self._circuit[target_index + 2]
+
+                qubits_ab = g1.target_qubits  # e.g. [0, 1]
+                qubits_ba = list(reversed(g1.target_qubits))  # e.g. [1, 0]
 
-                if (g1.name == "CNOT" and g2.name == "SWAP" and
-                        set(g1.target_qubits) == set(g2.target_qubits)):
+                if (g1.name == "CNOT" and g2.name == "CNOT" and g3.name == "CNOT" and
+                        g1.target_qubits == g3.target_qubits and
+                        g2.target_qubits == qubits_ba):
+                    self._circuit.pop(target_index + 2)
                     self._circuit.pop(target_index + 1)
                     self._circuit[target_index] = QuantumGate(
-                        name="CZ", target_qubits=g1.target_qubits
+                        name="SWAP", target_qubits=g1.target_qubits
                     )
-                    reward = 1.0
-                    action_result = "identity_cnot_swap_to_cz"
+                    reward = 2.0  # saves 2 gates, same as H-X-H identity
+                    action_result = "identity_3cnot_to_swap"
                     self._used_advanced_actions = True
 
         return self._build_observation(reward, action_result)
@@ -377,7 +393,7 @@ class QuantumCircuitOptimizationEnvironment(Environment):
             "1: Cancel identical self-inverse gates (H, X, Y, Z, CNOT, SWAP).\n\n"
             "2: Swap adjacent commuting gates (gates not sharing qubits).\n\n"
             "3: Replace an H-X-H sequence with a Z gate.\n\n"
-            "4: Replace a CNOT-SWAP sequence with a CZ gate.\n\n"
+            "4: Replace CNOT(a,b)→CNOT(b,a)→CNOT(a,b) with a single SWAP gate.\n\n"
             "CURRENT CIRCUIT STATE:\n\n"
         )