tianshou

import gymnasium as gym from tianshou.data import Collector, VectorReplayBuffer from tianshou.env import DummyVectorEnv from tianshou.policy import DQNPolicy from tianshou.trainer import offpolicy_trainer import torch env = gym.make('CartPole-v1') train_envs = DummyVectorEnv([lambda: gym.make('CartPole-v1') for _ in range(10)]) test_envs = DummyVectorEnv([lambda: gym.make('CartPole-v1') for _ in range(10)]) state_shape = env.observation_space.shape or env.observation_space.n action_shape = env.action_space.shape or env.action_space.n policy = DQNPolicy( state_shape=state_shape, action_shape=action_shape, model=torch.nn.Linear(state_shape, action_shape), optim=torch.optim.Adam(policy_network.parameters(), lr=1e-3), ).to('cpu') buffer = VectorReplayBuffer(total_size=20000, buffer_num=len(train_envs)) collector = Collector(policy, train_envs, buffer) test_collector = Collector(policy, test_envs) def stop_fn(reward): return reward >= 195 result = offpolicy_trainer( policy=policy, train_collector=collector, test_collector=test_collector, max_epoch=10, step_per_epoch=1000, step_per_collect=10, episode_per_test=5, batch_size=64, stop_fn=stop_fn, update_per_step=1, ) print(f'Finished training in {result.timing.total_time_seconds}s')