envs/base_env.py

# 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):
        super().__init__(env)
        self.env = env

    def set_state(self, state):
        self.env.state = state
        self.env.steps_beyond_terminated = None

class BaseEnv(gym.Wrapper):
    def __init__(self, env, translator):
        super().__init__(env)
        self.translator = translator
        self.env_name = super().spec.id
        self.transition_data = {}
        self.game_description = self.get_game_description()
        self.goal_description = self.get_goal_description()
        self.action_description = self.get_action_description()
        self.action_desc_dict = self.get_action_desc_dict()
        self.reward_desc_dict = self.get_reward_desc_dict()

    def reset(self, **kwargs):
        state, _ = super().reset(**kwargs)
        self.transition_data['state'] = state
        self.translator.obtain(self.transition_data)
        summary, future_summary = self.translator.translate()
        info = {
            'future_summary': future_summary
        }
        self.state = state
        return summary, info

    def step(self, action): 
        potential_next_state = self.get_potential_next_state(action)
        state, reward, terminated, _, info = super().step(action)
        self.transition_data['action'] = action
        self.transition_data['next_state'] = state
        self.transition_data['reward'] = reward 
        self.transition_data['terminated'] = terminated
        self.translator.update(self.transition_data)
        self.transition_data = {}
        self.transition_data['state'] = state
        self.translator.obtain(self.transition_data)
        summary, future_summary = self.translator.translate()
        info = {
            'future_summary': future_summary,
            'potential_state': potential_next_state
        }
        return summary, reward, terminated, _, info


    def step_llm(self, action): 
        potential_next_state = self.get_potential_next_state(action)
        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 
        self.transition_data['terminated'] = terminated
        self.translator.update(self.transition_data)
        self.transition_data = {}
        self.transition_data['state'] = state
        self.translator.obtain(self.transition_data)
        self.state = state
        summary, future_summary = self.translator.translate()
        info = {
            'future_summary': future_summary,
            'potential_state': potential_next_state,
        }
        return summary, reward, terminated, _, info

    def get_terminate_state(self, episode_len, max_episode_len):
        return self.translator.translate_terminate_state(self.state, episode_len, max_episode_len)

    def get_game_description(self,):
        return self.translator.describe_game()

    def get_goal_description(self,):
        return self.translator.describe_goal()
    
    def get_action_description(self,):
        return self.translator.describe_action()

    def get_action_desc_dict(self,):
        return self.translator.get_action_desc_dict()

    def get_reward_desc_dict(self,):
        return self.translator.get_reward_desc_dict()

    def get_potential_next_state(self, action):
        return self.translator.translate_potential_next_state(self.state, action)