Add a regularization term to the objective of a stable-baseline3 model

Viewed 6

I'm using stable-baseline3's PPO implementation (see here) and wanted to play with the model a little bit further. More specifically, I wanted to add a regularization term to the objective which should (in theory) smooth out the (continuous) actions. Supposedly such goal would be achieved by performing something along the lines of

import gym
from stable_baselines3 import PPO

# Initializing PPO with some arbitrary env #
env = MyCustomEnv(...)
model = PPO("MlpPolicy", env)

# Now adding the supposed pseudo-code that adds regularization: #
model.regularization = lambda action,prev_action: np.linalg.norm(action-prev_action)

# Training the model w.r.t the original objective + the regularization: 
model.learn(total_timesteps=10_000)

as I couldn't find a way to do it, I've added this condition into the reward mechanism, inside the step function of my custom environment:

class MyCustomEnv(Env):

def __init__(self,...):
    self.action = 0
    self.max_action_diff = 0.1

def step(self,action):
    reward = ...
    ...
    action_norm = np.linalg.norm(self.action - action)
    if action_norm > self.max_action_diff:
        reward -= action_norm * reward

But this doesn't really work as I'm not quite sure how exactly to implement the penalty itself. Is this really the way to go? Am I restricted from playing with the objective?

0 Answers
Related