test.py

import gymnasium as gym
import numpy as np
from gym_modifly.env import DroneEnv  # Import the custom drone environment

# Initialize the environment
env = DroneEnv()

episodes = 5  # Number of episodes
for episode in range(episodes):
    state, _ = env.reset()
    done = False
    truncate = False
    step = 0
    total_reward = 0
    
    while not (done or truncate):
        action = env.action_space.sample()  # Select a random action
        next_state, reward, done, truncate, _ = env.step(action)
        total_reward += reward
        step += 1
        print(f"Episode {episode + 1}, Step {step}: State={next_state}, Reward={reward}")
    
    print(f"Episode {episode + 1} finished with total reward: {total_reward}")

env.close()