Use configured observations in Pinball and fix render()

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.

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)

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)

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")