gen_few_shots_examples.py

import argparse
import envs
import deciders
from matplotlib import animation
import matplotlib.pyplot as plt
import os
import numpy as np
import torch as th
from envs.translator import InitSummarizer, CurrSummarizer, FutureSummarizer, Translator
from tianshou.data import Collector, VectorReplayBuffer, ReplayBuffer
from tianshou.policy import PPOPolicy
from RL_based.utils import (
    Net_GRU_Bert_tianshou,
    Net_Bert_CLS_tianshou,
    Net_Bert_CNN_tianshou,
    Net_GRU_nn_emb_tianshou,
)
from tianshou.utils.net.common import ActorCritic
from tianshou.utils.net.discrete import Actor, Critic
import gym
import json

ENV_CLASS = {'classic_control': ['CartPole', 'Acrobot', 'MountainCar'],
             'box2d': ['LunarLander'],
             'toy_text': ['Blackjack', 'Taxi', 'CliffWalking', 'FrozenLake']}

def get_env_class(env_name):
    for key, value in ENV_CLASS.items():
        if env_name in value:
            return key
    return None

def get_fewshot_example_path(env, decider):
    assert decider in ['random_actor', 'expert'], "decider must be random_actor or expert"
    prompt_level = 2 if decider == 'random_actor' else 4
    fewshot_example_path = os.path.join(
        'envs', get_env_class(env.spec.name), 'few_shot_examples',
        ''.join([env.spec.name.lower(), '_l', str(prompt_level), '.json']))
    return fewshot_example_path

# https://colab.research.google.com/drive/1DdWsGi10232orUv-reY4wsTmT0VMoHaX?usp=sharing#scrollTo=4OfVmDKk7XvG
# LLMs bias on 0 so make the actions greater than 1 instead.

def gen_expert_examples(environment, policy, file_path, max_episode_len=120, n_episodes=1):
    replaybuffer = ReplayBuffer(size=1000)
    test_collector_1 = Collector(policy, environment, replaybuffer)
    test_collector_1.reset_env()
    game_description = environment.get_game_description()
    goal_description = environment.get_goal_description()
    action_description = environment.get_action_description()
    policy.eval()
    data_lst = []

    for _ in range(n_episodes):
        test_collector_1.reset_buffer()
        result = test_collector_1.collect(n_episode=1)
        sample_result = replaybuffer.sample(0)[0]
        round = 0 
        utility = 0
        data = []
        for transition in sample_result:
            round += 1
            if round > max_episode_len:
                break
            question = f"{transition.obs} \n {goal_description} \n {action_description} "
            reward = transition.rew
            utility += reward

            answer = f"The final answer is: {transition.act + 1}"

            data.append(
                {   
                    "observation": transition.obs, 
                    "goal_description": goal_description, 
                    "action_description": action_description, 
                    "game_description": game_description,
                    "action": str(transition.act + 1),
                    "question": question,
                    "answer": answer,
                    "reward": reward,
                    "cum_reward": utility,
                }
            )
            print(f"Now it is round {round}")
        data_lst.append(data)
    # Return the final reward
    with open(file_path, "w") as outfile:
        json.dump(data_lst, outfile)
    return utility


def gen_examples(environment, decider, file_path, max_episode_len=200, n_episodes=1):
    game_description = environment.get_game_description()
    goal_description = environment.get_goal_description()
    action_description = environment.get_action_description()
    frames = []
    utilities = []
    data_lst = []

    for _ in range(n_episodes):
        # Reset the environment
        round = 0
        state_description, env_info = environment.reset()
        utility = 0
        data = []
        for _ in range(max_episode_len):
            # Keep asking ChatGPT for an action until it provides a valid one  
            asking_round = 0
            action, prompt, answer, _, _, _ = decider.act(
                state_description,
                action_description,
                env_info,
                game_description,
                goal_description,
            )
            # Perform the action in the environment
            state_description, reward, terminated, truncated, env_info = environment.step_llm(
                action
            )
            question = f"{state_description} \n {goal_description} \n {action_description} "
            utility += reward
            answer += f"The final answer is: {action}"

            data.append(
                {
                    "observation": state_description, 
                    "goal_description": goal_description, 
                    "action_description": action_description, 
                    "game_description": game_description,
                    "action": action,
                    "question": question,
                    "answer": answer,
                    "reward": reward,
                    "cum_reward": utility,
                }
            )
            print(f"Now it is round {round}")
            round += 1
            # If the game is over, break the loop
            if terminated or truncated:
                print(f"Terminated!")
                break
        utilities.append(utility)
        data_lst.append(data)
    # Return the final reward
    with open(file_path, "w") as outfile:
        json.dump(data_lst, outfile)
    return utility


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Generate few shots examples of a gym environment."
    )
    parser.add_argument(
        "--init_summarizer",
        type=str,
        required=True,
        help="The name of the init summarizer to use.",
    )
    parser.add_argument(
        "--curr_summarizer",
        type=str,
        required=True,
        help="The name of the curr summarizer to use.",
    )
    parser.add_argument(
        "--env",
        type=str,
        default="base_env",
        help="The name of the gym environment to use.",
    )
    parser.add_argument(
        "--decider",
        type=str,
        default="naive_actor",
        help="The actor used to select action",
    )
    parser.add_argument(
        "--env_name",
        type=str,
        default="CartPole-v0",
        help="The name of the gym environment to use.",
    )
    parser.add_argument(
        "--max_episode_len",
        type=int,
        default=200,
        help="The maximum number of steps in an episode.",
    )
    parser.add_argument(
        "--num_episodes",
        type=int,
        default=1,
        help="The number of episodes to collect data.",
    )
    parser.add_argument(
        "--max_length",
        type=int,
        default=128,
        help="The token length of the observation",
    )
    parser.add_argument(
        "--trans_model_name",
        type=str,
        default="/home/ubuntu/LLM-Decider-Bench/RL_based/transformer_offline_distilbert",
        help="The name of the pretrained transformer to use.",
    )
    parser.add_argument(
        "--policy_path",
        type=str,
        default=None,
        help="The path to the policy to be evaluated",
    )
    parser.add_argument(
        "--n_episodes",
        type=int,
        default=2,
        help="The number of episodes to collect data (for env where episode is too short).",
    )

    args = parser.parse_args()
    # Get the specified translator, environment, and ChatGPT model
    device = "cuda" if th.cuda.is_available() else "cpu"
    env_class = envs.REGISTRY[args.env]
    init_summarizer = InitSummarizer(envs.REGISTRY[args.init_summarizer])
    curr_summarizer = CurrSummarizer(envs.REGISTRY[args.curr_summarizer])
    translator = Translator(init_summarizer, curr_summarizer, None, env=None)
    environment = env_class(gym.make(args.env_name, render_mode=None), translator)

    fewshot_example_path = get_fewshot_example_path(environment, args.decider)

    if args.decider == "expert":
        net = Net_GRU_nn_emb_tianshou(
            hidden_sizes=[256, 128],
            device=device,
            max_length=args.max_length,
            trans_model_name=args.trans_model_name,
        )
        actor = Actor(net, environment.action_space.n, device=device).to(device)
        critic = Critic(net, device=device).to(device)
        actor_critic = ActorCritic(actor, critic)
        optim = th.optim.Adam(actor_critic.parameters(), lr=0.0003)

        # PPO policy
        dist = th.distributions.Categorical
        policy = PPOPolicy(
            actor,
            critic,
            optim,
            dist,
            action_space=environment.action_space,
            deterministic_eval=True,
        )
        policy.load_state_dict(th.load(args.policy_path))
        utility = gen_expert_examples(
            environment, policy, fewshot_example_path,
            max_episode_len=args.max_episode_len, n_episodes=args.n_episodes
        )
    else:
        decider_class = deciders.REGISTRY[args.decider]
        decider = decider_class(environment.env.action_space)
        # Evaluate the translator
        utility = gen_examples(
            environment, decider, fewshot_example_path,
            max_episode_len=args.max_episode_len, 
            n_episodes=args.n_episodes
        )
    print(f"(Avg.) Cummulative reward: {utility}")