fix(applications): RLRW returns the choices it simulates

CHANGELOG.md

@@ -5,7 +5,7 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
-## [0.26.0.dev1] - 2026-10-05
+## [0.26.0.dev2] - 2026-10-06
 
 ### Added
 
@@ -72,6 +72,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fixed `cpm.models.activation.ProspectUtility` with `weighting="prelec"` failing for options with several outcomes
 - Fixed `cpm.applications.decision_making.PTSM2025` stopping a fit with `ValueError` when its exponentials overflowed
 - Fixed `cpm.models.learning.HumbleTeacher` raising `IndexError` for 1D weights
+- Fixed `cpm.applications.reinforcement_learning.RLRW` not returning the choices it simulates with `generate=True`; it now outputs `response` ([#100](https://github.com/DevComPsy/cpm/issues/100))
 - Fixed `cpm.generators.Simulator` not raising its intended `TypeError` for an ungrouped `pandas.DataFrame`
 - Fixed `cpm.generators.Simulator` raising `AttributeError` for parameters in a `pandas.DataFrame` with grouped data
 - Fixed `cpm.generators.Simulator` raising `ValueError` for a single parameter set given as `Parameters`, `dict` or `pandas.Series`

cpm/__version__.py

@@ -1 +1 @@
-__version__ = "0.26.0.dev1"
+__version__ = "0.26.0.dev2"

cpm/applications/_sessions.py

@@ -68,7 +68,9 @@ def rlrw(alpha, temperature, initial_values, arms, rewards, response, generate,
 
     Returns
     -------
-    policy, reward, values, change, dependent : numpy.ndarray
+    policy, response, reward, values, change, dependent : numpy.ndarray
+        `response` is the arm chosen on each trial: the sampled one if
+        `generate`, otherwise the observed one.
     """
     n, k = arms.shape
     d = initial_values.shape[0]
@@ -78,6 +80,7 @@ def rlrw(alpha, temperature, initial_values, arms, rewards, response, generate,
         raise ValueError("RLRW needs at least two arms.")
     values = initial_values.copy()
     policy = np.empty((n, k))
+    out_response = np.empty(n, dtype=np.int64)
     reward = np.empty(n)
     history = np.empty((n, d))
     change = np.empty((n, d))
@@ -94,6 +97,9 @@ def rlrw(alpha, temperature, initial_values, arms, rewards, response, generate,
         choice = choose(policy[t], uniforms[t]) if generate else response[t]
         if not -k <= choice < k:
             raise IndexError("RLRW: response must be an arm, from 0 to the number of arms - 1.")
+        if choice < 0:
+            choice += k
+        out_response[t] = choice
         stimulus = arms[t, choice] - 1
         if stimulus < 0:
             stimulus += d
@@ -107,7 +113,7 @@ def rlrw(alpha, temperature, initial_values, arms, rewards, response, generate,
         change[t] = update[0]
         values += update[0]
         history[t] = values
-    return policy, reward, history, change, policy[:, 1].copy()
+    return policy, out_response, reward, history, change, policy[:, 1].copy()
 
 
 @njit

cpm/applications/reinforcement_learning.py

@@ -87,7 +87,7 @@ class _RLRWModel(SessionModel):
     """The model function of `RLRW`, for all trials at once."""
 
     def __call__(self, parameters, data):
-        policy, reward, values, change, dependent = self.sessions.rlrw(
+        policy, response, reward, values, change, dependent = self.sessions.rlrw(
             float(parameters.alpha),
             float(parameters.temperature),
             numpy.asarray(parameters.values, dtype=float),
@@ -99,6 +99,7 @@ def __call__(self, parameters, data):
         )
         return {
             "policy": policy,
+            "response": response,
             "reward": reward,
             "values": values,
             "change": change,
@@ -140,11 +141,13 @@ class RLRW(Application):
         The number of distinct stimuli present in the data.
     parameters_settings: list-like
         The parameters to be fit by the model. The parameters must be specified as a list of lists, with each list containing the value, lower, and upper bounds of the parameter. See Notes for more information on how to specify parameters and for the default settings.
+    generate: bool
+        If True, the model samples a choice on each trial from its policy instead of using the observed one, and learns from the reward of the sampled arm. Default is False.
 
     Returns
     -------
     cpm.generators.Wrapper
-        A cpm.generators.Wrapper object.
+        A cpm.generators.Wrapper object. Its outputs include `response`, the arm chosen on each trial: the sampled one if `generate` is True, otherwise the observed one. Simulated data can therefore be fitted directly.
 
     Examples
     --------

test/applications/test_backends.py

@@ -112,6 +112,16 @@ def cubic(x=None, alpha=None, lambda_loss=None):
     ("PTSM2025-standard", PTSM2025, risky, dict(parameters_settings=PT, variant="standard")),
 ]
 IDS = [case[0] for case in CASES]
+## outputs added since the reference was frozen, tested separately
+ADDED = {RLRW: ["response"]}
+
+
+def without_added(case, output):
+    """`output` (an export or a per-trial record) without the outputs added since the reference."""
+    added = ADDED.get(case[1], [])
+    if isinstance(output, pd.DataFrame):
+        return output.drop(columns=added)
+    return {key: value for key, value in output.items() if key not in added}
 
 
 def on_backend(model, backend):
@@ -151,7 +161,7 @@ def test_runs_equal_the_per_trial_implementation(case, backend):
         model.run()
         np.testing.assert_allclose(model.dependent, reference["dependent"][scale], **TOLERANCE)
         if scale in reference["exports"]:
-            expected, got = pd.DataFrame(reference["exports"][scale]), model.export()
+            expected, got = pd.DataFrame(reference["exports"][scale]), without_added(case, model.export())
             assert list(got.columns) == list(expected.columns)
             assert dict(got.dtypes) == dict(expected.dtypes)
             np.testing.assert_allclose(got.to_numpy(float), expected.to_numpy(float), **TOLERANCE)
@@ -161,7 +171,7 @@ def test_runs_equal_the_per_trial_implementation(case, backend):
             ## the per-trial records of the run, as a per-trial Wrapper kept them
             assert len(model.simulation) == len(model.dependent)
             for got, expected in zip(model.simulation[:3], reference["simulation"]):
-                assert_same_records(got, expected)
+                assert_same_records(without_added(case, got), expected)
 
 
 @pytest.mark.parametrize("backend", BACKENDS)
@@ -185,7 +195,7 @@ def test_the_per_trial_model_function_is_unchanged(case, backend):
     np.random.seed(5)
     first = unpack_trials(model.data, 0, model.__pandas__)
     got = model.model(parameters=model.parameters, trial=first)
-    assert_same_records(got, REFERENCE["cases"][case[0]]["model"])
+    assert_same_records(without_added(case, got), REFERENCE["cases"][case[0]]["model"])
 
 
 @pytest.mark.parametrize("backend", BACKENDS)
@@ -216,6 +226,50 @@ def test_simulations_equal_the_per_trial_implementation(backend):
                                   rtol=1e-12, atol=1e-12)
 
 
+def simulate_rlrw(backend, seed=0):
+    data = pd.concat([bandit(p).assign(ppt=p) for p in (1, 2, 3)], ignore_index=True)
+    draws = pd.DataFrame({"alpha": [0.3, 0.5, 0.7], "temperature": [2.0, 5.0, 8.0]})
+    wrapper = on_backend(RLRW(data=bandit(1), dimensions=4, parameters_settings=[[0.5, 0, 1], [5, 0, 10]],
+                              generate=True), backend)
+    simulator = Simulator(wrapper=wrapper, data=data.groupby("ppt"), parameters=draws)
+    np.random.seed(seed)
+    simulator.run()
+    return data, simulator.export()
+
+
+@pytest.mark.parametrize("backend", BACKENDS)
+def test_rlrw_returns_the_choices_it_simulates(backend):
+    """The sampled choices follow the policy, and the reward is that of the chosen arm."""
+    data, simulated = simulate_rlrw(backend)
+    response = simulated["response"].to_numpy()
+    assert simulated["response"].dtype == np.int64
+    ## the uniforms are drawn per participant in turn, one per trial, as `choose` uses them
+    np.random.seed(0)
+    uniforms = np.random.random_sample(len(simulated))
+    cdf = simulated[["policy_0", "policy_1"]].cumsum(axis=1).to_numpy()
+    np.testing.assert_array_equal(response, (cdf / cdf[:, -1:] <= uniforms[:, None]).sum(axis=1))
+    rewards = data[["reward_left", "reward_right"]].to_numpy()
+    np.testing.assert_array_equal(simulated["reward"], rewards[np.arange(len(data)), response])
+    ## not all the same choice
+    assert 0 < response.mean() < 1
+
+
+@pytest.mark.parametrize("backend", BACKENDS)
+def test_rlrw_returns_the_observed_choices_when_not_generating(backend):
+    data = bandit()
+    model = on_backend(RLRW(data=data, dimensions=4, parameters_settings=[[0.3, 0, 1], [4, 0, 10]]), backend)
+    model.run()
+    np.testing.assert_array_equal(model.export()["response"], data["response"])
+    assert model.simulation[0]["response"] == data["response"][0]
+
+
+@pytest.mark.skipif(not _jit.JIT_ENABLED, reason="numba is not installed or disabled")
+def test_rlrw_simulates_the_same_choices_with_and_without_numba():
+    _, python = simulate_rlrw("python", seed=11)
+    _, numba = simulate_rlrw("numba", seed=11)
+    np.testing.assert_array_equal(python["response"], numba["response"])
+
+
 @pytest.mark.skipif(not _jit.JIT_ENABLED, reason="numba is not installed or disabled")
 @pytest.mark.parametrize("case", CASES, ids=IDS)
 def test_backends_agree(case):