From 62a036c653e0f68ec9e8677e1b6e4e58a0daa007 Mon Sep 17 00:00:00 2001
From: NopeCharon <115794254+NopeCharon@users.noreply.github.com>
Date: Thu, 29 Oct 2026 02:21:14 +0530
Subject: [PATCH] Use configured observations in Pinball and fix render()

Pinball overrode get_observations() to always return the lift and drag
coefficients, bypassing the ObservationFunction that FlowConfig builds
from observation_type. num_outputs, and therefore the FlowEnv
observation space, still followed observation_type, so with probe
observations the environment returned 6 values for an observation space
of 2 * len(probes) entries and failed gymnasium's env checker. Drop the
override: the inherited FlowConfig.get_observations() already returns
(CL1, CL2, CL3, CD1, CD2, CD3) for the default "lift_drag" observation,
so the default behaviour is unchanged.

Pinball.render() referred to self.flow.u, self.flow.pressure_space,
self.flow.x0, self.flow.y0 and self.flow.rad, but Pinball is itself the
flow and has no flow attribute, so every call raised AttributeError.
Use the instance attributes, matching the other Firedrake flows, and
convert the radius Constant to a float for matplotlib.

Add tests for probe observations through FlowEnv and for render(), and
correct the documented ordering of the Pinball lift_drag observation.

Closes #361
---
 .../examples/firedrake/getting_started.md     |  2 +-
 examples/firedrake/getting_started/README.md  |  2 +-
 hydrogym/firedrake/envs/pinball/flow.py       | 10 ++----
 test/test_pinball.py                          | 31 +++++++++++++++++++
 4 files changed, 36 insertions(+), 9 deletions(-)

diff --git a/docs/docs/examples/firedrake/getting_started.md b/docs/docs/examples/firedrake/getting_started.md
index cf22a344..9588add9 100644
--- a/docs/docs/examples/firedrake/getting_started.md
+++ b/docs/docs/examples/firedrake/getting_started.md
@@ -139,7 +139,7 @@ python train_sb3_firedrake.py --env cylinder --algo PPO --total-timesteps 100000
 ## Observation types
 
 **Force-based**
-- `'lift_drag'` — returns `(CL, CD)` for single-cylinder cases; `(CL1, CD1, CL2, CD2, CL3, CD3)` for Pinball.
+- `'lift_drag'` — returns `(CL, CD)` for single-cylinder cases; `(CL1, CL2, CL3, CD1, CD2, CD3)` for Pinball.
 
 **Sensor-based**
 - `'stress_sensor'` — returns wall shear stress as a scalar.
diff --git a/examples/firedrake/getting_started/README.md b/examples/firedrake/getting_started/README.md
index 81c6ca81..bb6ad6ae 100644
--- a/examples/firedrake/getting_started/README.md
+++ b/examples/firedrake/getting_started/README.md
@@ -139,7 +139,7 @@ Open `config_reference.py` and copy the example that matches your use case.
 ## Observation Types
 
 ### 1. **Force-Based Observations**
-- `'lift_drag'` → Returns `(CL, CD)` for cylinder/rotary, `(CL1, CD1, CL2, CD2, CL3, CD3)` for pinball
+- `'lift_drag'` → Returns `(CL, CD)` for cylinder/rotary, `(CL1, CL2, CL3, CD1, CD2, CD3)` for pinball
 
 ### 2. **Sensor-Based Observations**
 - `'stress_sensor'` → Returns wall shear stress (scalar)
diff --git a/hydrogym/firedrake/envs/pinball/flow.py b/hydrogym/firedrake/envs/pinball/flow.py
index f0ed3d36..b0e4eb3d 100644
--- a/hydrogym/firedrake/envs/pinball/flow.py
+++ b/hydrogym/firedrake/envs/pinball/flow.py
@@ -105,10 +105,6 @@ def linearize_bcs(self, function_spaces=None):
         self.bcu_inflow.set_value(fd.Constant((0, 0)))
         self.bcu_freestream.set_value(fd.Constant(0.0))
 
-    def get_observations(self):
-        CL, CD = self.compute_forces()
-        return [*CL, *CD]
-
     def evaluate_objective(self, q=None):
         CL, CD = self.compute_forces(q=q)
         return sum(CD)
@@ -118,7 +114,7 @@ def render(self, mode="human", clim=None, levels=None, cmap="RdBu", **kwargs):
             clim = (-2, 2)
         if levels is None:
             levels = np.linspace(*clim, 10)
-        vort = fd.project(fd.curl(self.flow.u), self.flow.pressure_space)
+        vort = fd.project(fd.curl(self.u), self.pressure_space)
         im = tricontourf(
             vort,
             cmap=cmap,
@@ -129,6 +125,6 @@ def render(self, mode="human", clim=None, levels=None, cmap="RdBu", **kwargs):
             **kwargs,
         )
 
-        for x0, y0 in zip(self.flow.x0, self.flow.y0):
-            cyl = plt.Circle((x0, y0), self.flow.rad, edgecolor="k", facecolor="gray")
+        for x0, y0 in zip(self.x0, self.y0):
+            cyl = plt.Circle((x0, y0), float(self.rad), edgecolor="k", facecolor="gray")
             im.axes.add_artist(cyl)
diff --git a/test/test_pinball.py b/test/test_pinball.py
index 6ce3ff4e..9d3ad123 100644
--- a/test/test_pinball.py
+++ b/test/test_pinball.py
@@ -1,3 +1,4 @@
+import matplotlib.pyplot as plt
 import numpy as np
 
 import hydrogym.firedrake as hgym
@@ -83,3 +84,33 @@ def feedback_ctrl(y, K=None):
     for _ in range(10):
         y, reward, terminated, truncated, info = env.step(u)
         u = feedback_ctrl(y)
+
+
+def test_probe_observations():
+    probes = [(x, y) for x in np.linspace(2.0, 6.0, 3) for y in (-1.0, 1.0)]
+    env_config = {
+        "flow": hgym.Pinball,
+        "flow_config": {
+            "mesh": "medium",
+            "observation_type": "velocity_probes",
+            "probes": probes,
+        },
+        "solver": hgym.SemiImplicitBDF,
+        "solver_config": {
+            "dt": 1e-2,
+        },
+    }
+    env = hgym.FlowEnv(env_config)
+
+    y, info = env.reset()
+    assert env.observation_space.shape == (2 * len(probes),)
+    assert y.shape == env.observation_space.shape
+
+    y, reward, terminated, truncated, info = env.step(np.zeros(3))
+    assert y.shape == env.observation_space.shape
+
+
+def test_render():
+    flow = hgym.Pinball(mesh="medium")
+    flow.render()
+    plt.close("all")