ppo_continuous_action doesn't work on Pendulum-v1?

[RishiMalhotra920]

I've tried running CleanRL/PPO on Pendulum-v1 and I find it impossible to get an optimal solution. However, similar experiments on stable-baselines3/ppo worked.

I've used the CleanRL/ppo command below and several hyperparameter variants. I have tried toggling the annealing, clip_vloss but I just don't see where I am going wrong.

python -m cleanrl.ppo_continuous_action --env-id Pendulum-v1 --exp-name ppo --seed 1 --track --wandb-project-name cleanRL --gamma 0.99 --vf-coef 0.5 --ent-coef 0.01 --norm-adv --num-envs
4  --clip-coef 0.2 --num-steps 2048 --clip-vloss --gae-lambda 0.95 --learning-rate 3e-4 --anneal-lr --max-grad-norm 0.5 --update-epochs 20 --num-minibatches 4 --total-timesteps 5000000 --torch-deterministic --capture-video --exp-name testing_video --target-kl 0.01

The script below with stable-baselines3 works achieves the optimal solution. However, sb3 is too complex to understand which is why I opted in for cleanRL.

import os

import gymnasium as gym
from stable_baselines3 import PPO


def train():
    # Create environment (no render during training)
    env = gym.make("Pendulum-v1")

    # Create the PPO model with tuned hyperparameters
    model = PPO(
        "MlpPolicy",
        env,
        # Larger network
        policy_kwargs=dict(net_arch=dict(pi=[256, 256], vf=[256, 256])),
        # Learning parameters
        learning_rate=1e-3,
        n_steps=2048,
        batch_size=64,
        n_epochs=20,
        gamma=0.99,
        gae_lambda=0.95,
        clip_range=0.2,
        # Increase entropy coefficient to encourage exploration
        ent_coef=0.01,
        verbose=1,
        device="auto",
    )

    # Train for longer (200k steps)
    model.learn(total_timesteps=200000)

    # Save the model
    os.makedirs("models", exist_ok=True)
    model.save("models/ppo_pendulum")

    env.close()

In addition, cleanRL/rpo_continuous works too.

python -m cleanrl.rpo_continuous_action --no-cuda --track --wandb-project-name cleanRL --env-id Pendulum-v1 --seed 5

So what's the issue with cleanRL/ppo_continuous_action? Am I missing something?

[noahfarr]

Did you try running CleanRL PPO with the sb3 hyperparameters?

[RishiMalhotra920]

Yes, I tried that @noahfarr and @vwxyzjn. That would be the same command as above but with lr 1e-3 instead of 3e-4 and hidden layers of 256, 256 instead of 64, 64.

I must be missing something critical because the rewards aren't increasing. Policy loss is around 0 and VF loss converges.

[noahfarr]

@RishiMalhotra920 did you check the NN architecture ? Are you using the same activations and weight initializations?

[RishiMalhotra920]

Yes I have. in sb3,

pi = self.action_net(self.mlp_extractor.policy_net(inputs))
vf = self.value_net(self.mlp_extractor.value_net(inputs))

self.mlp_extractor is MlpExtractor(
  (policy_net): Sequential(
    (0): Linear(in_features=3, out_features=256, bias=True)
    (1): Tanh()
    (2): Linear(in_features=256, out_features=256, bias=True)
    (3): Tanh()
  )
  (value_net): Sequential(
    (0): Linear(in_features=3, out_features=256, bias=True)
    (1): Tanh()
    (2): Linear(in_features=256, out_features=256, bias=True)
    (3): Tanh()
  )
)
self.value_net is Linear(in_features=256, out_features=1, bias=True)
self.action_net is Linear(in_features=256, out_features=1, bias=True)

In cleanRL, i have:

        hidden_dim = 256
        self.critic = nn.Sequential(
            layer_init(nn.Linear(np.array(envs.single_observation_space.shape).prod(), hidden_dim)),
            nn.Tanh(),
            layer_init(nn.Linear(hidden_dim, hidden_dim)),
            nn.Tanh(),
            layer_init(nn.Linear(hidden_dim, 1), std=1.0),
        )
        self.actor_mean = nn.Sequential(
            layer_init(nn.Linear(np.array(envs.single_observation_space.shape).prod(), hidden_dim)),
            nn.Tanh(),
            layer_init(nn.Linear(hidden_dim, hidden_dim)),
            nn.Tanh(),
            layer_init(nn.Linear(hidden_dim, np.prod(envs.single_action_space.shape)), std=0.01),
            nn.Tanh(),
        )
        self.actor_logstd = nn.Parameter(torch.zeros(1, np.prod(envs.single_action_space.shape)))

they both also use state-independent log-std for the policy net.

[noahfarr]

What about weight initialization in sb3?

[RishiMalhotra920]

They use orthogonal initialization, just like cleanRL.

[noahfarr]

Also the same gains? And constant initialization for the biases?

[noahfarr]

Also make sure the CleanRL uses the same wrappers around the environment.

[RishiMalhotra920]

Looking into it

[RishiMalhotra920]

@noahfarr yeah same initialization:

CleanRL:

        hidden_dim = 256
        self.critic = nn.Sequential(
            layer_init(nn.Linear(np.array(envs.single_observation_space.shape).prod(), hidden_dim)),
            nn.Tanh(),
            layer_init(nn.Linear(hidden_dim, hidden_dim)),
            nn.Tanh(),
            layer_init(nn.Linear(hidden_dim, 1), std=1.0),
        )
        self.actor_mean = nn.Sequential(
            layer_init(nn.Linear(np.array(envs.single_observation_space.shape).prod(), hidden_dim)),
            nn.Tanh(),
            layer_init(nn.Linear(hidden_dim, hidden_dim)),
            nn.Tanh(),
            layer_init(nn.Linear(hidden_dim, np.prod(envs.single_action_space.shape)), std=0.01),
        )
        self.actor_logstd = nn.Parameter(torch.zeros(1, np.prod(envs.single_action_space.shape)))

SB3:

if self.ortho_init:
            # TODO: check for features_extractor
            # Values from stable-baselines.
            # features_extractor/mlp values are
            # originally from openai/baselines (default gains/init_scales).
            module_gains = {
                self.features_extractor: np.sqrt(2),
                self.mlp_extractor: np.sqrt(2),
                self.action_net: 0.01,
                self.value_net: 1,
            }
            if not self.share_features_extractor:
                # Note(antonin): this is to keep SB3 results
                # consistent, see GH#1148
                del module_gains[self.features_extractor]
                module_gains[self.pi_features_extractor] = np.sqrt(2)
                module_gains[self.vf_features_extractor] = np.sqrt(2)

            for module, gain in module_gains.items():
                module.apply(partial(self.init_weights, gain=gain))

[noahfarr]

@RishiMalhotra920 and did you check whether sb3 uses the same environment wrappers as CleanRL

[RishiMalhotra920]

sorry i missed that. let me check

[RishiMalhotra920]

Yes.

cleanRL:

def make_env(env_id, idx, capture_video, run_name, gamma):
    def thunk():
        if capture_video and idx == 0:
            env = gym.make(env_id, render_mode="rgb_array")
            env = gym.wrappers.RecordVideo(env, f"videos/{run_name}")
        else:
            env = gym.make(env_id)
        env = gym.wrappers.FlattenObservation(env)  # deal with dm_control's Dict observation space
        env = gym.wrappers.RecordEpisodeStatistics(env)
        env = gym.wrappers.ClipAction(env)
        env = gym.wrappers.NormalizeObservation(env)
        env = gym.wrappers.TransformObservation(env, lambda obs: np.clip(obs, -10, 10))
        env = gym.wrappers.NormalizeReward(env, gamma=gamma)
        env = gym.wrappers.TransformReward(env, lambda reward: np.clip(reward, -10, 10))
        return env

    return thunk

sb3 action clipping:

        if isinstance(self.action_space, spaces.Box):
            if self.squash_output:
                # Rescale to proper domain when using squashing
                actions = self.unscale_action(actions)  # type: ignore[assignment, arg-type]
            else:
                # Actions could be on arbitrary scale, so clip the actions to avoid
                # out of bound error (e.g. if sampling from a Gaussian distribution)
                actions = np.clip(actions, self.action_space.low, self.action_space.high)  # type: ignore[assignment, arg-type]

sb3 normalize and clip obs:

def _normalize_obs(self, obs: np.ndarray, obs_rms: RunningMeanStd) -> np.ndarray:
        """
        Helper to normalize observation.
        :param obs:
        :param obs_rms: associated statistics
        :return: normalized observation
        """
        return np.clip((obs - obs_rms.mean) / np.sqrt(obs_rms.var + self.epsilon), -self.clip_obs, self.clip_obs)

sb3 normalize and clip reward:

    def normalize_reward(self, reward: np.ndarray) -> np.ndarray:
        """
        Normalize rewards using this VecNormalize's rewards statistics.
        Calling this method does not update statistics.
        """
        if self.norm_reward:
            reward = np.clip(reward / np.sqrt(self.ret_rms.var + self.epsilon), -self.clip_reward, self.clip_reward)
        # Note: we cast to float32 as it correspond to Python default float type
        # This cast is needed because `RunningMeanStd` keeps stats in float64
        return reward.astype(np.float32)

[RishiMalhotra920]

FWIW, here's graphs from my training run with this command. This is intended to be an exact replication of the SB3 script i shared in my first post on this issue.

python -m cleanrl.ppo_continuous_action --env-id Pendulum-v1 --exp-name ppo --seed 1 --track --wandb-project-name cleanRL --gamma 0.99 --vf-coef 0.5 --ent-coef 0.01 --norm-adv --num-envs 4  --clip-coef 0.2 --num-steps 2048 --clip-vloss --gae-lambda 0.95 --learning-rate 1e-3 --no-anneal-lr --max-grad-norm 0.5 --update-epochs 20 --num-minibatches 32 --total-timesteps 5000000 --torch-deterministic --capture-video --exp-name adjusted_minibatch_size