How to Integrate Process Rewards in LangGraph Agents and Handle Node Execution Failures?

[dalek-who]

Hello, I’m exploring how to combine LangGraph with process rewards and handle node failures in sequential agents. I have two main questions that aren’t covered in existing tutorials:

Process reward + LangGraph

Is it possible to call emit_reward directly within an agent’s node to emit intermediate rewards during training? For example:

import agentlightning as agl
from langgraph.graph import END, START, MessagesState, StateGraph
from langgraph.graph.state import CompiledStateGraph

class Agent:
        def ___init___(...):
                 ...

        def node_1(self, state: State) -> State:
                result = self.invoke_prompt(prompt)
                reward_1 = ...
                agl.emit_reward(reward_1)
                new_state = ...
                return new_state

        def node_2(self, state: State) -> State:
                result = self.invoke_prompt(prompt)
                reward_2 = ...
                agl.emit_reward(reward_2)
                new_state = ...
                return new_state

        def graph(self) -> CompiledStateGraph[State]:
            builder = StateGraph(State)
            builder.add_node(self.node1)
            builder.add_node(self.node2)

            builder.add_edge(START, "node1")
            builder.add_edge("node1", "node2")
            builder.add_edge("node2", END)

            return builder.compile()


class LitSQLAgent(agl.LitAgent[Dict[str, Any]]):

    def __init__(...):
        ...

    def rollout(
        self,
        task: Dict[str, Any],
        resources: agl.NamedResources,
        rollout: agl.Rollout,
    ) -> None:
        agent = Agent(...).graph()
        try:
                # Required to make the langchain tracing work
                handler = self.tracer.get_langchain_handler()
                result = agent.invoke( 
                    {"question": question},
                    {"callbacks": [handler] if handler else [], "recursion_limit": 100},
                )
        except Exception as e:
                logger.exception(f"[Rollout {rollout_id}] Error during agent invocation: {e}")
                return
        return

Handle node failures

If I should add a confitional edge for each intermediate node?

class Agent:
    ...  # nearly the same as the agent above

        def check_1(self, state: State):
            if is_training and not is_success(state):
                    return END
            else:
                    return "node2"

        def graph(self) -> CompiledStateGraph[State]:
            builder = StateGraph(State)
            builder.add_node(self.node1)
            builder.add_node(self.node2)

            builder.add_edge(START, "node1")
            builder.add_edge("node2", END)
            builder.add_conditional_edges("node1", self.check_1)

            return builder.compile()

[beanie00]

For Process Reward + LangGraph, I’m not familiar with LangGraph specifically, but it seems possible to call emit_reward inside a node. However, you should verify that the reward span is being saved properly for example, by using:
spans = await self.store.query_spans(rollout.rollout_id, attempt_id="latest")

Since multiple nodes are involved, you may also need to link the corresponding call span and reward span together. If you can obtain a unique response ID from self.invoke_prompt(prompt), then you can associate the call span and reward span by calling:
emit_reward(reward_1, attributes=make_link_attributes({"response_id": unique_response_id}))
You can then create the link between spans later.

[dalek-who]

@beanie00 Is there any example I can follow?

[beanie00]

I am currently working on this PR: #367. It still needs to be cleaned up because it includes a large amount of simulation-environment logic. I plan to clean up the PR either today or by tomorrow.

In the meantime, you can refer to these three files in PR: #367:

• examples/simulation/simulation_agent.py
• agentlightning/adapter/triplet.py
• agentlightning/verl/daemon.py

I used the Autogen agent and the VERL trainer, and since it’s a multi-turn setting, I connected the spans based on the step count. I’ll let you know when the PR is ready.

[ultmaster]

I don't see any issues to call emit_reward inside node functions.

You need to cautious that if your node is auto-retried or whatever, you might get duplicated rewards.