import numpy as np
import gymnasium as gym
from ppo_torch import Agent
import time, telegram, asyncio
from wrappers import make_atari, wrap_deepmind
from gymnasium.vector import SyncVectorEnv, AsyncVectorEnv

channel_id = "-1002031363666"
bot = telegram.Bot(token="<TRUFFLEHOG_REDACTED_TELEGRAMBOTTOKEN>")

async def send_message_to_channel(message_text):
    await bot.send_message(channel_id, message_text)

num_envs = 8
if __name__ == '__main__':
    def create_env(seed):
        def _init():
            env = make_atari("PongNoFrameskip-v4", 100000, render=None)
            env = wrap_deepmind(env, True, True, True, True)
            env.seed(seed)
            return env
        return _init

    envs = AsyncVectorEnv([create_env(i) for i in range(num_envs)])


    # agent = Agent([4, 84, 84], 64, env.action_space, eps_dec=1e-6, eps_end=0.02, max_mem_size=100000, gamma=0.99, epsilon=1.0, lr=1e-4)
    agent = Agent(envs.single_action_space.n, [4, 84, 84])

    num_episodes = 10000
    max_steps = 500000
    N = 20
    total_steps = 0

    print(f"Training for {num_episodes} episodes.\nMax steps per episode: {max_steps}.\nGLHF.")
    episode_rewards = []
    start_training_time = time.time()

    for episode in range(num_episodes):
        episode_training_time = time.time()
        observations, infos = envs.reset()
        # Transpose observations if necessary
        observations = np.transpose(observations, [0, 3, 1, 2])  # [num_envs, channels, height, width]

        total_rewards = np.zeros(num_envs)
        done_flags = np.zeros(num_envs, dtype=bool)
        step = 0

        while not np.all(done_flags) and step < max_steps:
            actions, probs, vals = agent.choose_action(observations)
            observations_, rewards, terminated, truncated, infos = envs.step(actions)
            dones = np.logical_or(terminated, truncated)
            observations_ = np.transpose(observations_, [0, 3, 1, 2])

            agent.remember(observations, actions, probs, vals, rewards, dones)

            if total_steps % N == 0:
                agent.learn()

            total_rewards += rewards
            total_steps += 1
            step += 1

            # Update the observations for environments that are not done
            observations = observations_
            done_flags = np.logical_or(done_flags, dones)

            # Handle environments that have finished
            if np.any(dones):
                # Reset the environments that are done
                indices = np.where(dones)[0]
                envs.reset_async(indices=indices)  # Initiate async reset for the done environments
                new_observations, infos = envs.reset_wait()  # Complete the reset and get observations
                new_observations = np.transpose(new_observations, [0, 3, 1, 2])
                observations[indices] = new_observations

        mean_reward = np.mean(total_rewards)
        episode_rewards.append(mean_reward)
        print(
            f"Episode {episode + 1} terminated after {time.time() - episode_training_time:.2f} seconds."
            f" Mean total reward: {mean_reward:.2f}."
            f" Average score over last 50 episodes: {np.mean(episode_rewards[-50:]):.2f}."
        )
        if (episode + 1) % 500 == 0:
            agent.save_models()



I am haveing troubles with the reset async function. Can you help me fix it?