I am trying to use tianshou to integrate my env, but I am facing action_scaling can only be True when action_space is Box but got: Box(-1.0, 1.0, (2,), float32)

ENV

import sys
import os

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from market import GBMPriceSimulator
from amm.amm import AMM, SimpleFeeAMM
from amm.fee import PercentFee

from typing import Tuple
import numpy as np
from gym import spaces


EPSILON = 1e-5


class ArbitrageEnv:
    def __init__(self,
                 market: GBMPriceSimulator,
                 amm: AMM
                 ) -> None:
        self.amm = amm
        self.market = market
        self.cum_pnl = 0.
        self.observation_space = spaces.Box(low=np.array([0., 0., 0., 0.], dtype=np.float32),
                                            high=np.array([np.inf, np.inf, np.inf, np.inf], dtype=np.float32))
        self.action_space = spaces.Box(low=np.array([-1 + EPSILON, 0], dtype=np.float32),
                                                    high=np.array([1 - EPSILON, 1]), dtype=np.float32)

    # action: [trade_size_fraction, trade_decision (take opportunity when >0.5)]

    def step(self, action: np.array) -> Tuple[
        np.array, float, bool, bool, dict]:  # next_obs, rew, done, truncated, info
        trade_size_fraction = action[0]
        # Transform the second action component to the range [0, 1]
        trade_prob = (action[1] + 1) / 2  # Mapping from [-1, 1] to [0, 1]
        if trade_prob > 0.5:
            if trade_size_fraction > 0:
                asset_in, asset_out = 'B', 'A'
            else:
                asset_in, asset_out = 'A', 'B'
            size = abs(trade_size_fraction) * self.amm.portfolio[asset_out]
            print(f"asset_in: {asset_in}, asset_out: {asset_out}, size: {size}")
            success, info = self.amm.trade_swap(asset_in, asset_out, -size)  # take out -size shares of outing asset

            # calculate the reward
            asset_delta = info['asset_delta']
            fee = info['fee']
            print(f"asset_delta: {asset_delta}, fee: {fee}")
            print(self.amm)
            amm_order_cost = asset_delta[asset_in] + fee[asset_in]  # unit is always in B
            if asset_out == 'B':
                market_order_gain = abs(asset_delta[asset_out]) * self.market.current_price
            else:
                market_order_gain = abs(asset_delta[asset_out]) / self.market.current_price
            rew = market_order_gain - amm_order_cost
            print(f"market_order_gain: {market_order_gain} | "
                  f"amm_order_cost: {amm_order_cost} | "
                  f"reward: {rew}")
        else:
            success, info = True, {}
            rew = 0.

        self.cum_pnl += rew
        self.market.next()

        next_obs = self.get_obs()

        return next_obs, rew, False, not success, {'amm_trade_info': info}

    def reset(self):
        self.cum_pnl = 0
        self.amm.reset()
        self.market.reset()
        return self.get_obs()

    def get_obs(self) -> np.array:
        cur_market_price = self.market.current_price
        tmp = self.amm.portfolio
        cur_amm_price = tmp['A'] / tmp['B']

        return np.array([cur_market_price, cur_amm_price, tmp['A'], tmp['B']])

    def render(self, mode='human'):
        pass


if __name__ == '__main__':
    market = GBMPriceSimulator()
    amm_no_fee = SimpleFeeAMM(
        utility_func="constant_product",
        init_portfolio={'A': 1000, 'B': 1000, 'L': 1000},
        fee_structure=PercentFee(0.0)
    )
    env = ArbitrageEnv(market, amm_no_fee)

    # Reset the environment
    obs = env.reset()
    print(f"Initial observation: {obs}")

    # Perform a few steps with sample actions
    for step in range(5):
        # Generate a random action
        trade_size_fraction = np.random.uniform(-0.1, 0.1)  # Random number between -1 and 1
        trade_decision = np.random.uniform(0, 1)  # Random number between 0 and 1
        action = np.array([trade_size_fraction, trade_decision], dtype=np.float32)
        print(f"n--------------------"
              f"nStep {step + 1}:")
        print(f"action: {action}")
        next_obs, reward, done, truncated, info = env.step(action)
        print(f"Next observation: {next_obs}")
        print(f"Reward: {reward}")
        print(f"Done: {done}")
        print(f"Truncated: {truncated}")
        print(f"Info: {info}")

main.py

#!/usr/bin/env python3

import argparse
import datetime
import os
import sys
import pprint

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))


import numpy as np
import torch
from torch import nn
from torch.distributions import Distribution, Independent, Normal
from torch.optim.lr_scheduler import LambdaLR

from tianshou.data import Collector, ReplayBuffer, VectorReplayBuffer
from tianshou.highlevel.logger import LoggerFactoryDefault
from tianshou.policy import PPOPolicy
from tianshou.policy.base import BasePolicy
from tianshou.trainer import OnpolicyTrainer
from tianshou.utils.net.common import ActorCritic, Net
from tianshou.utils.net.continuous import ActorProb, Critic
from tianshou.env import DummyVectorEnv

from market import GBMPriceSimulator
from amm.amm import SimpleFeeAMM
from amm.fee import PercentFee
from AmmEnv import ArbitrageEnv

def get_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--task", type=str, default="ArbitrageEnv")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--buffer-size", type=int, default=4096)
    parser.add_argument("--hidden-sizes", type=int, nargs="*", default=[64, 64])
    parser.add_argument("--lr", type=float, default=3e-4)
    parser.add_argument("--gamma", type=float, default=0.99)
    parser.add_argument("--epoch", type=int, default=100)
    parser.add_argument("--step-per-epoch", type=int, default=30000)
    parser.add_argument("--step-per-collect", type=int, default=2048)
    parser.add_argument("--repeat-per-collect", type=int, default=10)
    parser.add_argument("--batch-size", type=int, default=64)
    parser.add_argument("--training-num", type=int, default=8)
    parser.add_argument("--test-num", type=int, default=10)
    parser.add_argument("--rew-norm", type=int, default=True)
    parser.add_argument("--vf-coef", type=float, default=0.25)
    parser.add_argument("--ent-coef", type=float, default=0.0)
    parser.add_argument("--gae-lambda", type=float, default=0.95)
    parser.add_argument("--bound-action-method", type=str, default="clip")
    parser.add_argument("--lr-decay", type=int, default=True)
    parser.add_argument("--max-grad-norm", type=float, default=0.5)
    parser.add_argument("--eps-clip", type=float, default=0.2)
    parser.add_argument("--dual-clip", type=float, default=None)
    parser.add_argument("--value-clip", type=int, default=0)
    parser.add_argument("--norm-adv", type=int, default=0)
    parser.add_argument("--recompute-adv", type=int, default=1)
    parser.add_argument("--logdir", type=str, default="log")
    parser.add_argument("--render", type=float, default=0.0)
    parser.add_argument(
        "--device",
        type=str,
        default="cuda" if torch.cuda.is_available() else "cpu",
    )
    parser.add_argument("--resume-path", type=str, default=None)
    parser.add_argument("--resume-id", type=str, default=None)
    parser.add_argument(
        "--logger",
        type=str,
        default="tensorboard",
        choices=["tensorboard", "wandb"],
    )
    parser.add_argument("--wandb-project", type=str, default="arbitrage.benchmark")
    parser.add_argument(
        "--watch",
        default=False,
        action="store_true",
        help="watch the play of pre-trained policy only",
    )
    return parser.parse_args()


def test_ppo(args: argparse.Namespace = get_args()) -> None:
    # Initialize the environment
    env = ArbitrageEnv(GBMPriceSimulator(), SimpleFeeAMM(
        utility_func="constant_product",
        init_portfolio={'A': 1000, 'B': 1000, 'L': 1000},
        fee_structure=PercentFee(0.0)
    ))

    # Create parallel environments
    train_envs = DummyVectorEnv([lambda: env for _ in range(args.training_num)])
    test_envs = DummyVectorEnv([lambda: env for _ in range(args.test_num)])

    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 = 1.0  # Since action is in [-1, 1] and [0, 1]

    print("Observations shape:", args.state_shape)
    print("Actions shape:", args.action_shape)
    print("Action range:", np.min(env.action_space.low), np.max(env.action_space.high))

    # Seed
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)

    # Model
    net_a = Net(
        args.state_shape,
        hidden_sizes=args.hidden_sizes,
        activation=nn.Tanh,
        device=args.device,
    )
    actor = ActorProb(
        net_a,
        args.action_shape,
        unbounded=True,
        device=args.device,
    ).to(args.device)
    net_c = Net(
        args.state_shape,
        hidden_sizes=args.hidden_sizes,
        activation=nn.Tanh,
        device=args.device,
    )
    critic = Critic(net_c, device=args.device).to(args.device)
    actor_critic = ActorCritic(actor, critic)

    torch.nn.init.constant_(actor.sigma_param, -0.5)
    for m in actor_critic.modules():
        if isinstance(m, torch.nn.Linear):
            torch.nn.init.orthogonal_(m.weight, gain=np.sqrt(2))
            torch.nn.init.zeros_(m.bias)
    for m in actor.mu.modules():
        if isinstance(m, torch.nn.Linear):
            torch.nn.init.zeros_(m.bias)
            m.weight.data.copy_(0.01 * m.weight.data)

    optim = torch.optim.Adam(actor_critic.parameters(), lr=args.lr)

    lr_scheduler = None
    if args.lr_decay:
        max_update_num = np.ceil(args.step_per_epoch / args.step_per_collect) * args.epoch
        lr_scheduler = LambdaLR(optim, lr_lambda=lambda epoch: 1 - epoch / max_update_num)

    def dist(loc_scale: tuple[torch.Tensor, torch.Tensor]) -> Distribution:
        loc, scale = loc_scale
        return Independent(Normal(loc, scale), 1)

    policy: PPOPolicy = PPOPolicy(
        actor=actor,
        critic=critic,
        optim=optim,
        dist_fn=dist,
        discount_factor=args.gamma,
        gae_lambda=args.gae_lambda,
        max_grad_norm=args.max_grad_norm,
        vf_coef=args.vf_coef,
        ent_coef=args.ent_coef,
        reward_normalization=args.rew_norm,
        action_scaling=False,
        action_bound_method=args.bound_action_method,
        lr_scheduler=lr_scheduler,
        action_space=env.action_space,
        eps_clip=args.eps_clip,
        value_clip=args.value_clip,
        dual_clip=args.dual_clip,
        advantage_normalization=args.norm_adv,
        recompute_advantage=args.recompute_adv,
    )

    # Load a previous policy if specified
    if args.resume_path:
        ckpt = torch.load(args.resume_path, map_location=args.device)
        policy.load_state_dict(ckpt["model"])
        train_envs.set_obs_rms(ckpt["obs_rms"])
        test_envs.set_obs_rms(ckpt["obs_rms"])
        print("Loaded agent from: ", args.resume_path)

    # Create collectors
    buffer: VectorReplayBuffer | ReplayBuffer
    if args.training_num > 1:
        buffer = VectorReplayBuffer(args.buffer_size, len(train_envs))
    else:
        buffer = ReplayBuffer(args.buffer_size)
    train_collector = Collector(policy, train_envs, buffer, exploration_noise=True)
    test_collector = Collector(policy, test_envs)

    # Logging setup
    now = datetime.datetime.now().strftime("%y%m%d-%H%M%S")
    args.algo_name = "ppo"
    log_name = os.path.join(args.task, args.algo_name, str(args.seed), now)
    log_path = os.path.join(args.logdir, log_name)

    # Logger
    logger_factory = LoggerFactoryDefault()
    if args.logger == "wandb":
        logger_factory.logger_type = "wandb"
        logger_factory.wandb_project = args.wandb_project
    else:
        logger_factory.logger_type = "tensorboard"

    logger = logger_factory.create_logger(
        log_dir=log_path,
        experiment_name=log_name,
        run_id=args.resume_id,
        config_dict=vars(args),
    )

    def save_best_fn(policy: BasePolicy) -> None:
        state = {"model": policy.state_dict(), "obs_rms": train_envs.get_obs_rms()}
        torch.save(state, os.path.join(log_path, "policy.pth"))

    if not args.watch:
        # Trainer
        result = OnpolicyTrainer(
            policy=policy,
            train_collector=train_collector,
            test_collector=test_collector,
            max_epoch=args.epoch,
            step_per_epoch=args.step_per_epoch,
            repeat_per_collect=args.repeat_per_collect,
            episode_per_test=args.test_num,
            batch_size=args.batch_size,
            step_per_collect=args.step_per_collect,
            save_best_fn=save_best_fn,
            logger=logger,
            test_in_train=False,
        ).run()
        pprint.pprint(result)

    # Watch performance
    test_envs.seed(args.seed)
    test_collector.reset()
    collector_stats = test_collector.collect(n_episode=args.test_num, render=args.render)
    print(collector_stats)


if __name__ == "__main__":
    test_ppo()

First, I got error “ValueError: action_scaling can only be True when action_space is Box but got: Box([-0.99999 0. ], [0.99999 1. ], (2,), float32)”, then I turn the action_scaling to false, and I got ” Unsupported action space: Box([-0.99999 0. ], [0.99999 1. ], (2,), float32).”