import gym
from stable_baselines3 import A2C
from gymnasium.envs.registration import register
from huggingface_sb3 import push_to_hub


config = {
    "policy_type": "MlpPolicy",
    "total_timesteps": 5000,
    "env_name": "CartPole-v1",
    }


register(id="CartPole-v1", entry_point='gym.envs.classic_control:CartPoleEnv', max_episode_steps=200)

env = gym.make("CartPole-v1", render_mode="rgb_array")

model = A2C("MlpPolicy", env, verbose=1)
model.learn(total_timesteps=config["total_timesteps"])

model.save("CartPole-v1.zip")

push_to_hub(
    repo_id="jossuaseksik/a2c-sb3_cartpole",
    filename="CartPole-v1.zip",
    commit_message="Adding CartPole model trained with A2C on HuggingFace",
    token="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"
)