rLLM Data Structures API

Module: rllm.agents.agent

Reference API for the core rollout and training data structures used across rLLM workflows.

---

Step

rllm.agents.agent.Step dataclass

Atomic unit of one model interaction in a rollout.

A step stores token-level fields from the rollout engine (prompt_ids, response_ids, logprobs) together with execution context and training signals.

Attributes:

Name	Type	Description
prompt_ids	list[int] | list[Any]	Input token IDs for this model call.
response_ids	list[int]	Generated token IDs for this model call.
logprobs	list[float]	Token-level log probabilities aligned with response_ids.
chat_completions	list[dict[str, str]]	Chat message history at this step.
observation	Any	Environment/workflow observation before generation.
thought	str	Optional model reasoning text.
action	Any	Parsed action emitted from the model response.
model_response	str	Raw response content from the model.
model_output	ModelOutput | None	Original rollout engine output payload.
info	dict	Additional per-step metadata.
reward	float	Step-level reward signal.
done	bool	Whether this step ends the trajectory.
mc_return	float	Monte-Carlo return (if computed by workflow logic).
advantage	list[float] | float | None	Step advantage signal; can be scalar or token-level list.

Source code in rllm/agents/agent.py

@dataclass
class Step:
    """
    Atomic unit of one model interaction in a rollout.

    A step stores token-level fields from the rollout engine (`prompt_ids`,
    `response_ids`, `logprobs`) together with execution context and training
    signals.

    Attributes:
        prompt_ids: Input token IDs for this model call.
        response_ids: Generated token IDs for this model call.
        logprobs: Token-level log probabilities aligned with `response_ids`.
        chat_completions: Chat message history at this step.
        observation: Environment/workflow observation before generation.
        thought: Optional model reasoning text.
        action: Parsed action emitted from the model response.
        model_response: Raw response content from the model.
        model_output: Original rollout engine output payload.
        info: Additional per-step metadata.
        reward: Step-level reward signal.
        done: Whether this step ends the trajectory.
        mc_return: Monte-Carlo return (if computed by workflow logic).
        advantage: Step advantage signal; can be scalar or token-level list.
    """

    # this is to accomodate the fact that for backend like `tinker`, the prompt_ids might contain special image blocks
    prompt_ids: list[int] | list[Any] = field(default_factory=list)
    response_ids: list[int] = field(default_factory=list)
    logprobs: list[float] = field(default_factory=list)

    chat_completions: list[dict[str, str]] = field(default_factory=list)

    observation: Any = None
    thought: str = ""
    action: Any = None
    model_response: str = ""
    model_output: ModelOutput | None = None
    info: dict = field(default_factory=dict)  # Store any additional info.

    # field below are filled by the engine
    reward: float = 0.0
    done: bool = False
    mc_return: float = 0.0

    # field below are filled by the advantage computer. Note when advantage is a list, it is per-token advantages.
    # TODO: potentially rename this as "advantages" so its clearer that it allows a generic list.
    advantage: list[float] | float | None = None

    def __post_init__(self):
        if self.model_output is None:
            return
        # backfill fields like prompt_ids, response_ids, logprobs, etc.
        if len(self.prompt_ids) == 0 and self.model_output.prompt_ids is not None:
            self.prompt_ids = self.model_output.prompt_ids
        if len(self.response_ids) == 0 and self.model_output.completion_ids is not None:
            self.response_ids = self.model_output.completion_ids
        if len(self.logprobs) == 0 and self.model_output.logprobs is not None:
            self.logprobs = self.model_output.logprobs

        # check that the token ids are filled
        # TODO(listar2000): this might cause compatibility issue. Double check if we should make these assertions.
        # assert len(self.prompt_ids) > 0, "prompt_ids is empty"
        # assert len(self.response_ids) > 0, "response_ids is empty"

        # check that the lengths would match up
        if len(self.logprobs) > 0:
            assert len(self.response_ids) == len(self.logprobs), f"length mismatch between response_ids and logprobs, got {len(self.response_ids)}, {len(self.logprobs)}"

    def to_dict(self) -> dict:
        return {
            "prompt_ids": self.prompt_ids,
            "response_ids": self.response_ids,
            "logprobs": self.logprobs,
            "chat_completions": self.chat_completions,
            "observation": self.observation,
            "thought": self.thought,
            "action": self.action.action if isinstance(self.action, Action) else self.action,
            "model_response": self.model_response,
            "model_output": self.model_output.to_dict() if self.model_output is not None else None,
            "info": self.info,
            "reward": self.reward,
            "done": self.done,
            "mc_return": self.mc_return,
            "advantage": self.advantage,
        }

    @classmethod
    def from_dict(cls, data: dict) -> Step:
        from rllm.engine.rollout import ModelOutput

        return cls(
            prompt_ids=data["prompt_ids"],
            response_ids=data["response_ids"],
            logprobs=data["logprobs"],
            chat_completions=data["chat_completions"],
            observation=data["observation"],
            thought=data["thought"],
            action=data["action"],
            model_response=data["model_response"],
            model_output=ModelOutput.from_dict(data["model_output"]) if data.get("model_output", None) is not None else None,
            info=data.get("info", {}),
            reward=data["reward"],
            done=data["done"],
            mc_return=data["mc_return"],
            advantage=data["advantage"],
        )

    @classmethod
    def from_model_output(cls, model_output: ModelOutput, messages: list[dict] | None = None, action: Any | None = None) -> Step:
        return cls(
            prompt_ids=model_output.prompt_ids or [],
            response_ids=model_output.completion_ids or [],
            logprobs=model_output.logprobs or [],
            chat_completions=(messages or []) + [{"role": "assistant", "content": model_output.content, "reasoning": model_output.reasoning}],
            thought=model_output.reasoning or "",
            action=action,
            model_response=model_output.content or "",
            model_output=model_output,
        )

---

Trajectory

rllm.agents.agent.Trajectory dataclass

Ordered sequence of steps for one role/agent thread.

A trajectory is the primary container for per-role rollout history and trajectory-level reward.

Attributes:

Name	Type	Description
uid	str	Unique trajectory identifier.
name	str	Role/name for grouping and metrics (for example, solver).
task	Any	Task payload associated with the trajectory.
steps	list[Step]	Ordered list of step records.
reward	float | None	Optional trajectory-level reward.
info	dict	Additional trajectory metadata.

Source code in rllm/agents/agent.py

@dataclass
class Trajectory:
    """
    Ordered sequence of steps for one role/agent thread.

    A trajectory is the primary container for per-role rollout history and
    trajectory-level reward.

    Attributes:
        uid: Unique trajectory identifier.
        name: Role/name for grouping and metrics (for example, `solver`).
        task: Task payload associated with the trajectory.
        steps: Ordered list of step records.
        reward: Optional trajectory-level reward.
        info: Additional trajectory metadata.
    """

    uid: str = field(default_factory=lambda: str(uuid.uuid4()))  # unique id to deduplicate on
    name: str = _DEFAULT_TRAJ_NAME
    task: Any = None
    steps: list[Step] = field(default_factory=list)
    reward: float | None = None  # it is possible that the trajectory-level reward does not exist
    info: dict = field(default_factory=dict)

    def to_dict(self):
        # Remove large/non-serializable payloads (e.g., images) from task
        def _sanitize_task(task_obj):
            if isinstance(task_obj, dict):
                cleaned = {k: v for k, v in task_obj.items() if k not in ("image", "images")}
                return cleaned
            return task_obj

        return {
            "uid": self.uid,
            "name": self.name,
            "task": _sanitize_task(self.task),
            "steps": [step.to_dict() for step in self.steps],
            "reward": float(self.reward) if self.reward is not None else None,
            "info": self.info,
        }

    @classmethod
    def from_dict(cls, data: dict) -> Trajectory:
        """Create Trajectory from dictionary, properly deserializing Step objects."""
        return cls(
            uid=data.get("uid", str(uuid.uuid4())),
            name=data["name"],
            task=data["task"],
            steps=[Step.from_dict(step_data) for step_data in data.get("steps", [])],
            reward=data["reward"],
            info=data.get("info", {}),
        )

    def is_cumulative(self) -> bool:
        """
        Returns True if for every step after the first, its chat_completions is an exact superset
        of the previous step's chat_completions (i.e., the previous chat_completions is a prefix).
        """
        prev = None
        for step in self.steps:
            if prev is not None:
                prev_cc = prev.chat_completions
                curr_cc = step.chat_completions
                if not (len(curr_cc) >= len(prev_cc) and curr_cc[: len(prev_cc)] == prev_cc):
                    return False
            prev = step
        return True

from_dict classmethod

from_dict(data: dict) -> Trajectory

Create Trajectory from dictionary, properly deserializing Step objects.

Source code in rllm/agents/agent.py

@classmethod
def from_dict(cls, data: dict) -> Trajectory:
    """Create Trajectory from dictionary, properly deserializing Step objects."""
    return cls(
        uid=data.get("uid", str(uuid.uuid4())),
        name=data["name"],
        task=data["task"],
        steps=[Step.from_dict(step_data) for step_data in data.get("steps", [])],
        reward=data["reward"],
        info=data.get("info", {}),
    )

is_cumulative

is_cumulative() -> bool

Returns True if for every step after the first, its chat_completions is an exact superset of the previous step's chat_completions (i.e., the previous chat_completions is a prefix).

Source code in rllm/agents/agent.py

def is_cumulative(self) -> bool:
    """
    Returns True if for every step after the first, its chat_completions is an exact superset
    of the previous step's chat_completions (i.e., the previous chat_completions is a prefix).
    """
    prev = None
    for step in self.steps:
        if prev is not None:
            prev_cc = prev.chat_completions
            curr_cc = step.chat_completions
            if not (len(curr_cc) >= len(prev_cc) and curr_cc[: len(prev_cc)] == prev_cc):
                return False
        prev = step
    return True

---

Episode

rllm.agents.agent.Episode dataclass

Workflow-level rollout output containing one or more trajectories.

Episode is the unit returned by Workflow.run(...) and usually represents one (task_id, rollout_idx) execution.

Attributes:

Name	Type	Description
id	str	Rollout identifier, typically {task_id}:{rollout_idx}.
task	Any	Original task payload.
termination_reason	TerminationReason | None	Workflow termination status.
is_correct	bool	Workflow-level correctness flag.
trajectories	list[Trajectory]	Trajectories generated in this rollout.
metrics	dict	Workflow-defined episode metrics.
info	dict	Additional episode metadata.

Source code in rllm/agents/agent.py

@dataclass
class Episode:
    """
    Workflow-level rollout output containing one or more trajectories.

    Episode is the unit returned by `Workflow.run(...)` and usually represents
    one `(task_id, rollout_idx)` execution.

    Attributes:
        id: Rollout identifier, typically `{task_id}:{rollout_idx}`.
        task: Original task payload.
        termination_reason: Workflow termination status.
        is_correct: Workflow-level correctness flag.
        trajectories: Trajectories generated in this rollout.
        metrics: Workflow-defined episode metrics.
        info: Additional episode metadata.
    """

    id: str = ""  # rollout id e.g., task_id:rollout_idx
    task: Any = None
    termination_reason: TerminationReason | None = None  # noqa: F821
    is_correct: bool = False
    trajectories: list[Trajectory] = field(default_factory=list)
    metrics: dict = field(default_factory=dict)
    info: dict = field(default_factory=dict)

    def to_dict(self):
        # Remove large/non-serializable payloads (e.g., images) from task
        def _sanitize_task(task_obj):
            if isinstance(task_obj, dict):
                cleaned = {k: v for k, v in task_obj.items() if k not in ("image", "images")}
                return cleaned
            return task_obj

        return {
            "id": self.id,
            "task": _sanitize_task(self.task),
            "termination_reason": self.termination_reason.value if self.termination_reason is not None else None,
            "is_correct": bool(self.is_correct),
            "trajectories": [trajectory.to_dict() for trajectory in self.trajectories],
            "metrics": self.metrics,
            "info": self.info,
        }

    @classmethod
    def from_dict(cls, data: dict) -> Episode:
        """Create Episode from dictionary, properly deserializing Trajectory objects."""
        from rllm.workflows.workflow import TerminationReason

        return cls(
            id=data["id"],
            task=data["task"],
            termination_reason=TerminationReason(data.get("termination_reason", TerminationReason.UNKNOWN)),
            is_correct=data["is_correct"],
            trajectories=[Trajectory.from_dict(trajectory_data) for trajectory_data in data["trajectories"]],
            metrics=data.get("metrics", {}),
            info=data.get("info", {}),
        )

    @cached_property
    def task_id(self) -> str:
        return self.id.split(":")[0]

    @cached_property
    def rollout_idx(self) -> str:
        return self.id.split(":")[1]

from_dict classmethod

from_dict(data: dict) -> Episode

Create Episode from dictionary, properly deserializing Trajectory objects.

Source code in rllm/agents/agent.py

@classmethod
def from_dict(cls, data: dict) -> Episode:
    """Create Episode from dictionary, properly deserializing Trajectory objects."""
    from rllm.workflows.workflow import TerminationReason

    return cls(
        id=data["id"],
        task=data["task"],
        termination_reason=TerminationReason(data.get("termination_reason", TerminationReason.UNKNOWN)),
        is_correct=data["is_correct"],
        trajectories=[Trajectory.from_dict(trajectory_data) for trajectory_data in data["trajectories"]],
        metrics=data.get("metrics", {}),
        info=data.get("info", {}),
    )

---

TrajectoryGroup

rllm.agents.agent.TrajectoryGroup dataclass

A group of trajectories for advantage computation.

Unlike Episode (which represents raw rollout data), TrajectoryGroup is specifically structured for advantage computation. All trajectories in a group will have their rewards compared to compute advantages (e.g., via GRPO).

Attributes:

Name	Type	Description
trajectories	list[Trajectory]	List of trajectories to compare for advantage computation
group_id	str	Optional identifier for the group (e.g., "task1:agent_0")
metadata	list[dict]	List of metadata for each trajectory in the group

Source code in rllm/agents/agent.py

@dataclass
class TrajectoryGroup:
    """
    A group of trajectories for advantage computation.

    Unlike Episode (which represents raw rollout data), TrajectoryGroup is specifically
    structured for advantage computation. All trajectories in a group will have their
    rewards compared to compute advantages (e.g., via GRPO).

    Attributes:
        trajectories: List of trajectories to compare for advantage computation
        group_id: Optional identifier for the group (e.g., "task1:agent_0")
        metadata: List of metadata for each trajectory in the group
    """

    trajectories: list[Trajectory]
    group_id: str = ""
    metadata: list[dict] = field(default_factory=list)

    @cached_property
    def group_role(self) -> str:
        return self.group_id.split(":")[1] if ":" in self.group_id[:-1] else "all_groups"

    @cached_property
    def task_id(self) -> str:
        return self.group_id.split(":")[0]