Resolve review comments

Author: rchardx

areal/api/workflow_api.py

@@ -1,7 +1,6 @@
 from __future__ import annotations  # noqa
 
 from abc import ABC, abstractmethod
-from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any
 
 from areal.experimental.openai.types import InteractionWithTokenLogpReward
@@ -10,23 +9,6 @@
     from areal.api.engine_api import InferenceEngine
 
 
-@dataclass(slots=True)
-class WorkflowTaskInput:
-    """Input payload provided to :class:`RolloutWorkflow` implementations.
-
-    Parameters
-    ----------
-    data : dict[str, Any]
-        Original sample provided by the dataloader or caller.
-    session_id : int | None, optional
-        Identifier registered with the global session tracer when tracing is
-        enabled.
-    """
-
-    data: dict[str, Any]
-    session_id: int | None = None
-
-
 class RolloutWorkflow(ABC):
     @abstractmethod
     async def arun_episode(

areal/workflow/vision_rlvr.py

@@ -157,7 +157,7 @@ async def arun_episode(
                 "multi_modal_input": multi_modal_input,
                 "versions": torch.tensor(versions, dtype=torch.int32).unsqueeze(0),
                 "attention_mask": torch.ones(len(seq), dtype=torch.bool).unsqueeze(0),
-                "rewards": torch.tensor([reward], dtype=torch.float32),
+                "rewards": torch.tensor(reward, dtype=torch.float32).unsqueeze(0),
             }
             results.append(res)