SAC and continuous control · Learning

Execution conditions

• Required packages: gymnasium, stable-baselines3, and PyTorch.
• SB3 policy: MlpPolicy.
• Gymnasium step: 5 return values.
• Public environment visualization: Add gymnasium[classic-control].

If the required package is missing, run %pip install control==0.10.2 gymnasium==1.3.0 stable-baselines3==2.9.2 in a Colab or Jupyter code cell.

Topics covered

Consider both exploration and input cost in a continuous-input control problem, and compare with a simple controller.

Continuous control

SAC learns a stochastic policy for continuous actions, considering both reward and entropy. Distinguish numerical action ranges from physical units such as torque or voltage.

Problem mapping

Map Pendulum's angle representation, reward, and time limit to the motor angular velocity tracking problem. Even with the same algorithm, different inputs and objectives require different evaluation metrics.

Worked example

Synthetic motor, dt=0.1 s, input $u\in[-2,2]$, and training reference angular velocity 1. Evaluate targets 0.5, 1, and 1.5.

"""Synthetic motor; not a hardware interface. dt=0.1 s, scalar speed and bounded input."""
import numpy as np
class Motor:
    def reset(self, seed=0, reference=1., noise=.01):
        self.rng = np.random.default_rng(seed)
        self.omega, self.t = 0., 0
        self.reference, self.noise = reference, noise
        return np.array([self.omega / 5, reference / 2], dtype=np.float32)
    def step(self, action):
        u = 2 * float(np.clip(np.asarray(action).item(), -1, 1))
        self.omega += .1 * (-.4 * self.omega + .8 * u + self.rng.uniform(-self.noise, self.noise))
        self.t += 1; error = self.reference - self.omega
        reward = -(error**2 + .05 * u**2)
        return np.array([self.omega / 5, self.reference / 2], dtype=np.float32), reward, abs(self.omega) > 5, self.t >= 100, {'error': error, 'u': u}

def motor_evaluate(policy, references=(.5, 1., 1.5)):
    rows=[]
    for reference in references:
        env=Motor(); obs=env.reset(100, reference)
        errors=[]; inputs=[]; total=0.
        while True:
            obs,r,terminated,truncated,info=env.step(policy(obs))
            errors.append(info['error']);inputs.append(info['u']);total+=r
            if terminated or truncated:break
        rows.append([reference, np.sqrt(np.mean(np.square(errors))),np.mean(np.square(inputs)),int(terminated),total])
    return np.array(rows)

def proportional(obs):return np.clip((2 * obs[1] - 5 * obs[0]), -1, 1)
print('reference / RMSE / mean u² / safety termination / Return')
print(motor_evaluate(proportional))

# Run the Queue or Motor definition first. Install gymnasium and stable-baselines3 if absent.
import gymnasium as gym
from gymnasium import spaces
from stable_baselines3 import A2C, PPO, SAC
class QueueGym(gym.Env):
    def __init__(self):
        self.model=Queue();self.action_space=spaces.Discrete(3)
        self.observation_space=spaces.Box(0.,1.,(1,),dtype=np.float32)
    def reset(self,seed=None,options=None):
        super().reset(seed=seed)
        return np.array([self.model.reset(seed)/20],dtype=np.float32),{}
    def step(self,a):
        q,r,te,tr,info=self.model.step(a)
        return np.array([q/20],dtype=np.float32),r/20,te,tr,info
class MotorGym(gym.Env):
    def __init__(self):
        self.model=Motor();self.action_space=spaces.Box(-1.,1.,(1,),dtype=np.float32)
        self.observation_space=spaces.Box(-np.inf,np.inf,(2,),dtype=np.float32)
    def reset(self,seed=None,options=None):
        super().reset(seed=seed)
        return self.model.reset(seed),{}
    def step(self,a):return self.model.step(a)

def fit_sb3(algorithm='PPO',seed=0,steps=2048):
    cls={'A2C':A2C,'PPO':PPO,'SAC':SAC}[algorithm]
    env=MotorGym() if algorithm=='SAC' else QueueGym()
    options={'policy_kwargs':dict(net_arch=[32,32]),'seed':seed,'verbose':0,'device':'cpu'}
    if algorithm=='PPO':options.update(n_steps=128,batch_size=64,n_epochs=3)
    if algorithm=='A2C':options.update(n_steps=32)
    if algorithm=='SAC':options.update(learning_starts=100,buffer_size=5000,batch_size=64)
    model=cls('MlpPolicy',env,**options);model.learn(total_timesteps=steps)
    return model

model=fit_sb3('SAC')
print(motor_evaluate(lambda obs:model.predict(obs,deterministic=True)[0]))

reference / RMSE / mean u² / safety termination / Return
[[  0.5          0.12626501   0.07338515   0.          -1.96121101]
 [  1.           0.25292968   0.29432127   0.          -7.86894852]
 [  1.5          0.3970162    0.6158097    0.         -18.84123509]]
[[  0.5          0.81118223   0.63537708   0.         -68.97854718]
 [  1.           0.42250626   0.65318108   0.         -21.11705962]
 [  1.5          0.42756512   0.65434743   0.         -21.55293066]]

Official documentation and data

• Pendulum · Environment definition
• SB3 · SAC

Reference example

View more examples

• Soft Actor-Critic (SAC)