atari visualization with Gradio

.gitignore

@@ -186,4 +186,7 @@ main_jarvis.sh
 test*.py
 *.zip
 test_
-*.ipynb
+*.ipynb
+
+# gradio
+flagged

deciders/parser.py

@@ -7,10 +7,10 @@ class DisActionModel(BaseModel):
     @classmethod
     def create_validator(cls, max_action):
         @validator('action', allow_reuse=True)
-        def action_is_valid(cls, field):
-            if field not in range(1, max_action + 1):
+        def action_is_valid(cls, info):
+            if info not in range(1, max_action + 1):
                 raise ValueError(f"Action is not valid ([1, {max_action}])!")
-            return field
+            return info
         return action_is_valid
 
 # Generate classes dynamically

environment.yaml

@@ -186,4 +186,5 @@ dependencies:
       - win32-setctime==1.1.0
       - yarl==1.9.2
       - zipp==3.15.0
-      - git+ssh://git@github.com/hyyh28/atari-representation-learning.git
+      - git+ssh://git@github.com/hyyh28/atari-representation-learning.git
+      - gradio==4.13.0

gradio_reflexion.py

@@ -0,0 +1,312 @@
+import envs
+import deciders
+import distillers
+import prompts as task_prompts
+import datetime
+import time
+from envs.translator import InitSummarizer, CurrSummarizer, FutureSummarizer, Translator
+import gym
+import pandas as pd
+import random
+import datetime
+from loguru import logger
+from argparse import Namespace
+import gradio as gr
+
+
+def set_seed(seed):
+    random.seed(seed)         
+
+def main_progress(env_name, decider, prompt_level, num_trails, seed):
+    init_summarizer = env_name.split("-")[0] + '_init_translator'
+    curr_summarizer = env_name.split("-")[0] + '_basic_translator'
+    args = Namespace(
+        env_name=env_name,
+        init_summarizer=init_summarizer,
+        curr_summarizer=curr_summarizer,
+        decider=decider,
+        prompt_level=prompt_level,
+        num_trails=num_trails,
+        seed=seed,
+        future_summarizer=None,
+        env="base_env",
+        gpt_version="gpt-3.5-turbo",
+        render="rgb_array",
+        max_episode_len=200,
+        max_query_tokens=5000,
+        max_tokens=2000,
+        distiller="traj_distiller",
+        prompt_path=None,
+        use_short_mem=1,
+        short_mem_num=10,
+        is_only_local_obs=1,
+        api_type="azure",
+    )
+
+    if args.api_type != "azure" and args.api_type != "openai":
+        raise ValueError(f"The {args.api_type} is not supported, please use 'azure' or 'openai' !")
+    
+    # Please note when using "azure", the model name is gpt-35-turbo while using "openai", the model name is "gpt-3.5-turbo"
+    if args.api_type == "azure":
+        if args.gpt_version == "gpt-3.5-turbo":
+            args.gpt_version = 'gpt-35-turbo'
+    elif args.api_type == "openai":
+        if args.gpt_version == "gpt-35-turbo":
+            args.gpt_version = 'gpt-3.5-turbo'
+
+    # Get the specified translator, environment, and ChatGPT model
+    env_class = envs.REGISTRY[args.env]
+    init_summarizer = InitSummarizer(envs.REGISTRY[args.init_summarizer], args)
+    curr_summarizer = CurrSummarizer(envs.REGISTRY[args.curr_summarizer])
+    
+    if args.future_summarizer:
+        future_summarizer = FutureSummarizer(
+            envs.REGISTRY[args.future_summarizer],
+            envs.REGISTRY["cart_policies"],
+            future_horizon=args.future_horizon,
+        )
+    else:
+        future_summarizer = None
+
+    decider_class = deciders.REGISTRY[args.decider]
+    distiller_class = distillers.REGISTRY[args.distiller]
+    sampling_env = envs.REGISTRY["sampling_wrapper"](gym.make(args.env_name))
+    if args.prompt_level == 5:
+        prompts_class = task_prompts.REGISTRY[(args.env_name,args.decider)]()
+    else:
+        prompts_class = task_prompts.REGISTRY[(args.decider)]()
+    translator = Translator(
+        init_summarizer, curr_summarizer, future_summarizer, env=sampling_env
+    )
+    environment = env_class(
+        gym.make(args.env_name, render_mode=args.render), translator
+    )
+
+    logfile = (
+        f"llm.log/output-{args.env_name}-{args.decider}-{args.gpt_version}-l{args.prompt_level}"
+        f"-{datetime.datetime.now().timestamp()}.log"
+    )
+
+    logfile_reflexion = (
+        f"llm.log/memory-{args.env_name}-{args.decider}-{args.gpt_version}-l{args.prompt_level}"
+        f"-{datetime.datetime.now().timestamp()}.log"
+    )
+    my_distiller = distiller_class(logfile=logfile_reflexion,args=args)
+
+    args.game_description = environment.game_description
+    args.goal_description = environment.goal_description
+    args.action_description = environment.action_description
+    args.action_desc_dict = environment.action_desc_dict
+    args.reward_desc_dict = environment.reward_desc_dict
+
+    logger.add(logfile, colorize=True, enqueue=True, filter=lambda x: '[Reflexion Memory]' not in x['message'])
+
+    decider = decider_class(environment.env.action_space, args, prompts_class, my_distiller, temperature=0.0, logger=logger, max_tokens=args.max_tokens)
+    
+    # Evaluate the translator
+    utilities = []
+    df = pd.read_csv('record_reflexion.csv', sep=',')
+    filtered_df = df[(df['env'] == args.env_name) & (df['decider'] == 'expert') & (df['level'] == 1)]
+    expert_score = filtered_df['avg_score'].item()
+    seeds = [i for i in range(1000)]
+    # prompt_file = "prompt.txt"
+    # f = open(prompt_file,"w+")
+    num_trails = args.num_trails
+    if not "Blackjack" in args.env_name:
+        curriculums = 1
+    else:
+        curriculums = 20
+    for curriculum in range(curriculums):
+        for trail in range(num_trails): 
+            if "Blackjack" in args.env_name:
+                seed = seeds[curriculum*curriculums + num_trails - trail - 1]
+            else:
+                seed = args.seed
+
+            # single run
+            # Reset the environment
+            if not "Blackjack" in args.env_name:
+                set_seed(args.seed)
+                seed = args.seed
+                # Reset the environment
+                state_description, env_info = environment.reset(seed=args.seed)
+            else:
+                set_seed(seed)
+                # Reset the environment
+                state_description, env_info = environment.reset(seed=seed)
+            game_description = environment.get_game_description()
+            goal_description = environment.get_goal_description()
+            action_description = environment.get_action_description()
+
+            # Initialize the statistics
+            frames = []
+            utility = 0
+            current_total_tokens = 0
+            current_total_cost = 0
+            start_time = datetime.datetime.now()
+            # Run the game for a maximum number of steps
+            for round in range(args.max_episode_len):
+                # Keep asking ChatGPT for an action until it provides a valid one
+                error_flag = True
+                retry_num = 1
+                for error_i in range(retry_num):
+                    try:
+                        action, prompt, response, tokens, cost = decider.act(
+                            state_description,
+                            action_description,
+                            env_info,
+                            game_description,
+                            goal_description,
+                            logfile
+                        )
+
+                        state_description, reward, termination, truncation, env_info = environment.step_llm(
+                            action
+                        )
+                        if "Cliff" in args.env_name or "Frozen" in args.env_name:
+                            decider.env_history.add('reward', env_info['potential_state'] + environment.reward_desc_dict[reward])
+                        else:
+                            decider.env_history.add('reward', f"The player get rewards {reward}.")
+                            
+                        utility += reward
+
+                        # Update the statistics
+                        current_total_tokens += tokens
+                        current_total_cost += cost
+                        error_flag = False
+                        break
+                    except Exception as e:
+                        print(e)
+                        if error_i < retry_num-1:
+                            if "Cliff" in args.env_name or "Frozen" in args.env_name:
+                                decider.env_history.remove_invalid_state()
+                            decider.env_history.remove_invalid_state()
+                        if logger:
+                            logger.debug(f"Error: {e}, Retry! ({error_i+1}/{retry_num})")
+                        continue
+                if error_flag:
+                    action = decider.default_action
+                    state_description, reward, termination, truncation, env_info = environment.step_llm(
+                            action
+                        )
+
+                    decider.env_history.add('action', decider.default_action)
+
+                    if "Cliff" in args.env_name or "Frozen" in args.env_name:
+                        # decider.env_history.add('reward', reward)
+                        decider.env_history.add('reward', env_info['potential_state'] + environment.reward_desc_dict[reward])
+                    utility += reward
+
+                    
+                    logger.info(f"Seed: {seed}")
+                    logger.info(f'The optimal action is: {decider.default_action}.')
+                    logger.info(f"Now it is round {round}.")
+                else:
+                    current_total_tokens += tokens
+                    current_total_cost += cost
+                    logger.info(f"Seed: {seed}")
+                    logger.info(f"current_total_tokens: {current_total_tokens}")
+                    logger.info(f"current_total_cost: {current_total_cost}")
+                    logger.info(f"Now it is round {round}.")
+
+                # return results
+                yield environment.render(), state_description, prompt, response, action
+
+                if termination or truncation:
+                    if logger:
+                        logger.info(f"Terminated!")
+                    break
+                time.sleep(10)
+            decider.env_history.add(
+                'terminate_state', environment.get_terminate_state(round+1, args.max_episode_len))
+            decider.env_history.add("cummulative_reward", str(utility))
+            # Record the final reward
+            if logger:
+                logger.info(f"Cummulative reward: {utility}.")
+                end_time = datetime.datetime.now()
+                time_diff = end_time - start_time
+                logger.info(f"Time consumer: {time_diff.total_seconds()} s")
+            
+            utilities.append(utility)
+            # TODO: set env sucess utility threshold
+            if trail < num_trails -1:
+                if args.decider in ['reflexion']:
+                    if utility < expert_score: 
+                        decider.update_mem() 
+                else:
+                    decider.update_mem() 
+        decider.clear_mem()
+    return utilities
+
+# def pause():
+#     for i in range(31415926):
+#         time.sleep(0.1)
+#         yield i            
+
+if __name__ == "__main__":
+    custom_css = """
+        #render {
+            flex-grow: 1;
+        }
+        #input_text .tabs {
+            display: flex;
+            flex-direction: column;
+            flex-grow: 1;
+        }
+        #input_text .tabitem[style="display: block;"] {
+            flex-grow: 1;
+            display: flex !important;
+        }
+        #input_text .gap {
+            flex-grow: 1;
+        }
+        #input_text .form {
+            flex-grow: 1 !important;
+        }
+        #input_text .form > :last-child{
+            flex-grow: 1;
+        }
+    """
+
+    with gr.Blocks(theme=gr.themes.Monochrome(), css=custom_css) as demo:
+        with gr.Row():
+            env_name = gr.Dropdown(
+                ["RepresentedBoxing-v0", 
+                 "RepresentedPong-v0", 
+                 "RepresentedMsPacman-v0", 
+                 "RepresentedMontezumaRevenge-v0"], 
+                label="Environment Name")
+            decider = gr.Dropdown(
+                ["naive_actor", 
+                 "cot_actor", 
+                 "spp_actor", 
+                 "reflexion_actor"], 
+                 label="Decider")
+            prompt_level = gr.Dropdown([1, 2, 3, 4, 5], label="Prompt Level")
+        with gr.Row():
+            num_trails = gr.Slider(1, 100, 1, label="Number of Trails", scale=2)
+            seed = gr.Slider(1, 1000, 1, label="Seed", scale=2)
+            run = gr.Button("Run", scale=1)
+            # pause_ = gr.Button("Pause")
+            # resume = gr.Button("Resume")
+            stop = gr.Button("Stop", scale=1)
+        with gr.Row():
+            with gr.Column():
+                render = gr.Image(label="render", elem_id="render")
+            with gr.Column(elem_id="input_text"):
+                state = gr.Textbox(label="translated state")
+                prompt = gr.Textbox(label="prompt", max_lines=100)
+        with gr.Row():
+            response = gr.Textbox(label="response")
+            action = gr.Textbox(label="parsed action")
+        run_event = run.click(
+            fn=main_progress, 
+            inputs=[env_name, decider, prompt_level, num_trails, seed], 
+            outputs=[render, state, prompt, response, action])
+        stop.click(fn=None, inputs=None, outputs=None, cancels=[run_event])
+        # pause_event = pause_.click(fn=pause, inputs=None, outputs=None)
+        # resume.click(fn=None, inputs=None, outputs=None, cancels=[pause_event])
+
+    demo.launch(server_name="0.0.0.0", server_port=7860)
+
+