refactor: unify rollout onto inference engine KV-cache path

- RolloutGenerator now delegates prefill/decode to InferenceScheduler.run_batch (sync API, no background thread), sharing one KV-cache code path with the inference server and eliminating O(n^2) recompute in rollout
- Add sample(return_logprobs=) and Executor.execute_decode(return_logprobs=) to expose behaviour-policy log-probs through the engine; Task gains output_logprobs
- RolloutResult now subclasses RawRollout (adds rewards only), removing duplicated fields
- RolloutRunner.__call__ returns (result, is_fresh) instead of relying on object identity, removing the fragile refresh-detection contract
- Remove O(n^2) generate_responses helper and dead code (_tokenize_prompts, unused old_model arg)
- train_context.py wires InferenceScheduler directly instead of hand-rolling SamplingPipeline
- Tests: +11 covering return_logprobs, run_batch, and KV-cache-backed rollout semantics; 404 pass

astrai/trainer/strategy.py

@@ -109,7 +109,6 @@ class BaseStrategy(ABC):
         self.executor = kwargs.pop("executor", None)
         self.extra_kwargs = kwargs
         self._rollout_runner = None
-        self._prev_rollout_result = None
 
     @abstractmethod
     def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
@@ -159,10 +158,9 @@ class BaseStrategy(ABC):
         if self._rollout_runner is None:
             return self.compute_loss(batch)
 
-        result = self._rollout_runner(batch)
-        if result is not self._prev_rollout_result:
+        result, is_fresh = self._rollout_runner(batch)
+        if is_fresh:
             self._on_rollout_refresh()
-            self._prev_rollout_result = result
             if self.executor and self.executor.sync_gradients:
                 self._rollout_runner.step()