Accomplish the exp scripts v1; Add video recording and trajectory recording of desktop agent; Fix minor bugs

desktop_env/controllers/python.py

@@ -243,6 +243,32 @@ class PythonController:
         else:
             raise Exception(f"Unknown action type: {action_type}")
 
+    # Record video
+    def start_recording(self):
+        """
+        Starts recording the screen.
+        """
+        response = requests.post(self.http_server + "/start_recording")
+        if response.status_code == 200:
+            logger.info("Recording started successfully")
+        else:
+            logger.error("Failed to start recording. Status code: %d", response.status_code)
+
+    def end_recording(self, dest: str):
+        """
+        Ends recording the screen.
+        """
+        response = requests.post(self.http_server + "/end_recording")
+        if response.status_code == 200:
+            logger.info("Recording stopped successfully")
+            with open(dest, 'wb') as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    if chunk:
+                        f.write(chunk)
+        else:
+            logger.error("Failed to stop recording. Status code: %d", response.status_code)
+            return None
+
     # Additional info
     def get_vm_platform(self):
         """

desktop_env/controllers/setup.py

@@ -209,10 +209,6 @@ class SetupController:
         if not command:
             raise Exception("Empty command to launch.")
 
-        if isinstance(command, str) and len(command.split()) > 1:
-            logger.warning("Command should be a list of strings. Now it is a string. Will split it by space.")
-            command = command.split()
-
         payload = json.dumps({"command": command})
         headers = {"Content-Type": "application/json"}
 

desktop_env/evaluators/metrics/vlc.py

@@ -86,6 +86,9 @@ def compare_images(image1_path, image2_path):
     # score = compare_images('path_to_image1', 'path_to_image2')
     # print("Similarity score:", score)
 
+    if not image1_path or not image2_path:
+        return 0
+
     # Open the images and convert to grayscale
     image1 = Image.open(image1_path).convert('L')
     image2 = Image.open(image2_path).convert('L')
@@ -119,6 +122,9 @@ def compare_audios(audio_path_1, audio_path_2, max_distance=1000):
     # print(f'Similarity Score: {similarity}')
 
     # Convert to common format if necessary and load audio
+    if not audio_path_1 or not audio_path_2:
+        return 0
+
     y1, sr1 = librosa.load(audio_path_1)
     y2, sr2 = librosa.load(audio_path_2)
 

desktop_env/server/README.md

@@ -78,3 +78,8 @@ Activating the window manager control requires the installation of `wmctrl`:
 ```bash
 sudo apt install wmctrl
 ```
+
+To enable recording in the virtual machine, you need to install `ffmpeg`:
+```bash
+sudo apt install ffmpeg
+```

desktop_env/server/main.py

@@ -1,6 +1,7 @@
 import ctypes
 import os
 import platform
+import shlex
 import subprocess
 from pathlib import Path
 from typing import Any, Optional
@@ -13,7 +14,7 @@ import pyautogui
 import requests
 from PIL import Image
 from Xlib import display, X
-from flask import Flask, request, jsonify, send_file, abort
+from flask import Flask, request, jsonify, send_file, abort, send_from_directory
 from lxml.etree import _Element
 from pyatspi import Accessible, StateType
 from pyatspi import Action as ATAction
@@ -29,7 +30,8 @@ pyautogui.PAUSE = 0
 pyautogui.DARWIN_CATCH_UP_TIME = 0
 
 logger = app.logger
-
+recording_process = None  # fixme: this is a temporary solution for recording, need to be changed to support multiple-process
+recording_path = "/tmp/recording.mp4"
 
 @app.route('/setup/execute', methods=['POST'])
 @app.route('/execute', methods=['POST'])
@@ -39,6 +41,9 @@ def execute_command():
     shell = data.get('shell', False)
     command = data.get('command', "" if shell else [])
 
+    if isinstance(command, str):
+        command = shlex.split(command)
+
     # Execute the command without any safety checks.
     try:
         result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=shell, text=True)
@@ -60,6 +65,9 @@ def launch_app():
     data = request.json
     command: List[str] = data.get("command", [])
 
+    if isinstance(command, str):
+        command = shlex.split(command)
+
     try:
         subprocess.Popen(command)
         return "{:} launched successfully".format(" ".join(command))
@@ -604,5 +612,42 @@ def activate_window():
     return "File opened successfully", 200
 
 
+@app.route('/start_recording', methods=['POST'])
+def start_recording():
+    global recording_process
+    if recording_process:
+        return jsonify({'status': 'error', 'message': 'Recording is already in progress.'}), 400
+
+    d = display.Display()
+    screen_width = d.screen().width_in_pixels
+    screen_height = d.screen().height_in_pixels
+
+    start_command = f"ffmpeg -y -f x11grab -draw_mouse 1 -s {screen_width}x{screen_height} -i :0.0 -c:v libx264 -r 30 {recording_path}"
+
+    recording_process = subprocess.Popen(shlex.split(start_command), stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+
+    return jsonify({'status': 'success', 'message': 'Started recording.'})
+
+
+@app.route('/end_recording', methods=['POST'])
+def end_recording():
+    global recording_process
+
+    if not recording_process:
+        return jsonify({'status': 'error', 'message': 'No recording in progress to stop.'}), 400
+
+    recording_process.terminate()
+    recording_process.wait()
+    return_code = recording_process.returncode
+    output, error = recording_process.communicate()
+    recording_process = None
+
+    # return recording video file
+    if os.path.exists(recording_path):
+        return send_file(recording_path, as_attachment=True)
+    else:
+        return abort(404, description="Recording failed")
+
+
 if __name__ == '__main__':
     app.run(debug=True, host="0.0.0.0")