Merge pull request #15 from Astera-org/headlessScreenshot

.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+    "editor.formatOnSave": false
+}

CMakeLists.txt

@@ -39,6 +39,7 @@ endif()
 
 # Configuration options
 set(BUILD_CLIENT TRUE CACHE BOOL "Build client")
+set(BUILD_HEADLESS TRUE CACHE BOOL "Build in headless mode")
 set(BUILD_SERVER FALSE CACHE BOOL "Build server")
 set(BUILD_UNITTESTS TRUE CACHE BOOL "Build unittests")
 set(BUILD_BENCHMARKS FALSE CACHE BOOL "Build benchmarks")

cmake/Modules/FindZmq.cmake

@@ -1,10 +1,10 @@
 mark_as_advanced(ZMQ_LIBRARY ZMQ_INCLUDE_DIR)
 
 find_path(ZMQ_INCLUDE_DIR NAMES zmq.h)
-message(STATUS ${ZMQPP_INCLUDE_DIR})
+message(STATUS ${ZMQ_INCLUDE_DIR})
 
 find_library(ZMQ_LIBRARY NAMES zmq)
-message(STATUS ${ZMQPP_LIBRARY})
+message(STATUS ${ZMQ_LIBRARY})
 
 include(FindPackageHandleStandardArgs)
 find_package_handle_standard_args(Zmq DEFAULT_MSG ZMQ_LIBRARY ZMQ_INCLUDE_DIR)

gym_client_headless.py

@@ -0,0 +1,37 @@
+import gymnasium as gym
+import pygame
+import numpy as np
+from minetester.minetest_env import KEY_MAP
+import sys
+
+# The Makefile puts the binary into build/macos
+minetest_executable = "/home/simon/minetest/bin/minetest"
+if sys.platform == "darwin":
+    minetest_executable = (
+        "/Users/siboehm/repos/minetest/build/macos/minetest.app/Contents/MacOS/minetest"
+    )
+
+env = gym.make(
+    "minetest",
+    minetest_executable=minetest_executable,
+    render_mode="human",
+    display_size=(1600, 1200),
+    start_xvfb=True,
+    headless=True,
+)
+env.reset()
+
+for i in range(200):
+    print(f"i: {i}")
+    state, reward, terminated, truncated, info = env.step(
+        {
+            "KEYS": np.zeros(len(KEY_MAP), dtype=bool),
+            "MOUSE": np.array([0.0, 0.0]),
+        }
+    )
+    print(
+        f"R: {reward} Term: {terminated} Trunc: {truncated} AllBlack: {state.sum() == 0}"
+    )
+
+env.close()
+pygame.quit()

handshaker.py

@@ -0,0 +1,54 @@
+import capnp
+import pkg_resources
+import zmq
+import numpy as np
+from PIL import Image
+
+"""
+Useful for debugging:
+- lldb -- /home/simon/minetest/bin/minetest --go --worldname test_world --config /home/simon/minetest/artifacts/2dd22d78-8c03-445e-83ad-8fff429569d4.conf --remote-input localhost:5555 --headless
+- Then run handshaker.py
+"""
+
+def unpack_pb_obs(received_obs: str):
+    with remoteclient_capnp.Observation.from_bytes(received_obs) as obs_proto:
+        # Convert the response to a numpy array
+        img = obs_proto.image
+        img_data = np.frombuffer(img.data, dtype=np.uint8).reshape(
+            (img.height, img.width, 3)
+        )
+        # Reshape the numpy array to the correct dimensions
+        reward = obs_proto.reward
+        done = obs_proto.done
+    return img_data, reward, done
+
+
+remoteclient_capnp = capnp.load(
+    pkg_resources.resource_filename(
+        "minetester", "../../src/network/proto/remoteclient.capnp"
+    )
+)
+context = zmq.Context()
+socket = context.socket(zmq.REQ)
+socket.connect("tcp://localhost:5555")
+
+i = 0
+inp = ""
+while inp != "stop":
+    pb_action = remoteclient_capnp.Action.new_message()
+    socket.send(pb_action.to_bytes())
+
+    byte_obs = socket.recv()
+    (
+        obs,
+        _,
+        _,
+    ) = unpack_pb_obs(byte_obs)
+
+    # Save the observation as a PNG file
+    if obs.size > 0:
+        image = Image.fromarray(obs)
+        image.save(f"observation_{i}.png")
+
+    i += 1
+    inp = input("Stop with 'stop':")

python/minetester/minetest_env.py

@@ -60,6 +60,9 @@ def __init__(
         game_id: str = "minetest",
         client_name: str = "minetester",
         config_dict: Dict[str, Any] = None,
+        start_xvfb: bool = False,
+        x_display: Optional[int] = None,
+        headless: bool = False,
     ):
         if config_dict is None:
             config_dict = {}
@@ -69,6 +72,12 @@ def __init__(
         self.fov_y = fov
         self.fov_x = self.fov_y * self.display_size.width / self.display_size.height
         self.render_mode = render_mode
+        self.headless = headless
+        self.start_xvfb = start_xvfb and self.headless
+
+        assert (
+            not self.start_xvfb or sys.platform == "linux"
+        ), "Xvfb is only supported on Linux."
 
         if render_mode == "human":
             self._start_pygame()
@@ -131,6 +140,16 @@ def __init__(
         # Configure game and mods
         self.game_id = game_id
 
+        # Start X server virtual frame buffer
+        self.default_display = x_display or 0
+        if "DISPLAY" in os.environ:
+            self.default_display = int(os.environ["DISPLAY"].split(":")[1])
+        self.x_display = x_display or self.default_display
+        self.xserver_process = None
+        if self.start_xvfb:
+            self.x_display = x_display or self.default_display + 4
+            self.xserver_process = start_xserver(self.x_display, self.display_size)
+
     def _configure_spaces(self):
         # Define action and observation space
         self.max_mouse_move_x = self.display_size[0] // 2
@@ -206,6 +225,7 @@ def _reset_minetest(self):
             self.config_path,
             log_path,
             f"localhost:{self.env_port}",
+            display=self.x_display,
         )
 
     def _perform_client_handshake(self):
@@ -401,14 +421,11 @@ def start_minetest_client(
     log_path: str,
     client_socket: str,
     display: int = None,
+    headless: bool = True,
+    set_gpu_vars = True,
+    set_vsync_vars = True,
 ):
-    virtual_display = []
-    if sys.platform == "linux":
-        virtual_display = ["xvfb-run", "--auto-servernum"]
-    else:
-        print(f"Warning: virtual display not supported on {sys.platform}")
-
-    cmd = virtual_display + [
+    cmd = [
         minetest_executable,
         "--go",
         "--worldname",
@@ -417,16 +434,30 @@ def start_minetest_client(
         config_path,
         "--remote-input",
         client_socket,
-        "--verbose",
+        "--verbose"
     ]
+    if headless:
+        # don't render to screen
+        cmd.append("--headless")
 
     stdout_file = log_path.format("client_stdout")
     stderr_file = log_path.format("client_stderr")
     with open(stdout_file, "w") as out, open(stderr_file, "w") as err:
-        # client_env = os.environ.copy()
-        # if display is not None:
-        #     client_env["DISPLAY"] = ":" + str(display)
-        client_process = subprocess.Popen(cmd, stdout=out, stderr=err)
+        client_env = os.environ.copy()
+        if display is not None:
+            client_env["DISPLAY"] = ":" + str(display)
+        if set_gpu_vars:
+            # enable GPU usage
+            client_env["__GLX_VENDOR_LIBRARY_NAME"] = "nvidia"
+            client_env["__NV_PRIME_RENDER_OFFLOAD"] = "1"
+        if set_vsync_vars:
+            # disable vsync
+            client_env["__GL_SYNC_TO_VBLANK"] = "0"
+            client_env["vblank_mode"] = "0"
+        out.write(f"Starting client with command: {' '.join(str(x) for x in cmd)}\n")
+        out.write(f"Client environment: {client_env}\n")
+        client_process = subprocess.Popen(cmd, stdout=out, stderr=err, env=client_env)
+        out.write(f"Client started with pid {client_process.pid}\n")
     return client_process
 
 
@@ -455,3 +486,19 @@ def write_config_file(file_path, config):
     with open(file_path, "w") as f:
         for key, value in config.items():
             f.write(f"{key} = {value}\n")
+
+
+def start_xserver(
+    display_idx: int = 1,
+    display_size: Tuple[int, int] = (1024, 600),
+    display_depth: int = 24,
+):
+    cmd = [
+        "Xvfb",
+        f":{display_idx}",
+        "-screen",
+        "0",  # screennum param
+        f"{display_size[0]}x{display_size[1]}x{display_depth}",
+    ]
+    xserver_process = subprocess.Popen(cmd)
+    return xserver_process

src/CMakeLists.txt

@@ -36,6 +36,19 @@ set(CMAKE_BUILD_TYPE "${CMAKE_BUILD_TYPE}" CACHE STRING
 # Set some random things default to not being visible in the GUI
 mark_as_advanced(EXECUTABLE_OUTPUT_PATH LIBRARY_OUTPUT_PATH)
 
+if(BUILD_CLIENT AND BUILD_HEADLESS)
+	find_package(SDL2 REQUIRED NO_DEFAULT_PATH)
+	# SDL2 exports targets since 2.0.6, but some distributions override config.
+	if(NOT TARGET SDL2::SDL2)
+		add_library(SDL2::SDL2 INTERFACE IMPORTED)
+		set_target_properties(SDL2::SDL2 PROPERTIES
+			INTERFACE_INCLUDE_DIRECTORIES ${SDL2_INCLUDE_DIRS}
+			INTERFACE_LINK_LIBRARIES ${SDL2_LIBRARIES}
+		)
+	endif()
+	get_target_property(SDL2_LIB_LOCATION SDL2::SDL2 LOCATION)
+	message(STATUS "SDL2 library location: ${SDL2_LIB_LOCATION}")
+endif()
 
 if(NOT (BUILD_CLIENT OR BUILD_SERVER))
 	message(WARNING "Neither BUILD_CLIENT nor BUILD_SERVER is set! Setting BUILD_SERVER=true")
@@ -557,15 +570,20 @@ if(BUILD_CLIENT)
 		add_executable(${PROJECT_NAME} ${client_SRCS} ${extra_windows_SRCS})
 	endif()
 	add_dependencies(${PROJECT_NAME} GenerateVersion)
+	if(BUILD_HEADLESS)
+		set(SDL2_TARGET SDL2::SDL2)
+	endif()
 	target_link_libraries(
 		${PROJECT_NAME}
 		${ZLIB_LIBRARIES}
 		IrrlichtMt::IrrlichtMt
+		${SDL2_TARGET}
 		${ZSTD_LIBRARY}
 		${ZMQ_LIBRARY}
 		${ZMQPP_LIBRARY}
 		${SOUND_LIBRARIES}
 		${SQLITE3_LIBRARY}
+		${SDL2_LIBRARIES}
 		${LUA_LIBRARY}
 		${GMP_LIBRARY}
 		${JSON_LIBRARY}

src/client/client.cpp

@@ -1884,7 +1884,13 @@ float Client::getCurRate()
 void Client::makeScreenshot()
 {
 	irr::video::IVideoDriver *driver = m_rendering_engine->get_video_driver();
-	irr::video::IImage* const raw_image = driver->createScreenShot();
+	irr::video::IImage* raw_image;
+
+	if(m_rendering_engine->headless){
+		raw_image = m_rendering_engine->get_screenshot();
+	} else {
+		raw_image = driver->createScreenShot();
+	}
 
 	if (!raw_image)
 		return;
@@ -1951,6 +1957,11 @@ void Client::makeScreenshot()
 	raw_image->drop();
 }
 
+RenderingEngine *Client::getRenderingEngine()
+{
+	return m_rendering_engine;
+}
+
 bool Client::shouldShowMinimap() const
 {
 	return !m_minimap_disabled_by_server;

src/client/client.h

@@ -414,6 +414,8 @@ class Client : public con::PeerHandler, public InventoryManager, public IGameDef
 
 	void showMinimap(bool show = true);
 
+	RenderingEngine *getRenderingEngine();
+
 	const Address getServerAddress();
 
 	const std::string &getAddressName() const

src/client/clientlauncher.cpp

@@ -346,6 +346,8 @@ void ClientLauncher::init_args(GameStartData &start_data, const Settings &cmd_ar
 			remote_input_socket.clear();
 		}
 	}
+
+	start_data.headless = isRemote && cmd_args.getFlag("headless");
 }
 
 bool ClientLauncher::init_engine()

src/client/game.cpp

@@ -1161,7 +1161,7 @@ bool Game::startup(bool *kill,
 		return false;
 	input->registerLocalPlayer(client->getEnv().getLocalPlayer());
 
-	m_rendering_engine->initialize(client, hud);
+	m_rendering_engine->initialize(client, hud, start_data.isHeadless());
 
 	return true;
 }