refactor: use more of a MDP style structure

agent_state_actor.py

@@ -0,0 +1,34 @@
+from datetime import datetime
+import ray
+from copy import deepcopy
+
+class AgentState:
+    def __init__(self):
+        self.timestamp = datetime.utcnow()
+        self.input_stt_preview = ''
+        self.prompts = []
+        self.video_input = None
+
+@ray.remote
+class AgentStateActor:
+    def __init__(self):
+        self.begin_step()
+
+    def begin_step(self):
+        self.state = AgentState()
+        # self.state = deepcopy(self.state)
+        # self.state.timestamp = datetime.utcnow()
+
+    def set_input_stt_preview(self, input_stt_preview):
+        self.state.input_stt_preview = input_stt_preview
+
+    def add_input_stt_prompt(self, prompt):
+        self.state.prompts.append(prompt)
+
+    def add_video_input(self, video_input):
+        self.state.video_input = video_input
+
+    def get_state(self):
+        return self.state
+
+

app.py

@@ -53,7 +53,9 @@ async def main():
 
     with col1:
         listening = st.checkbox("Listen", value=True)
-        system_one_audio_history_output = st.empty()
+        looking = st.checkbox("Look", value=False)
+        charles_actor_debug_output = st.empty()
+        environment_state_ouput = st.empty()
 
     # Initialize resources if not already done
     system_one_audio_status.write("Initializing streaming")
@@ -98,7 +100,7 @@ async def main():
     try:
         while True:
             if "streamlit_av_queue" in st.session_state:
-                st.session_state.streamlit_av_queue.set_listening(listening)
+                st.session_state.streamlit_av_queue.set_looking_listening(looking, listening)
             if not webrtc_ctx.state.playing:
                 system_one_audio_status.write("Stopped.")
                 await asyncio.sleep(0.1)
@@ -112,8 +114,10 @@ async def main():
                     pass
             if charles_actor is not None: 
                 try:
-                    audio_history = await charles_actor.get_system_one_audio_history_output.remote()
-                    system_one_audio_history_output.markdown(audio_history)
+                    new_environment_state = await charles_actor.get_environment_state.remote()
+                    environment_state_ouput.markdown(f"{new_environment_state}")
+                    charles_debug_str = await charles_actor.get_charles_actor_debug_output.remote()
+                    charles_actor_debug_output.markdown(charles_debug_str)
                 except Exception as e:
                     # assume we disconnected
                     charles_actor = None

charles_actor.py

@@ -1,23 +1,30 @@
+import json
 import ray
 import time
 import asyncio
 import os
 from clip_transform import CLIPTransform
+from environment_state_actor import EnvironmentStateActor, EnvironmentState
+from agent_state_actor import AgentStateActor
+import asyncio
 
 @ray.remote
 class CharlesActor:
     def __init__(self):
         self._needs_init = True
-        self._system_one_audio_history_output = ""
+        self._charles_actor_debug_output = ""
+        self._environment_state:EnvironmentState = EnvironmentState(episode=0, step=0)  # Initialize as EnvironmentState
         self._state = "Initializing"
         self._clip_transform = CLIPTransform()
         
     def get_state(self):
         return self._state
     
+    def get_charles_actor_debug_output(self):
+        return self._charles_actor_debug_output
     
-    def get_system_one_audio_history_output(self):
-        return self._system_one_audio_history_output
+    def get_environment_state(self)->EnvironmentState:
+        return self._environment_state
     
     async def _initalize_resources(self):
         # Initialize resources
@@ -28,7 +35,9 @@ class CharlesActor:
 
         print("001 - create RespondToPromptActor")
         from respond_to_prompt_actor import RespondToPromptActor
-        self._respond_to_prompt_actor = RespondToPromptActor.remote(self._out_audio_queue)
+        self._environment_state_actor = EnvironmentStateActor.remote()
+        self._agent_state_actor = AgentStateActor.remote()
+        self._respond_to_prompt_actor = RespondToPromptActor.remote(self._environment_state_actor, self._out_audio_queue)
 
         print("002 - create SpeechToTextVoskActor")
         from speech_to_text_vosk_actor import SpeechToTextVoskActor
@@ -51,7 +60,20 @@ class CharlesActor:
         if self._needs_init:
             await self._initalize_resources()
             
-        system_one_audio_history = []
+        debug_output_history = []
+        def add_debug_output(output):
+            debug_output_history.append(output)
+            if len(debug_output_history) > 10:
+                debug_output_history.pop(0)
+            table_content = "| Charles Actor debug history |\n| --- |\n"
+            table_content += "\n".join([f"| {item} |" for item in reversed(debug_output_history)])
+            self._charles_actor_debug_output = table_content
+        def preview_debug_output(output):
+            table_content = "| Charles Actor debug history |\n| --- |\n"
+            debug_output_history_copy = debug_output_history.copy()
+            debug_output_history_copy.append(output)
+            table_content += "\n".join([f"| {item} |" for item in reversed(debug_output_history_copy)])
+            self._charles_actor_debug_output = table_content
         
         self._state = "Waiting for input"
         total_video_frames = 0
@@ -67,7 +89,13 @@ class CharlesActor:
             if len(self._debug_queue) > 0:
                 prompt = self._debug_queue.pop(0)
                 await self._respond_to_prompt_actor.enqueue_prompt.remote(prompt)
+            
+            env_state = await self._environment_state_actor.begin_next_step.remote()
+            self._environment_state = env_state
+            self._agent_state_actor.begin_step.remote()
             audio_frames = await self._streamlit_av_queue.get_in_audio_frames_async()    
+            video_frames = await self._streamlit_av_queue.get_video_frames_async()
+
             if len(audio_frames) > 0:
                 total_audio_frames += len(audio_frames)
                 # Concatenate all audio frames into a single buffer
@@ -76,6 +104,20 @@ class CharlesActor:
                 process_speech_to_text_future.append(future)
             # audio_frames_task = None
 
+            if len(video_frames) > 0:
+                vector_debug = f"found {len(video_frames)} video frames"
+                total_video_frames += 1
+                skipped_video_frames += (len(video_frames) -1)
+                image_as_array = video_frames[-1]
+                image_vector = self._clip_transform.image_to_embeddings(image_as_array)
+                image_vector = image_vector[0]
+                distances, closest_item_key, distance_debug_str = self._prototypes.get_distances(image_vector)
+                vector_debug = f"{closest_item_key} {distance_debug_str}"
+
+
+            human_preview_text = ""
+            robot_preview_text = ""
+
             if len(process_speech_to_text_future) > 0:
                 ready, _ = ray.wait([process_speech_to_text_future[0]], timeout=0)
                 if ready:
@@ -87,32 +129,26 @@ class CharlesActor:
                     if speaker_finished and len(prompt) > 0 and prompt not in prompts_to_ignore:
                         print(f"Prompt: {prompt}")
                         # system_one_audio_history.append("... " + str(raw_json))
-                        system_one_audio_history.append(prompt)
-                        while len(system_one_audio_history) > 10:
-                            system_one_audio_history = system_one_audio_history[-10:]
-                        table_content = "| System 1 Audio History |\n| --- |\n"
-                        table_content += "\n".join([f"| {item} |" for item in reversed(system_one_audio_history)])
-                        self._system_one_audio_history_output = table_content
+                        add_debug_output(f"👨 {prompt}")
                         await self._respond_to_prompt_actor.enqueue_prompt.remote(prompt)
-                    # else:
-                    #     print(f"not ready... " + str(raw_json))
-                        
-            video_frames = await self._streamlit_av_queue.get_video_frames_async()
-            if len(video_frames) > 0:
-                vector_debug = f"found {len(video_frames)} video frames"
-                total_video_frames += 1
-                skipped_video_frames += (len(video_frames) -1)
-                image_as_array = video_frames[-1]
-                image_vector = self._clip_transform.image_to_embeddings(image_as_array)
-                image_vector = image_vector[0]
-                distances, closest_item_key, distance_debug_str = self._prototypes.get_distances(image_vector)
-                vector_debug = f"{closest_item_key} {distance_debug_str}"
+                    elif len(prompt) > 0 and prompt not in prompts_to_ignore:
+                        human_preview_text = f"👨❓ {prompt}"
+
+            for new_response in env_state.llm_responses:
+                add_debug_output(f"🤖 {new_response}")
+            if len(env_state.llm_preview):
+                robot_preview_text = f"🤖❓ {env_state.llm_preview}"
+
+            if len(human_preview_text) > 0:
+                preview_debug_output(human_preview_text)
+            elif len(robot_preview_text) > 0:
+                preview_debug_output(robot_preview_text)
             
             await asyncio.sleep(0.01)
             loops+=1
             self._state = f"Processed {total_video_frames} video frames and {total_audio_frames} audio frames, loops: {loops}. loops per second: {loops/(time.time()-start_time):.2f}. {vector_debug}"
 
-if __name__ == "__main__":
+async def main():
     if not ray.is_initialized():
         # Try to connect to a running Ray cluster
         ray_address = os.getenv('RAY_ADDRESS')
@@ -127,6 +163,8 @@ if __name__ == "__main__":
         ).remote() 
     future = charles_actor.start.remote()
 
+    last_step = -1
+    last_episode = -1
     try:
         while True:
             ready, _ = ray.wait([future], timeout=0)
@@ -141,9 +179,30 @@ if __name__ == "__main__":
                 break
             else:
                 # The start method is still running. You can poll for debug information here.
-                time.sleep(1)
-                state = charles_actor.get_state.remote()
-                print(f"Charles is in state: {ray.get(state)}")
+                await asyncio.sleep(1)
+                state = await charles_actor.get_state.remote()
+                env_state = await charles_actor.get_environment_state.remote()
+                if (env_state.episode != last_episode) or (env_state.step != last_step):
+                    last_episode = env_state.episode
+                    last_step = env_state.step
+                    print(f"Charles is in state: {state}")
+                    # if len(env_state.llm_preview):
+                    #     print (f"llm_preview: {env_state.llm_preview}")
+                    # if len(env_state.llm_responses):
+                    #     print (f"llm_responses: {env_state.llm_responses}")
+                    # if len(env_state.tts_raw_chunk_ids):
+                    #     for chunk_json in env_state.tts_raw_chunk_ids:
+                    #         chunk = json.loads(chunk_json)
+                    #         prompt = chunk['prompt']
+                    #         line = chunk['llm_sentence_id']
+                    #         chunk_id = chunk['chunk_count']
+                    #         print(f"Prompt: {prompt}, Line: {line}, Chunk: {chunk_id}")                            
+
     except KeyboardInterrupt as e:
         print("Script was manually terminated")
         raise(e)
+    
+
+if __name__ == "__main__":
+    loop = asyncio.get_event_loop()
+    loop.run_until_complete(main())

chat_service.py

@@ -4,7 +4,6 @@ import json
 import os
 import torch
 import openai
-from agent_response import AgentResponse
 
 class ChatService:
     def __init__(self, api="openai", model_id = "gpt-3.5-turbo"):
@@ -102,8 +101,8 @@ I fell off the pink step, and I had an accident.
             return True
         return False
 
-    async def get_responses_as_sentances_async(self, agent_response:AgentResponse, cancel_event):
-        self._messages.append({"role": "user", "content": agent_response['prompt']})
+    async def get_responses_as_sentances_async(self, prompt, cancel_event):
+        self._messages.append({"role": "user", "content": prompt})
         llm_response = ""
         current_sentence = ""
         delay = 0.1
@@ -125,21 +124,17 @@ I fell off the pink step, and I had an accident.
                         chunk_text = chunk_message['content']
                         current_sentence += chunk_text
                         llm_response += chunk_text
-                        agent_response['llm_preview'] = current_sentence
                         text_to_speak = self._should_we_send_to_voice(current_sentence)
                         if text_to_speak:
                             current_sentence = current_sentence[len(text_to_speak):]
-                            agent_response['llm_preview'] = ''
-                            agent_response['llm_sentence'] = text_to_speak
-                            agent_response['llm_sentences'].append(text_to_speak)
-                            yield agent_response
-                            agent_response['llm_sentence_id'] += 1
+                            yield text_to_speak, True
+                        else:
+                            yield current_sentence, False
 
                 if cancel_event.is_set():
                     return
                 if len(current_sentence) > 0:
-                    agent_response['llm_sentence'] = current_sentence
-                    yield agent_response
+                    yield current_sentence, True
                 self._messages.append({"role": "assistant", "content": llm_response})
                 return
 

environment_state_actor.py

@@ -0,0 +1,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

prototypes.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d59c424a695d82b6346742f635d1c4b469d5ad72c8dd47b590886d745681fd44
+size 6900

prototypes.py

@@ -10,8 +10,7 @@ from torch.nn import functional as F
 class Prototypes:
     def __init__(self):
         self._clip_transform = CLIPTransform()
-        self._prepare_prototypes()
-
+        self._load_prototypes()
 
     def _prepare_prototypes(self):
         image_embeddings = self.load_images_from_folder('prototypes')
@@ -32,6 +31,15 @@ class Prototypes:
 
         self.prototype_keys = ["person", "no_person"]
         self.prototypes = torch.stack([person_embedding, no_person_embedding])
+        # save prototypes to file
+        torch.save(self.prototypes, 'prototypes.pt')
+
+    def _load_prototypes(self):
+        # check if file exists
+        if not os.path.exists('prototypes.pt'):
+            self._prepare_prototypes()
+        self.prototypes = torch.load('prototypes.pt')
+        self.prototype_keys = ["person", "no_person"]
 
 
     def load_images_from_folder(self, folder):

respond_to_prompt_actor.py

@@ -8,28 +8,45 @@ import asyncio
 # from ray.actor import ActorHandle
 from ffmpeg_converter_actor import FFMpegConverterActor
 from agent_response import AgentResponse
-
+from environment_state_actor import EnvironmentStateActor
+from agent_state_actor import AgentStateActor
+from agent_state_actor import AgentState
+import json
 
 @ray.remote
 class PromptToLLMActor:
-    def __init__(self, input_queue:Queue, output_queue:Queue):
+    def __init__(
+            self, 
+            environment_state_actor:EnvironmentStateActor, 
+            input_queue:Queue,
+            output_queue:Queue):
         load_dotenv()
         self.input_queue = input_queue
         self.output_queue = output_queue
         self.chat_service = ChatService()
         self.cancel_event = None
+        self.environment_state_actor = environment_state_actor
 
     async def run(self):
         while True:
             prompt = await self.input_queue.get_async()
             self.cancel_event = asyncio.Event()
             agent_response = AgentResponse(prompt)
-            async for sentence_response in self.chat_service.get_responses_as_sentances_async(agent_response, self.cancel_event):
-                if self.chat_service.ignore_sentence(sentence_response['llm_sentence']):
+            async for text, is_complete_sentance in self.chat_service.get_responses_as_sentances_async(prompt, self.cancel_event):
+                if self.chat_service.ignore_sentence(text):
+                    is_complete_sentance = False
+                if not is_complete_sentance:
+                    agent_response['llm_preview'] = text
+                    await self.environment_state_actor.set_llm_preview.remote(text)
                     continue
-                print(f"{sentence_response['llm_sentence']} id: {agent_response['llm_sentence_id']} from prompt: {agent_response['prompt']}")
-                sentence_response = sentence_response.make_copy()
+                agent_response['llm_preview'] = ''
+                agent_response['llm_sentence'] = text
+                agent_response['llm_sentences'].append(text)
+                await self.environment_state_actor.add_llm_response_and_clear_llm_preview.remote(text)
+                print(f"{agent_response['llm_sentence']} id: {agent_response['llm_sentence_id']} from prompt: {agent_response['prompt']}")
+                sentence_response = agent_response.make_copy()
                 await self.output_queue.put_async(sentence_response)
+                agent_response['llm_sentence_id'] += 1
                 
     async def cancel(self):
         if self.cancel_event:
@@ -41,20 +58,35 @@ class PromptToLLMActor:
 
 @ray.remote
 class LLMSentanceToSpeechActor:
-    def __init__(self, input_queue, output_queue, voice_id):
+    def __init__(
+            self,
+            environment_state_actor:EnvironmentStateActor, 
+            input_queue, 
+            output_queue, 
+            voice_id):
         load_dotenv()
         self.input_queue = input_queue
         self.output_queue = output_queue
         self.tts_service = TextToSpeechService(voice_id=voice_id)
         self.cancel_event = None
+        self.environment_state_actor = environment_state_actor
 
     async def run(self):
         while True:
             sentence_response = await self.input_queue.get_async()
             self.cancel_event = asyncio.Event()
+            chunk_count = 0
             async for chunk_response in self.tts_service.get_speech_chunks_async(sentence_response, self.cancel_event):
                 chunk_response = chunk_response.make_copy()
                 await self.output_queue.put_async(chunk_response)
+                chunk_response = {
+                    'prompt': sentence_response['prompt'],
+                    'llm_sentence_id': sentence_response['llm_sentence_id'],
+                    'chunk_count': chunk_count,
+                }
+                chunk_id_json = json.dumps(chunk_response)
+                await self.environment_state_actor.add_tts_raw_chunk_id.remote(chunk_id_json)
+                chunk_count += 1
 
     async def cancel(self):
         if self.cancel_event:
@@ -85,7 +117,10 @@ class LLMSentanceToSpeechActor:
 
 @ray.remote
 class SpeechToConverterActor:
-    def __init__(self, input_queue:Queue, ffmpeg_converter_actor:FFMpegConverterActor):
+    def __init__(
+            self,
+            input_queue:Queue,
+            ffmpeg_converter_actor:FFMpegConverterActor):
         load_dotenv()
         self.input_queue = input_queue
         self.ffmpeg_converter_actor = ffmpeg_converter_actor
@@ -105,18 +140,31 @@ class SpeechToConverterActor:
 
 @ray.remote
 class RespondToPromptActor:
-    def __init__(self, out_audio_queue):
+    def __init__(
+            self, 
+            environment_state_actor:EnvironmentStateActor, 
+            out_audio_queue):
         voice_id="2OviOUQc1JsQRQgNkVBj"
         self.prompt_queue = Queue(maxsize=100)
         self.llm_sentence_queue = Queue(maxsize=100)
         self.speech_chunk_queue = Queue(maxsize=100)
+        self.environment_state_actor = environment_state_actor
 
         self.ffmpeg_converter_actor = FFMpegConverterActor.remote(out_audio_queue)
         
-        self.prompt_to_llm = PromptToLLMActor.remote(self.prompt_queue, self.llm_sentence_queue)
-        self.llm_sentence_to_speech = LLMSentanceToSpeechActor.remote(self.llm_sentence_queue, self.speech_chunk_queue, voice_id)
+        self.prompt_to_llm = PromptToLLMActor.remote(
+            self.environment_state_actor,
+            self.prompt_queue, 
+            self.llm_sentence_queue)
+        self.llm_sentence_to_speech = LLMSentanceToSpeechActor.remote(
+            self.environment_state_actor,
+            self.llm_sentence_queue, 
+            self.speech_chunk_queue, 
+            voice_id)
         # self.speech_output = SpeechToSpeakerActor.remote(self.speech_chunk_queue, voice_id)
-        self.speech_output = SpeechToConverterActor.remote(self.speech_chunk_queue, self.ffmpeg_converter_actor)
+        self.speech_output = SpeechToConverterActor.remote(
+            self.speech_chunk_queue, 
+            self.ffmpeg_converter_actor)
 
         # Start the pipeline components.
         self.prompt_to_llm.run.remote()

streamlit_av_queue.py

@@ -15,14 +15,16 @@ class StreamlitAVQueue:
         self._output_channels = 2
         self._audio_bit_rate = audio_bit_rate
         self._listening = True
+        self._looking = True
         self._lock = threading.Lock()
         self.queue_actor = WebRtcAVQueueActor.options(
             name="WebRtcAVQueueActor", 
             get_if_exists=True,
             ).remote() 
         
-    def set_listening(self, listening: bool):
+    def set_looking_listening(self, looking, listening: bool):
         with self._lock:
+            self._looking = looking
             self._listening = listening
             
     async def queued_video_frames_callback(
@@ -30,10 +32,13 @@ class StreamlitAVQueue:
                 frames: List[av.VideoFrame],
             ) -> av.VideoFrame:
         try:
-            for frame in frames:
-                shared_tensor = frame.to_ndarray(format="rgb24")
-                shared_tensor_ref = ray.put(shared_tensor)
-                await self.queue_actor.enqueue_in_video_frame.remote(shared_tensor_ref)
+            with self._lock:
+                should_look = self._looking
+            if len(frames) > 0 and should_look:
+                for frame in frames:
+                    shared_tensor = frame.to_ndarray(format="rgb24")
+                    shared_tensor_ref = ray.put(shared_tensor)
+                    await self.queue_actor.enqueue_in_video_frame.remote(shared_tensor_ref)
                 # print (f"tesnor len: {len(shared_tensor)}, tensor shape: {shared_tensor.shape}, tensor type:{shared_tensor.dtype} tensor ref: {shared_tensor_ref}")
         except Exception as e:
             print (e)