Trace, the next AutoDiff

We tested Trace in a traffic control problem which is an instance of hyper-parameter tuning. We used UXSim to simulate traffic at a four-way intersection, where the trainable parameters are 2 integers in [15,90], which are the green light duration for each direction of traffic flow.

Image sourced from UXSim Website

The feedback is the estimated delay experienced by all vehicles due to intersections, and the goal of an optimizer is to minimize the delay using the fewest number of traffic simulations. To this end, this optimizer must find the right trade-off for temporally distributed and variable demands.

We report the performance of a SOTA heuristic from the traffic control literature, SCATS as well as two black-box optimization techniques: Gaussian Process Minimization (GP) and Particle Swarm Optimization (PSO). All methods use the same starting parameters.

GP and PSO appear bad because 50 iterations are insufficient for their convergence; given enough iterations, both will eventually perform well. Trace is quickly competitive with the SCATS heuristic, whereas OPRO is not. We show the code sketch below. Trace sends a node object into the simulator and let the environment operate on it. The underlying operation logic is automatically revealed to the Trace optimizer.

                                                    
    import trace

    def traffic_simulation(EW_green_time, NS_green_time):
        W = None
        try:
            W = create_world(EW_green_time, NS_green_time)
        except Exception as e:
            e_node = ExceptionNode(
                e,
                inputs={"EW_green_time": EW_green_time, "NS_green_time": NS_green_time},
                description="[exception] Simulation raises an exception with these inputs.",
                name="exception_step",
            )
            return e_node
        W.data.exec_simulation()
        return_dict = analyze_world(W, verbosity)

        return return_dict

    EW_x = trace.node(MIN_GREEN_TIME, trainable=True, constraint=f"[{MIN_GREEN_TIME},{MAX_GREEN_TIME}]")
    NS_x = trace.node(MIN_GREEN_TIME, trainable=True, constraint=f"[{MIN_GREEN_TIME},{MAX_GREEN_TIME}]")

    optimizer.objective = (
                "You should suggest values for the variables so that the OVERALL SCORE is as small as possible.\n"
                + "There is a trade-off in setting the green light durations.\n"
                + "If the green light duration for a given direction is set too low, then vehicles will queue up over time and experience delays, thereby lowering the score for the intersection.\n"
                + "If the green light duration for a given direction is set too high, vehicles in the other direction will queue up and experience delays, thereby lowering the score for the intersection.\n"
                + "The goal is to find a balance for each direction (East-West and North-South) that minimizes the overall score of the intersection.\n"
                + optimizer.default_objective
        )

    for i in range(num_iter):
        result = traffic_simulation(EW_x, NS_x)
        # some steps are skipped for simplicity
        feedback = result.data
        optimizer.zero_feedback()
        optimizer.backward(result, feedback, visualize=True)
        optimizer.step()
                                            

LLM agents today have many components. Most libraries provide optimization tools to optimize a small portion of their workflows, predominantly the prompt that goes into an LLM call. However, for building self-adapting agents that can modify their own behavior, only allowing the change to one part of a workflow but not others seems limiting.

In this experiment, we test Trace's ability in joint prompt optimization and code generation. Specifically, we optimize a given DSPy-based LLM agent and tunes its three components: the meta-prompt prompt_template, a function create_prompt that modifies the prompt with the current question, and a function extract_answer that post-processes the output of an LLM call. We use Big-Bench Hard (BBH) as the problem source. The typical setup of BBH evaluation is 3-shot. We instead choose the more challenging 0-shot setting. The 0-shot setup requires the agent to conform to the correct answer format without any example, challenging for any method that just directly prompts LLM, but not for an agent where a complete workflow is optimized.

We compare Trace with DSPy’s COPRO module (which optimizes the meta-prompt). In the Table below, we show that Trace is able to optimize a DSPy program beyond what DSPy’s COPRO optimizer can offer, especially on algorithmic tasks.

                                                
    import trace

    @trace.model
    class Predict(LLMCallable):
        def __init__(self):
            super().__init__()

            self.demos = []
            self.prompt_template = dedent("""
            Given the fields `question`, produce the fields `answer`.
            ---
            Follow the following format.

            Question:
            Answer:
            ---
            Question: {}
            Answer:""")
            self.prompt_template = trace.node(self.prompt_template, trainable=True)

        @trace.bundle(trainable=True)
        def extract_answer(self, prompt_template, question, response):
            """
            Need to read in the response, which can contain additional thought, delibration and an answer.
            Use code to process the response and find where the answer is.
            Can use self.call_llm("Return the answer from this text: " + response) again to refine the answer if necessary.

            Args:
                prompt_template: The prompt that was used to query LLM to get the response
                question: Question has a text describing the question but also "Options"
                response: LLM returned a string response
            """
            answer = response.split("Answer:")[1].strip()
            return answer

        @trace.bundle(trainable=True)
        def create_prompt(self, prompt_template, question):
            """
            The function takes in a question and then add to the prompt for LLM to answer.
            Args:
                prompt_template: some guidance/hints/suggestions for LLM
                question: the question for the LLM to answer
            """
            return prompt_template.format(question)

        def forward(self, question):
            """
            question: text

            We read in a question and produces a response
            """
            user_prompt = self.create_prompt(self.prompt_template, question)
            response = self.call_llm(user_prompt)
            answer = self.extract_answer(self.prompt_template, question, response)
            return answer
                                                
                                            

In this example, we want to learn a policy code for controlling a robotic manipulator. Compared with the previous Battleship example, the problem here has a longer horizon, since the policy would need to drive the robot for multiple time steps. Traditionally such a problem is framed as a reinforcement learning (RL) problem and usually learning a policy with RL requires tens of thousands of practice episodes. We show Trace can be used to effectively solve such a problem in just a dozen of episodes -- a 1000X speed up -- since it can end-to-end optimize the control system as opposed to treating the system like a black-box as RL does. We trace the steps of the entire practice episode and perform end-to-end update (using the same optimizer OPTO-Prime) through these steps. In this way, effectively, Trace performs back-propagation through time (BPTT).

We conduct experiments using a simulated Sawyer robot arm in the Meta-World environment of LLF-Bench. The agent policy needs to decide a target pose (end-effector position and the gripper state) for the robot, which will then be used as a set point for a low-level P controller, to perform a pick-and-place task. Each episode has 10 timesteps and tracing through the AI system’s rollout would result in a graph of depth around 30 for an episode. The agent receives intermediate language feedback as observations (from LLF-Bench) and finally feedback about success and return at the end of the episode in texts. Like the Battleship example, we initiate the policy code to be a dummy function and let it adapt through interactions.

We repetitively train the agent start from one initial condition and then test it on 10 new held-out initial conditions for generalization. Trace rapidly learns a robot controller in the MetaWorld simulated environment, that generalizes to new initial conditions. The video shows Trace learns a policy to successfully perform the pick-place task after 13 episodes.

Iteration 0 (Initial Policy)

Iteration 1 (Learned to reach goal but forgot to pick up the object)

Iteration 6 (Dropped the object too early but attempted to recover)

Iteration 13 (100% success rate on configurations unseen in training)

                                                
    import trace

    @trace.bundle(trainable=True)
    def controller(obs):
        """
        A feedback controller that computes the action based on the observation.

        Args:
            obs: (dict) The observation from the environment. Each key is a string (indicating a type of observation) and the value is a list of floats.
        Output:
            action: (list or nd.array) A 4-dimensional vector.
        """
        return [0, 0, 0, 0]

    def rollout(env, horizon, controller):
        """Rollout a controller in an env for horizon steps."""
        traj = dict(observation=[], action=[], reward=[], termination=[], truncation=[], success=[], input=[], info=[])

        # Initialize the environment
        obs, info = env.reset()
        traj["observation"].append(obs)

        # Rollout
        for t in range(horizon):
            controller_input = obs["observation"]
            error = None
            try:  # traced
                action = controller(controller_input)
                next_obs, reward, termination, truncation, info = env.step(action)
            except trace.ExecutionError as e:
                error = e
                break

            if error is None:
                # code skipped...logging
                if termination or truncation or info["success"]:
                    break
                obs = next_obs
        return traj, error

    optimizer = optimizer_cls(controller.parameters())
    env = TracedEnv(env_name, seed=seed, feedback_type=feedback_type, relative=relative)

    print("Optimization Starts")
    for i in range(n_optimization_steps):
        # Rollout and collect feedback
        traj, error = rollout(env, horizon, controller)
        feedback = construct_feedback(traj, error)

        # we provide a task-specific single-line hint for the optimizer
        optimizer.objective = hint + optimizer.default_objective

        # Optimization step
        optimizer.zero_feedback()
        optimizer.backward(target, feedback)
        optimizer.step()