Spaces:

jbilcke-hf
/

train-robots-with-mujoco

Paused

Julian Bilcke

more logs

9b8c4e3 4 months ago

7.6 kB

	#!/usr/bin/env python3
	"""
	GPU initialization script for MuJoCo EGL rendering in containerized environments.
	This should be run before starting the notebook to ensure GPU is properly set up.
	"""

	import os
	import subprocess
	import sys

	def check_nvidia_driver():
	"""Check if NVIDIA driver is accessible."""
	try:
	result = subprocess.run(['nvidia-smi'], capture_output=True, text=True)
	if result.returncode == 0:
	print("✓ NVIDIA driver accessible")
	print(result.stdout.split('\n')[2]) # Driver info line
	return True
	else:
	print("✗ NVIDIA driver not accessible")
	return False
	except FileNotFoundError:
	print("✗ nvidia-smi not found")
	return False

	def check_egl_libs():
	"""Check if EGL libraries are available."""
	import ctypes
	libs_to_check = [
	'libEGL.so.1',
	'libGL.so.1',
	'libEGL_nvidia.so.0'
	]

	for lib in libs_to_check:
	try:
	ctypes.CDLL(lib)
	print(f"✓ {lib} loaded successfully")
	except OSError as e:
	print(f"✗ Failed to load {lib}: {e}")

	def test_egl_device():
	"""Test EGL device creation with multiple approaches."""
	try:
	from OpenGL import EGL
	import ctypes

	# Method 1: Try platform device display (preferred for headless)
	try:
	display = EGL.eglGetPlatformDisplay(EGL.EGL_PLATFORM_DEVICE_EXT,
	EGL.EGL_DEFAULT_DISPLAY,
	None)
	if display != EGL.EGL_NO_DISPLAY:
	major = ctypes.c_long()
	minor = ctypes.c_long()
	if EGL.eglInitialize(display, ctypes.byref(major), ctypes.byref(minor)):
	print(f"✓ EGL platform device initialized (version {major.value}.{minor.value})")
	EGL.eglTerminate(display)
	return True
	except Exception as e:
	print(f" Platform device method failed: {e}")

	# Method 2: Try default display
	try:
	display = EGL.eglGetDisplay(EGL.EGL_DEFAULT_DISPLAY)
	if display != EGL.EGL_NO_DISPLAY:
	major = ctypes.c_long()
	minor = ctypes.c_long()
	if EGL.eglInitialize(display, ctypes.byref(major), ctypes.byref(minor)):
	print(f"✓ EGL default display initialized (version {major.value}.{minor.value})")
	EGL.eglTerminate(display)
	return True
	except Exception as e:
	print(f" Default display method failed: {e}")

	# Method 3: Try surfaceless context (what MuJoCo likely uses)
	try:
	os.environ['EGL_PLATFORM'] = 'surfaceless'
	display = EGL.eglGetDisplay(EGL.EGL_DEFAULT_DISPLAY)
	if display != EGL.EGL_NO_DISPLAY:
	major = ctypes.c_long()
	minor = ctypes.c_long()
	if EGL.eglInitialize(display, ctypes.byref(major), ctypes.byref(minor)):
	print(f"✓ EGL surfaceless display initialized (version {major.value}.{minor.value})")
	EGL.eglTerminate(display)
	return True
	except Exception as e:
	print(f" Surfaceless method failed: {e}")

	print("✗ All EGL initialization methods failed")
	return False

	except Exception as e:
	print(f"✗ EGL test failed: {e}")
	return False

	def test_mujoco_rendering():
	"""Test MuJoCo rendering capability with different approaches."""
	try:
	import mujoco

	# Create a simple model
	xml = """
	<mujoco>
	<worldbody>
	<body>
	<geom type="box" size="1 1 1"/>
	</body>
	</worldbody>
	</mujoco>
	"""

	model = mujoco.MjModel.from_xml_string(xml)

	# Try different rendering approaches
	approaches = [
	("Small resolution", {"height": 64, "width": 64}),
	("Default resolution", {"height": 240, "width": 320}),
	("Large resolution", {"height": 480, "width": 640}),
	]

	for name, kwargs in approaches:
	try:
	renderer = mujoco.Renderer(model, **kwargs)
	data = mujoco.MjData(model)
	renderer.update_scene(data)
	pixels = renderer.render()
	print(f" ✓ {name} ({kwargs['width']}x{kwargs['height']}): SUCCESS")
	print(f" Image shape: {pixels.shape}, dtype: {pixels.dtype}")

	# Test if we got actual rendered content (not all zeros)
	if pixels.max() > 0:
	print(f" ✓ Non-zero pixels detected (max value: {pixels.max()})")
	else:
	print(f" ⚠️ All pixels are zero - may indicate rendering issue")

	# Clean up
	del renderer
	return True

	except Exception as e:
	print(f" ✗ {name}: {e}")
	continue

	print("✗ All MuJoCo rendering approaches failed")
	return False

	except ImportError:
	print("✗ MuJoCo not installed")
	return False
	except Exception as e:
	print(f"✗ MuJoCo test failed: {e}")
	return False

	def main():
	"""Run all GPU initialization checks."""
	print("🔧 Initializing GPU for MuJoCo rendering...")
	print("=" * 50)

	# Set environment variables
	os.environ['MUJOCO_GL'] = 'egl'
	os.environ['PYOPENGL_PLATFORM'] = 'egl'
	os.environ['EGL_PLATFORM'] = 'surfaceless' # Better for headless

	print("Environment variables set:")
	print(f" MUJOCO_GL: {os.environ.get('MUJOCO_GL')}")
	print(f" PYOPENGL_PLATFORM: {os.environ.get('PYOPENGL_PLATFORM')}")
	print(f" EGL_PLATFORM: {os.environ.get('EGL_PLATFORM')}")
	print()

	# Run checks
	checks = [
	("NVIDIA Driver", check_nvidia_driver),
	("EGL Libraries", lambda: check_egl_libs() or True), # Always continue
	("EGL Device", test_egl_device),
	("MuJoCo Rendering", test_mujoco_rendering),
	]

	results = []
	for name, check_func in checks:
	print(f"Checking {name}...")
	try:
	result = check_func()
	results.append((name, result))
	except Exception as e:
	print(f"✗ {name} check failed with exception: {e}")
	results.append((name, False))
	print()

	# Summary
	print("=" * 50)
	print("🔍 Summary:")

	# Check if the critical test (MuJoCo) passed
	mujoco_passed = any(name == "MuJoCo Rendering" and passed for name, passed in results)

	for name, passed in results:
	status = "✓ PASS" if passed else "✗ FAIL"
	print(f" {name}: {status}")

	if mujoco_passed:
	print("\n🎉 MuJoCo rendering works! The notebook should work even if some EGL tests fail.")
	print("💡 Note: EGL device tests may fail but MuJoCo can still render successfully.")
	return 0
	else:
	print("\n⚠️ MuJoCo rendering failed. GPU rendering will not work properly.")
	print("💡 Try checking the container GPU configuration or driver compatibility.")
	return 1

	if __name__ == "__main__":
	sys.exit(main())