Merge pull request #13 from mail-ecnu/parser-new

deciders/act.py

@@ -3,7 +3,7 @@
 import openai
 from .gpt import gpt 
 from loguru import logger
-from .parser import PARSERS
+from .parser import DISPARSERS, CONPARSERS
 from langchain.output_parsers import PydanticOutputParser
 from langchain.output_parsers import OutputFixingParser
 from langchain.chat_models import AzureChatOpenAI, ChatOpenAI
@@ -12,13 +12,18 @@ import tiktoken
 import json
 import re
 from .utils import run_chain
+from gym.spaces import Discrete
 
 class RandomAct():
     def __init__(self, action_space):
         self.action_space = action_space
     
     def act(self, state_description, action_description, env_info, game_description=None, goal_description=None):
-        return self.action_space.sample()+1, '', '', '', 0, 0
+        if isinstance(self.action_space, Discrete):
+            action = self.action_space.sample()+1
+        else:
+            action = self.action_space.sample()
+        return action, '', '', '', 0, 0
 
 class NaiveAct(gpt):
     def __init__(self, action_space, args, prompts, distiller, temperature=0.0, max_tokens=2048, logger=None):
@@ -37,7 +42,10 @@ class NaiveAct(gpt):
         super().__init__(args)
         self.distiller = distiller
         self.fewshot_example_initialization(args.prompt_level, args.prompt_path, distiller = self.distiller)
-        self.default_action = 1
+        if isinstance(self.action_space, Discrete):
+            self.default_action = 1
+        else:
+            self.default_action = [0 for ind in range(self.action_space.shape[0])]
         self.parser = self._parser_initialization()
         self.irr_game_description = ''
         self.memory = []
@@ -82,11 +90,12 @@ class NaiveAct(gpt):
 
 
     def _parser_initialization(self):
-        if hasattr(self.action_space, 'n'):
-            assert self.action_space.n in PARSERS.keys(), f'Action space {self.action_space} is not supported.'
+        if isinstance(self.action_space, Discrete): 
+            PARSERS = DISPARSERS
             num_action = self.action_space.n
-        else:
-            num_action = 1
+        else: 
+            PARSERS = CONPARSERS
+            num_action = self.action_space.shape[0]
 
         if self.args.api_type == "azure":
             autofixing_chat = AzureChatOpenAI(
@@ -200,42 +209,16 @@ class NaiveAct(gpt):
                 my_mem += f"\nBelow are the latest {min(self.mem_num, len(self.env_history))} historical data entries:\n"
                 my_mem += f"{self.env_history.get_histories(self.mem_num)}"
 
-        while asking_round < 3: 
-            prompt, res = self.response(state_description, action_description, env_info, game_description, goal_description, my_mem)
-            action_str = res.choices[0].text.strip()
-            print(f'my anwser is {action_str}')
-            # import pdb; pdb.set_trace()
-            try: 
-                if "Continuous" in self.args.env_name:
-                    action = float(re.findall(r"[-+]?\d*\.\d+", action_str)[0])
-                    
-                else:
-                    action = int(re.findall(r"\d+", action_str)[0])
-            except: 
-                action = None
-                asking_round += 1
-                continue
-
-            if "Continuous" not in self.args.env_name:
-                if (action-1) in self.action_space:
-                    break
-                else:
-                    asking_round += 1
-                    action = None
-            else:
-                if action >= self.action_space.low and action <= self.action_space.high:
-                    break
-                else:
-                    asking_round += 1
-                    action = None
-                
-        if action is None: 
-            print('err on selecting action')
-            action = self.default_action
+        
+        prompt, res = self.response(state_description, action_description, env_info, game_description, goal_description, my_mem)
+        action_str = res.choices[0].text.strip()
+        print(f'my anwser is {action_str}')
+        action = self.parser.parse(response).action
         self._add_history_after_action(action)
-        self.logger.info(f'\n{prompt}')
-        self.logger.info(f'The GPT response is: {res}.')
+        self.logger.info(f'The GPT response is: {response}.')
         self.logger.info(f'The optimal action is: {action}.')
+        if env_info.get('history'):
+            self.logger.info(f'History: {history_to_str(env_info["history"])}')
         return action, prompt, res, 0, 0
 
     def _read_mem(self, ):

deciders/parser.py

@@ -1,72 +1,54 @@
 from pydantic import BaseModel, Field, validator
+from typing import List
 
-# Define your desired data structure.
-class TwoAction(BaseModel):
-    action: int = Field(description="the choosed action to perform")
-    
-    # You can add custom validation logic easily with Pydantic.
-    @validator('action')
-    def action_is_valid(cls, field):
-        if field not in [1, 2]:
-            raise ValueError("Action is not valid ([1, 2])!")
-        return field
-    
-class ThreeAction(BaseModel):
-    action: int = Field(description="the choosed action to perform")
-    
-    # You can add custom validation logic easily with Pydantic.
-    @validator('action')
-    def action_is_valid(cls, field):
-        if field not in [1, 2, 3]:
-            raise ValueError("Action is not valid ([1, 2, 3])!")
-        return field
-    
-class FourAction(BaseModel):
-    action: int = Field(description="the choosed action to perform")
-    
-    # You can add custom validation logic easily with Pydantic.
-    @validator('action')
-    def action_is_valid(cls, field):
-        if field not in [1, 2, 3, 4]:
-            raise ValueError("Action is not valid ([1, 2, 3, 4])!")
-        return field
-    
-class SixAction(BaseModel):
-    action: int = Field(description="the choosed action to perform")
+class DisActionModel(BaseModel):
+    action: int = Field(description="the chosen action to perform")
+
+    @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):
+                raise ValueError(f"Action is not valid ([1, {max_action}])!")
+            return field
+        return action_is_valid
+
+# Generate classes dynamically
+def generate_action_class(max_action):
+    return type(f"{max_action}Action", (DisActionModel,), {'action_is_valid': DisActionModel.create_validator(max_action)})
     
-    # You can add custom validation logic easily with Pydantic.
-    @validator('action')
-    def action_is_valid(cls, field):
-        if field not in [1, 2, 3, 4, 5, 6]:
-            raise ValueError("Action is not valid ([1, 2, 3, 4, 5, 6])!")
-        return field
+# Dictionary of parsers with dynamic class generation
+DISPARSERS = {num: generate_action_class(num) for num in [2, 3, 4, 6, 9, 18]}
 
+class ContinuousActionBase(BaseModel):
+    action: List[float] = Field(description="the chosen continuous actions to perform")
 
-class NineAction(BaseModel):
-    action: int = Field(description="the choosed action to perform")
+    @classmethod
+    def set_expected_length(cls, length):
+        cls.expected_length = length
 
-    # You can add custom validation logic easily with Pydantic.
-    @validator('action')
-    def action_is_valid(cls, field):
-        if field not in [1, 2, 3, 4, 5, 6, 7, 8, 9]:
-            raise ValueError("Action is not valid ([1, 2, 3, 4, 5, 6, 7, 8, 9])!")
-        return field
+    @validator('action', pre=True)
+    def validate_length(cls, action):
+        if len(action) != cls.expected_length:
+            raise ValueError(f"The action list must have exactly {cls.expected_length} items.")
+        return action
 
-class FullAtariAction(BaseModel):
-    action: int = Field(description="the choosed action to perform")
-    @validator('action')
-    def action_is_valid(cls, info):
-        if info not in [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18]:
-            raise ValueError("Action is not valid ([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18])!")
-        return info
-    
-class ContinuousAction(BaseModel):
-    action: float = Field(description="the choosed action to perform")
-    # You can add custom validation logic easily with Pydantic.
-    @validator('action')
-    def action_is_valid(cls, field):
-        if not (field >= -1 and field <= 1):
-            raise ValueError("Action is not valid ([-1,1])!")
-        return field
-    
-PARSERS = {1:ContinuousAction, 2: TwoAction, 3: ThreeAction, 4: FourAction, 6: SixAction, 9:NineAction, 18: FullAtariAction}
+    @validator('action', each_item=True)
+    def action_is_valid(cls, item):
+        if not -1 <= item <= 1:
+            raise ValueError("Each action dimension must be in the range [-1, 1]!")
+        return item
+
+# Generate classes dynamically
+def generate_continuous_action_class(expected_length):
+    NewClass = type(
+        f"{expected_length}DContinuousAction",
+        (ContinuousActionBase,),
+        {}
+    )
+    NewClass.set_expected_length(expected_length)
+    return NewClass
+
+
+# Dictionary of parsers with dynamic class generation
+CONPARSERS = {length: generate_continuous_action_class(length) for length in range(1, 17)}

envs/base_env.py

@@ -1,6 +1,7 @@
 # This file contains functions for interacting with the CartPole environment
 
 import gym
+from gym.spaces import Discrete
 
 class SettableStateEnv(gym.Wrapper):
     def __init__(self, env):
@@ -55,10 +56,11 @@ class BaseEnv(gym.Wrapper):
 
     def step_llm(self, action): 
         potential_next_state = self.get_potential_next_state(action)
-        if "Continuous" in self.env_name:
-            state, reward, terminated, _, info = super().step(action)
-        else:
+        if isinstance(self.action_space, Discrete):
             state, reward, terminated, _, info = super().step(action-1)
+        else:
+            state, reward, terminated, _, info = super().step(action)
+            
         self.transition_data['action'] = action
         self.transition_data['next_state'] = state
         self.transition_data['reward'] = reward 

main_reflexion.py

@@ -17,6 +17,9 @@ import random
 import numpy as np
 import datetime
 from loguru import logger
+from gym.spaces import Discrete
+
+
 
 
 def set_seed(seed):
@@ -109,9 +112,6 @@ def _run(translator, environment, decider, max_episode_len, logfile, args, trail
                     logfile
                 )
 
-                if "Continuous" in args.env_name:
-                    action = [action]
-
                 state_description, reward, termination, truncation, env_info = environment.step_llm(
                     action
                 )
@@ -137,10 +137,7 @@ def _run(translator, environment, decider, max_episode_len, logfile, args, trail
                     logger.debug(f"Error: {e}, Retry! ({error_i+1}/{retry_num})")
                 continue
         if error_flag:
-            if "Continuous" in args.env_name:
-                action = [decider.default_action]
-            else:
-                action = decider.default_action
+            action = decider.default_action
             state_description, reward, termination, truncation, env_info = environment.step_llm(
                     action
                 )
@@ -164,7 +161,7 @@ def _run(translator, environment, decider, max_episode_len, logfile, args, trail
             logger.info(f"current_total_cost: {current_total_cost}")
             logger.info(f"Now it is round {round}.")
 
-        frames.append(environment.render())
+        # frames.append(environment.render())
         if termination or truncation:
             if logger:
                 logger.info(f"Terminated!")

record_reflexion.csv

@@ -10,4 +10,5 @@ FrozenLake-v1,1,expert,200.0
 MountainCarContinuous-v0,1,expert,200.0
 RepresentedBoxing-v0,1,expert,200.0
 RepresentedPong-v0,1,expert,200.0
+Ant-v4,1,expert,5000