LLM Agent의 효율적 탐색(Exploration)을 위한 PSRL 구현

# LLM-based PSRL 핵심 루프 (의사코드)

SYSTEM_POSTERIOR_UPDATER = """
You are a Bayesian posterior distribution for a sequential decision-making problem.
Given a current prior belief and trajectory observation, produce the updated posterior
that accurately reflects knowledge about environment transitions and rewards.
Never discard prior knowledge — only update it.
Environment: {env_description}
"""

SYSTEM_POSTERIOR_SAMPLER = """
Given the current posterior belief, generate ONE plausible hypothesis
(a concrete sample) for how transitions and rewards work in this environment.
Your sample must be consistent with the posterior constraints.
Start with 'You think' and provide only the hypothesis.
"""

SYSTEM_OPTIMAL_POLICY = """
Environment: {env_description}
Always select optimal actions that maximize value according to this hypothesis:
{posterior_sample}
Just say the action after 'Action:' and nothing else.
"""

def llm_psrl_episode(prior, env, llm_updater, llm_sampler, llm_policy):
    # Step 1: Sample one hypothesis from current posterior
    posterior_sample = llm_sampler.call(
        system=SYSTEM_POSTERIOR_SAMPLER,
        user=f"Current posterior: {prior}"
    )
    
    # Step 2: Act optimally w.r.t. sampled hypothesis for entire episode
    trajectory = []
    state = env.reset()
    for step in range(env.horizon):
        action = llm_policy.call(
            system=SYSTEM_OPTIMAL_POLICY.format(
                env_description=env.description,
                posterior_sample=posterior_sample
            ),
            user=f"Current state: {state}"
        )
        next_state, reward = env.step(action)
        trajectory.append((state, action, reward, next_state))
        state = next_state
    
    # Step 3: Update posterior with full trajectory
    new_posterior = llm_updater.call(
        system=SYSTEM_POSTERIOR_UPDATER.format(
            env_description=env.description
        ),
        user=f"Prior: {prior}\nTrajectory: {format_trajectory(trajectory)}"
    )
    
    return new_posterior, trajectory

# 메인 루프
prior = "Beta(1,1) prior for each arm — all actions equally likely to be optimal"
for episode in range(K):
    prior, traj = llm_psrl_episode(prior, env, updater_llm, sampler_llm, policy_llm)