DDPG Tutorial and Custom Environment

ItsVik02: <p>Hi everybody, I trying to implement my own DDPG agent to control an unstable system taken from MATLAB. In particular, my goal is to stabilize some plasma velocity, keeping a low current in the control circuit and using a limited control action (the tension applied to such circuit).<br>
As references, i used this tutorial for the agent : <a href="https://pytorch.org/tutorials/intermediate/coding_ddpg.html" class="inline-onebox" rel="noopener nofollow ugc">TorchRL objectives: Coding a DDPG loss — PyTorch Tutorials 2.0.1+cu117 documentation</a><br>
And this one for the environment: <a href="https://pytorch.org/tutorials/advanced/pendulum.html" class="inline-onebox" rel="noopener nofollow ugc">Pendulum: Writing your environment and transforms with TorchRL — PyTorch Tutorials 2.4.0+cu121 documentation</a><br>
The program runs correctly but the agent doesn’t learn to solve the task. It’s been a few weeks of debugging, however I don’t have a clue where the problem is at.  I’ll share with you the Environment as the DDPG part of my code (included training) is identical to the tutorial.<br>
A few differences from the tutorials is that I don’t use the “get_env_stats()” method as I already know the mean and variance of my observations, and that I’m not using multiple environments in parallel (env_per_collector = 1), because I get some batch size error that I wasn’t able to figure out. Thanks in advance  for any reply!</p>
<p><span class="hashtag-raw">#MATLAB</span><br>
<span class="hashtag-raw">#C2D</span> ‘ZOH’ Ts = 0.0025<br>
A = torch.tensor(np.array([[0.9755, -168.6], [-6.01e-07, 1.01]])).to(dtype=torch.float32)<br>
B = torch.tensor(np.array([[2.037], [ 5.004e-05]])).to(dtype=torch.float32)<br>
C = torch.tensor(np.identity(2)).to(dtype=torch.float32)<br>
D = torch.tensor(np.zeros((2, 1))).to(dtype=torch.float32)</p>
<p><span class="hashtag-raw">#Constants</span><br>
is_fork = torch.multiprocessing.get_start_method() == “fork”<br>
device = (<br>
torch.device(0)<br>
if torch.cuda.is_available() and not is_fork<br>
else torch.device(“cpu”)<br>
)</p>
<p>collector_device = “cpu”  # or “CUDA”<br>
k1 = 1<br>
k2 = 2<br>
k3 = 1<br>
MAX_SPEED = 10.0  # [m/s]<br>
MAX_VOLTAGE = 2300.0  # [V]<br>
MAX_CURRENT = 10_000.0  # [A] (max current for the vs circuit)<br>
CURRENT_THRESHOLD = 500<br>
SPEED_THRESHOLD = 0.5</p>
<p>def generate_params(batch_size=None) → TensorDictBase:<br>
if batch_size is None:<br>
batch_size = <span class="chcklst-box fa fa-square-o fa-fw"></span><br>
tens_dict = TensorDict(<br>
{<br>
“params”: TensorDict(<br>
{<br>
“max_speed”: MAX_SPEED,<br>
“max_current”: MAX_CURRENT,<br>
“max_voltage”: MAX_VOLTAGE,<br>
},<br>
<span class="chcklst-box fa fa-square-o fa-fw"></span><br>
)<br>
},<br>
<span class="chcklst-box fa fa-square-o fa-fw"></span>,<br>
)<br>
if batch_size:<br>
tens_dict = tens_dict.expand(batch_size).contiguous()<br>
return tens_dict</p>
<p>def _reset(self, tensordict):<br>
<span class="hashtag-raw">#print</span>(‘resetting’)</p>
<pre><code>if tensordict is None or tensordict.is_empty():
    tensordict = self.generate_params()

max_init_current = torch.tensor(100, device=self.device)
max_init_speed = torch.tensor(0.1, device=self.device)
min_init_current = -max_init_current
min_init_speed = -max_init_speed
ivs3 = (
    torch.rand(tensordict.shape, generator=self.rng, device=self.device)
    * (max_init_current-min_init_current) + min_init_current
)
zc_dot = (
    torch.rand(tensordict.shape, generator=self.rng, device=self.device)
    * (max_init_speed-min_init_speed) + min_init_speed
)

VS_Env.current_time = 0

shape = torch.zeros(size=(*tensordict.shape, 1))
done = torch.zeros_like(shape, dtype=torch.bool)
#terminated = torch.zeros_like(shape, dtype=torch.bool)

out = TensorDict(
    {
        "ivs3": ivs3,
        "zc_dot": zc_dot,
        "params": tensordict["params"],
        "done": done,
        #"terminated": terminated,
    },
    batch_size=tensordict.shape
)
return out
</code></pre>
<p><span class="hashtag-raw">#ENV</span> SPEC:<br>
def make_composite_from_td(td):<br>
# Custom function to convert a <code>tensordict</code> in a similar spec structure<br>
# of unbounded values.<br>
composite = CompositeSpec(<br>
{<br>
key: make_composite_from_td(tensor)<br>
if isinstance(tensor, TensorDictBase)<br>
else UnboundedContinuousTensorSpec(<br>
dtype=tensor.dtype, device=tensor.device, shape=tensor.shape<br>
)<br>
for key, tensor in td.items()<br>
},<br>
shape=td.shape,<br>
)<br>
return composite</p>
<p>def _make_spec(self, td_params):<br>
# STATE:<br>
self.observation_spec = CompositeSpec(<br>
zc_dot=BoundedTensorSpec(<br>
low=-td_params[“params”, “max_speed”],<br>
high=td_params[“params”, “max_speed”],<br>
shape=(),<br>
dtype=torch.float32,<br>
),<br>
ivs3=BoundedTensorSpec(<br>
low=-td_params[“params”, “max_current”],<br>
high=td_params[“params”, “max_current”],<br>
shape=(),<br>
dtype=torch.float32,<br>
),<br>
# We need to add the <code>params</code> to the observation specs, as we want<br>
# to pass it at each step during a rollout<br>
params=make_composite_from_td(td_params[“params”]),<br>
shape=(),<br>
)<br>
# Since the environment is stateless, we expect the previous output as input.<br>
# For this, <code>EnvBase</code> expects some state_spec to be available<br>
self.state_spec = self.observation_spec.clone()</p>
<pre><code>#ACTION:
self.action_spec = BoundedTensorSpec(
    low=-td_params["params", "max_voltage"],
    high=td_params["params", "max_voltage"],
    shape=(1,),
    dtype=torch.float32,
)
# REWARD:
self.reward_spec = UnboundedContinuousTensorSpec(shape=(*td_params.shape, 1))
</code></pre>
<p><span class="hashtag-raw">#SEEDING:</span></p>
<p>def _set_seed(self, seed: Optional[int]):<br>
rng = torch.manual_seed(seed)  # Returns a torch.Generator object<br>
self.rng = rng</p>
<p>class VS_Env(EnvBase):<br>
metadata = {<br>
“render_modes”: [“human”, “rgb_array”],<br>
“render_fps”: 30,<br>
}<br>
batch_locked = False<br>
current_time = 0<br>
total_time = 2.5<br>
time_step = 0.0025</p>
<pre><code>def __init__(self, td_params=None, seed=None, device="cpu"):
    if td_params is None:
        td_params = generate_params()
    super().__init__(device=device, batch_size=[])
    self._make_spec(td_params)
    if seed is None:
        seed = torch.empty((), dtype=torch.int64).random_().item()
    self.set_seed(seed)


# Helpers: _make_spec and gen_params
generate_params = staticmethod(generate_params)
_make_spec = _make_spec

# Mandatory methods: _step, _reset and _set_seed
_reset = _reset
#_step = staticmethod(_step)
_set_seed = _set_seed
@staticmethod
def _step(tensordict):
    # Var di stato per bsys
    ivs3, zc_dot = tensordict["ivs3"], tensordict["zc_dot"]

    max_speed = tensordict["params", "max_speed"]
    max_current = tensordict["params", "max_current"]
    max_voltage = tensordict["params", "max_voltage"]
    # Action
    u = tensordict["action"].squeeze(-1)
    u = u.clamp(-max_voltage, max_voltage)
    shape = torch.zeros(size=(*tensordict.shape, 1))

    cost_function = (k1 * ((zc_dot / max_speed).pow(2)) + k2 * ((ivs3 / max_current).pow(2)) +
                     k3 * ((u / max_voltage).pow(2)))

    #terminated = torch.zeros_like(shape, dtype=torch.bool)
    #truncated = torch.zeros_like(shape, dtype=torch.bool)
    done = torch.zeros_like(shape, dtype=torch.bool)
    
    state = torch.stack((ivs3, zc_dot)).squeeze(-1)

    new_ivs3 = (A[0, 0] * state[0] + A[0, 1] * state[1] + B[0] * u).squeeze(-1)
    new_zc_dot = (A[1, 0] * state[0] + A[1, 1] * state[1] + B[1] * u).squeeze(-1)


    if abs(new_ivs3) &gt; max_current or abs(new_zc_dot) &gt; max_speed:
        done = torch.ones_like(shape, dtype=torch.bool)
        cost_function = cost_function + 10
        print('disruption  -&gt;', new_ivs3)
        if abs(new_zc_dot) &gt; max_speed:
            print('disruption vel -&gt;', new_zc_dot)

    """
    if abs(new_ivs3) &lt; CURRENT_THRESHOLD and abs(new_zc_dot) &lt; SPEED_THRESHOLD:
        cost_function = cost_function - 2
    """

    if VS_Env.current_time &gt;= VS_Env.total_time:
        done = torch.ones_like(shape, dtype=torch.bool)
        print('done')

    VS_Env.current_time = VS_Env.current_time + VS_Env.time_step
    """
    if abs(new_ivs3) &gt; 5000 or abs(zc_dot) &gt; 5:
        cost_function = cost_function + 5
    

    if abs(new_ivs3) &gt; max_current or abs(new_zc_dot) &gt; max_speed:
        truncated = torch.ones_like(shape, dtype=torch.bool)
        cost_function = cost_function + 10
    if abs(new_ivs3) &lt; CURRENT_THRESHOLD and abs(new_zc_dot) &lt; SPEED_THRESHOLD:
        cost_function = cost_function - 2
    if VS_Env.current_time &gt;= VS_Env.total_time:
        terminated = torch.ones_like(shape, dtype=torch.bool)
    """

    reward = -cost_function.view(*tensordict.shape, 1).to(dtype=torch.float32)

    # Solo reward e terminated devono essere torch.Size([1]), il resto deve essere
    # torch.Size([])

    # Uscita di step
    out = TensorDict(
        {
            "ivs3": new_ivs3,
            "zc_dot": new_zc_dot,
            "params": tensordict["params"],
            "reward": reward,
            #"terminated": terminated,
            #"truncated": truncated,
            "done": done,
        },
        tensordict.shape,
    )
    return out
</code></pre>
<p>def make_env():<br>
env = VS_Env()<br>
return env</p>
<p><span class="hashtag-raw">#Transforming</span> the env:</p>
<p>def env_transform(env):<br>
reward_scaling = 5<br>
<span class="hashtag-raw">#max_steps</span> = 2_000</p>
<pre><code>env = TransformedEnv(
    env,
    # ``Unsqueeze`` the observations that we will concatenate
    UnsqueezeTransform(
        dim=-1,
        in_keys=["ivs3", "zc_dot"],
        in_keys_inv=["ivs3", "zc_dot"],
    ),
)

env.append_transform(RewardScaling(loc=0.0, scale=reward_scaling))

cat_transform = CatTensors(
    in_keys=["ivs3", "zc_dot"], dim=-1, out_key="observation_vector", del_keys=False
)
env.append_transform(cat_transform)
env.append_transform(ObservationNorm(in_keys=["observation_vector"], loc=[-175, 0], scale=[1.11e+04, 2.8305], standard_normal=True))
env.append_transform(DoubleToFloat())
#env.append_transform(StepCounter(max_steps))
# We need a marker for the start of trajectories for our Ornstein-Uhlenbeck (OU)
# exploration:
env.append_transform(InitTracker())
return env
</code></pre>