File size: 5,120 Bytes
2ddbf65 7d40aab 2ddbf65 | 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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | import numpy as np
import pybullet as p
import time
import pickle
import os
from gym_pybullet_drones.envs.VelocityAviary import VelocityAviary
from gym_pybullet_drones.utils.enums import Physics
from utils import step_compute
def run_realtime_simulation():
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
components_path = os.path.join(BASE_DIR, "W_drone_components.pkl")
if not os.path.exists(components_path):
print(f"{components_path} not found!")
return
with open(components_path, "rb") as f:
components = pickle.load(f)
W, Win, Wout, mask, alpha, Nx = components['W'], components['Win'], components['Wout'], components['mask'], components['alpha'], components['W'].shape[0]
# Realistic physics engine features: Ground Effect, Drag, and Downwash
env = VelocityAviary(
gui=True,
initial_xyzs=np.array([[0, 0, .5]]),
initial_rpys=np.array([[0, 0, 0]]),
physics=Physics.PYB_GND_DRAG_DW
)
# Close PyBullet GUI side panels and enable shadows for realistic visuals
p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0, physicsClientId=env.CLIENT)
p.configureDebugVisualizer(p.COV_ENABLE_SHADOWS, 1, physicsClientId=env.CLIENT)
obs, info = env.reset()
state = obs[0] # VelocityAviary state is size 20
x = np.random.random((Nx, 1))
# --- VIDEO RECORDING SETTINGS ---
# 1. Third Person camera tracking the drone
log_id_tp = p.startStateLogging(p.STATE_LOGGING_VIDEO_MP4, "drone_third_person.mp4")
# 2. First Person camera (ego-view)
log_id_fp = p.startStateLogging(p.STATE_LOGGING_VIDEO_MP4, "drone_first_person.mp4")
print("Simulation starting, recording video...")
try:
for i in range(43200): # Match simulation length
# Blind Flight: Hide global X and Y positions from the network
state[0] = 0
state[1] = 0
x, y = step_compute(x, W, Win, Wout, state, alpha, mask)
# Velocity references are directly generated by ESN (Autonomous flight)
obs, reward, terminated, truncated, info = env.step(np.array([y.flatten()]))
state = obs[0]
# --- SPECTATOR CAMERA (FREECAM) CONTROL ---
cam_info = p.getDebugVisualizerCamera(physicsClientId=env.CLIENT)
cam_yaw = cam_info[8]
cam_pitch = cam_info[9]
cam_dist = cam_info[10]
cam_target = list(cam_info[11])
forward = np.array(cam_info[5]) # Camera forward direction
right = np.array(cam_info[6]) # Camera right direction
# Reset Z axis for flat movement in X/Y plane
forward[2] = 0
if np.linalg.norm(forward) > 0: forward = forward / np.linalg.norm(forward)
right[2] = 0
if np.linalg.norm(right) > 0: right = right / np.linalg.norm(right)
keys = p.getKeyboardEvents()
cam_speed = 0.05
moved = False
# Arrow Keys (Forward-Backward, Left-Right based on camera direction)
if p.B3G_UP_ARROW in keys and keys[p.B3G_UP_ARROW] & p.KEY_IS_DOWN:
cam_target[0] += forward[0] * cam_speed
cam_target[1] += forward[1] * cam_speed
moved = True
if p.B3G_DOWN_ARROW in keys and keys[p.B3G_DOWN_ARROW] & p.KEY_IS_DOWN:
cam_target[0] -= forward[0] * cam_speed
cam_target[1] -= forward[1] * cam_speed
moved = True
if p.B3G_RIGHT_ARROW in keys and keys[p.B3G_RIGHT_ARROW] & p.KEY_IS_DOWN:
cam_target[0] += right[0] * cam_speed
cam_target[1] += right[1] * cam_speed
moved = True
if p.B3G_LEFT_ARROW in keys and keys[p.B3G_LEFT_ARROW] & p.KEY_IS_DOWN:
cam_target[0] -= right[0] * cam_speed
cam_target[1] -= right[1] * cam_speed
moved = True
# Move Up (Space) / Move Down (Shift)
if 32 in keys and keys[32] & p.KEY_IS_DOWN:
cam_target[2] += cam_speed
moved = True
if p.B3G_SHIFT in keys and keys[p.B3G_SHIFT] & p.KEY_IS_DOWN:
cam_target[2] -= cam_speed
moved = True
if moved:
p.resetDebugVisualizerCamera(
cameraDistance=cam_dist,
cameraYaw=cam_yaw,
cameraPitch=cam_pitch,
cameraTargetPosition=cam_target,
physicsClientId=env.CLIENT
)
time.sleep(env.CTRL_TIMESTEP)
except KeyboardInterrupt:
print("Stopped by user.")
finally:
# --- CLOSE VIDEO RECORDING ---
p.stopStateLogging(log_id_tp)
p.stopStateLogging(log_id_fp)
env.close()
print("Videos saved: drone_third_person.mp4 and drone_first_person.mp4")
if __name__ == "__main__":
run_realtime_simulation() |