text: import os
  import gym
  import torch
  import pprint
  import argparse
  import numpy as np
  from torch.utils.tensorboard import SummaryWriter
  
  from tianshou.policy import SACPolicy
  from tianshou.utils import BasicLogger
  from tianshou.utils.net.common import Net
  from tianshou.env import SubprocVectorEnv
  from tianshou.trainer import offpolicy_trainer
  from tianshou.data import Collector, VectorReplayBuffer
  from tianshou.utils.net.continuous import ActorProb, Critic
  
  
  def get_args():
      parser = argparse.ArgumentParser()
      parser.add_argument('--task', type=str, default="BipedalWalkerHardcore-v3")
      parser.add_argument('--seed', type=int, default=0)
      parser.add_argument('--buffer-size', type=int, default=1000000)
      parser.add_argument('--actor-lr', type=float, default=3e-4)
      parser.add_argument('--critic-lr', type=float, default=1e-3)
      parser.add_argument('--gamma', type=float, default=0.99)
      parser.add_argument('--tau', type=float, default=0.005)
      parser.add_argument('--alpha', type=float, default=0.1)
      parser.add_argument('--auto-alpha', type=int, default=1)
      parser.add_argument('--alpha-lr', type=float, default=3e-4)
      parser.add_argument('--epoch', type=int, default=100)
      parser.add_argument('--step-per-epoch', type=int, default=100000)
      parser.add_argument('--step-per-collect', type=int, default=10)
      parser.add_argument('--update-per-step', type=float, default=0.1)
      parser.add_argument('--batch-size', type=int, default=128)
      parser.add_argument('--hidden-sizes', type=int,
                          nargs='*', default=[128, 128])
      parser.add_argument('--training-num', type=int, default=10)
      parser.add_argument('--test-num', type=int, default=100)
      parser.add_argument('--logdir', type=str, default='log')
      parser.add_argument('--render', type=float, default=0.)
      parser.add_argument('--n-step', type=int, default=4)
      parser.add_argument(
          '--device', type=str,
          default='cuda' if torch.cuda.is_available() else 'cpu')
      parser.add_argument('--resume-path', type=str, default=None)
      return parser.parse_args()
  
  
  class Wrapper(gym.Wrapper):
      """Env wrapper for reward scale, action repeat and removing done penalty"""
  
      def __init__(self, env, action_repeat=3, reward_scale=5, rm_done=True):
          super().__init__(env)
          self.action_repeat = action_repeat
          self.reward_scale = reward_scale
          self.rm_done = rm_done
  
      def step(self, action):
          r = 0.0
          for _ in range(self.action_repeat):
              obs, reward, done, info = self.env.step(action)
              # remove done reward penalty
              if not done or not self.rm_done:
                  r = r + reward
              if done:
                  break
          # scale reward
          return obs, self.reward_scale * r, done, info
  
  
  def test_sac_bipedal(args=get_args()):
      env = Wrapper(gym.make(args.task))
  
      args.state_shape = env.observation_space.shape or env.observation_space.n
      args.action_shape = env.action_space.shape or env.action_space.n
      args.max_action = env.action_space.high[0]
  
      train_envs = SubprocVectorEnv([
          lambda: Wrapper(gym.make(args.task))
          for _ in range(args.training_num)])
      # test_envs = gym.make(args.task)
      test_envs = SubprocVectorEnv([
          lambda: Wrapper(gym.make(args.task), reward_scale=1, rm_done=False)
          for _ in range(args.test_num)])
  
      # seed
      np.random.seed(args.seed)
      torch.manual_seed(args.seed)
      train_envs.seed(args.seed)
      test_envs.seed(args.seed)
  
      # model
      net_a = Net(args.state_shape, hidden_sizes=args.hidden_sizes,
                  device=args.device)
      actor = ActorProb(
          net_a, args.action_shape, max_action=args.max_action,
          device=args.device, unbounded=True).to(args.device)
      actor_optim = torch.optim.Adam(actor.parameters(), lr=args.actor_lr)
  
      net_c1 = Net(args.state_shape, args.action_shape,
                   hidden_sizes=args.hidden_sizes,
                   concat=True, device=args.device)
      critic1 = Critic(net_c1, device=args.device).to(args.device)
      critic1_optim = torch.optim.Adam(critic1.parameters(), lr=args.critic_lr)
  
      net_c2 = Net(args.state_shape, args.action_shape,
                   hidden_sizes=args.hidden_sizes,
                   concat=True, device=args.device)
      critic2 = Critic(net_c2, device=args.device).to(args.device)
      critic2_optim = torch.optim.Adam(critic2.parameters(), lr=args.critic_lr)
  
      if args.auto_alpha:
          target_entropy = -np.prod(env.action_space.shape)
          log_alpha = torch.zeros(1, requires_grad=True, device=args.device)
          alpha_optim = torch.optim.Adam([log_alpha], lr=args.alpha_lr)
          args.alpha = (target_entropy, log_alpha, alpha_optim)
  
      policy = SACPolicy(
          actor, actor_optim, critic1, critic1_optim, critic2, critic2_optim,
          tau=args.tau, gamma=args.gamma, alpha=args.alpha,
          estimation_step=args.n_step, action_space=env.action_space)
      # load a previous policy
      if args.resume_path:
          policy.load_state_dict(torch.load(args.resume_path))
          print("Loaded agent from: ", args.resume_path)
  
      # collector
      train_collector = Collector(
          policy, train_envs,
          VectorReplayBuffer(args.buffer_size, len(train_envs)),
          exploration_noise=True)
      test_collector = Collector(policy, test_envs)
      # train_collector.collect(n_step=args.buffer_size)
      # log
      log_path = os.path.join(args.logdir, args.task, 'sac')
      writer = SummaryWriter(log_path)
      logger = BasicLogger(writer)
  
      def save_fn(policy):
          torch.save(policy.state_dict(), os.path.join(log_path, 'policy.pth'))
  
      def stop_fn(mean_rewards):
          return mean_rewards >= env.spec.reward_threshold
  
      # trainer
      result = offpolicy_trainer(
          policy, train_collector, test_collector, args.epoch,
          args.step_per_epoch, args.step_per_collect, args.test_num, args.batch_size,
          update_per_step=args.update_per_step, test_in_train=False,
          stop_fn=stop_fn, save_fn=save_fn, logger=logger)
  
      if __name__ == '__main__':
          pprint.pprint(result)
          # Let's watch its performance!
          policy.eval()
          test_envs.seed(args.seed)
          test_collector.reset()
          result = test_collector.collect(n_episode=args.test_num,
                                          render=args.render)
          rews, lens = result["rews"], result["lens"]
          print(f"Final reward: {rews.mean()}, length: {lens.mean()}")
  
  
  if __name__ == '__main__':
      test_sac_bipedal()
  