Spaces:
Sleeping
Sleeping
File size: 1,810 Bytes
149eeaf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
import ray
from datetime import datetime
from agent_state_actor import AgentState
class EnvironmentState:
def __init__(self, episode, step):
self.agent_state = None
self.timestamp = datetime.utcnow()
self.episode = episode
self.step = step
self.reward = 0
self.llm_preview = ''
self.llm_responses = []
self.tts_raw_chunk_ids = []
def __str__(self):
state = ', '.join(f'{k}={v}' for k, v in self.__dict__.items() if k not in {'episode', 'step', 'timestamp', 'reward'})
return f'episode={self.episode}, step={self.step}, timestamp={self.timestamp}, \nreward={self.reward}\nstate=({state})'
@ray.remote
class EnvironmentStateActor:
def __init__(self):
self.episode = 0
self.step = 0
self.state = None
self.reset_episode()
def reset_episode(self):
self.episode += 1
self.step = 0
self.state = EnvironmentState(self.episode, self.step)
return self.state
def begin_next_step(self)->EnvironmentState:
previous_state = self.state
self.step += 1
self.state = EnvironmentState(self.episode, self.step)
return previous_state
def add_reward(self, reward):
self.state.reward += reward
def set_llm_preview(self, llm_preview):
self.state.llm_preview = llm_preview
def add_llm_response_and_clear_llm_preview(self, llm_response):
self.state.llm_responses.append(llm_response)
self.state.llm_preview = ''
def add_tts_raw_chunk_id(self, chunk_id):
self.state.tts_raw_chunk_ids.append(chunk_id)
def add_agent_state(self, agent_state:AgentState):
self.state.agent_state = agent_state
def get_state(self)->EnvironmentState:
return self.state |