File size: 1,161 Bytes
037569d | 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 | import time
import pybullet as p
from stable_baselines3 import PPO
from kuka_env import KukaReachingEnv
if __name__ == "__main__":
print("Loading Environment in GUI mode...")
env = KukaReachingEnv(render_mode="human")
print("Loading Vision-Trained Model...")
try:
model = PPO.load("models/ppo_vision_kuka", env=env)
except FileNotFoundError:
print("Error: Could not find the model.")
env.close()
exit()
obs, info = env.reset()
print("Watch the AI control the arm using only raw pixels!")
try:
for i in range(5000):
if not p.isConnected():
break
action, _states = model.predict(obs, deterministic=True)
obs, reward, terminated, truncated, info = env.step(action)
if terminated:
print(f"Target Hit! (Distance: {info['distance']:.3f})")
time.sleep(0.5)
obs, info = env.reset()
elif truncated:
obs, info = env.reset()
except KeyboardInterrupt:
pass
finally:
env.close() |