Reward averages to 0 instead of increasing

ChrisAB: <p>Hello, I have a PPO network that is meant to receive noise, block error rate(BLER) and difference between current BLER and previous BLER, but reward is moving towards 0 instead of increasing.<br>
(It is <a href="https://pytorch.org/rl/stable/tutorials/coding_ppo.html" rel="noopener nofollow ugc">this example</a> but modified)</p>
<pre><code class="lang-auto">num_cells = 2048  # number of cells in each layer i.e. output dim.
lr = 1e-4
max_grad_norm = 1.0

frames_per_batch = 1000
# For a complete training, bring the number of frames up to 1M
total_frames = 1_00_000

sub_batch_size = 512  # cardinality of the sub-samples gathered from the current data in the inner loop
num_epochs = 10  # optimization steps per batch of data collected
clip_epsilon = (
    0.2  # clip value for PPO loss: see the equation in the intro for more context.
)
gamma = 0.99
lmbda = 0.95
entropy_eps = 1e-4

env = MatlabEnv(device=device)

env = TransformedEnv(
    env,
    Compose(
        # normalize observations
        UnsqueezeTransform(
            dim=-1,
            in_keys=["sinr", "blera", "blerc"],
            in_keys_inv=["sinr", "blera", "blerc"],
        ),
        CatTensors(in_keys=["sinr", "blera", "blerc"], dim=-1, out_key="observation", del_keys=False),
        ObservationNorm(in_keys=["observation"]),
        DoubleToFloat(),
        StepCounter(),
    ),
)
env.transform[2].init_stats(num_iter=1000, reduce_dim=0, cat_dim=0)

actor_net = nn.Sequential(
    nn.LazyLinear(num_cells, device=device),
    nn.Tanh(),
    nn.LazyLinear(num_cells, device=device),
    nn.Tanh(),
    nn.LazyLinear(num_cells, device=device),
    nn.Tanh(),
    nn.LazyLinear(2 * env.action_spec.shape[-1], device=device),
    NormalParamExtractor(),
)
policy_module = TensorDictModule(
    actor_net, in_keys=["observation"], out_keys=["loc", "scale"]
)

policy_module = ProbabilisticActor(
    module=policy_module,
    spec=env.action_spec,
    in_keys=["loc", "scale"],
    distribution_class=TanhNormal,
    distribution_kwargs={
        "low": env.action_spec.space.low,
        "high": env.action_spec.space.high,
    },
    return_log_prob=True,
    # we'll need the log-prob for the numerator of the importance weights
)

value_net = nn.Sequential(
    nn.LazyLinear(num_cells, device=device),
    nn.Tanh(),
    nn.LazyLinear(num_cells, device=device),
    nn.Tanh(),
    nn.LazyLinear(num_cells, device=device),
    nn.Tanh(),
    nn.LazyLinear(1, device=device),
)
value_module = ValueOperator(
    module=value_net,
    in_keys=["observation"],
)

collector = SyncDataCollector(
    env,
    policy_module,
    frames_per_batch=frames_per_batch,
    total_frames=total_frames,
    split_trajs=False,
    device=device,
)
replay_buffer = ReplayBuffer(
    storage=LazyTensorStorage(max_size=frames_per_batch),
    sampler=SamplerWithoutReplacement(),
)

advantage_module = GAE(
    gamma=gamma, lmbda=lmbda, value_network=value_module, average_gae=True
)
loss_module = ClipPPOLoss(
    actor_network=policy_module,
    critic_network=value_module,
    clip_epsilon=clip_epsilon,
    entropy_bonus=bool(entropy_eps),
    entropy_coef=entropy_eps,
    # these keys match by default but we set this for completeness
    critic_coef=1.0,
    loss_critic_type="smooth_l1",
)

optim = torch.optim.Adam(loss_module.parameters(), lr)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    optim, total_frames // frames_per_batch, 0.0
)

logs = defaultdict(list)
pbar = tqdm(total=total_frames)
eval_str = ""
</code></pre>
<p>And training loop:</p>
<pre><code class="lang-auto"># designed to collect:
for i, tensordict_data in enumerate(collector):
    # we now have a batch of data to work with. Let's learn something from it.
    for _ in range(num_epochs):
        # We'll need an "advantage" signal to make PPO work.
        # We re-compute it at each epoch as its value depends on the value
        # network which is updated in the inner loop.
        advantage_module(tensordict_data)
        data_view = tensordict_data.reshape(-1)
        # print(data_view)
        replay_buffer.extend(data_view.cpu())
        for _ in range(frames_per_batch // sub_batch_size):
            subdata = replay_buffer.sample(sub_batch_size)
            loss_vals = loss_module(subdata.to(device))
            loss_value = (
                loss_vals["loss_objective"]
                + loss_vals["loss_critic"]
                + loss_vals["loss_entropy"]
            )

            # Optimization: backward, grad clipping and optimization step
            loss_value.backward()
            # this is not strictly mandatory but it's good practice to keep
            # your gradient norm bounded
            torch.nn.utils.clip_grad_norm_(loss_module.parameters(), max_grad_norm)
            optim.step()
            optim.zero_grad()

    logs["reward"].append(tensordict_data["next", "reward"].mean().item())
    logs["sinr"].append(tensordict_data["sinr"].tolist()[0])
    pbar.update(tensordict_data.numel())
    cum_reward_str = (
        f"average reward={logs['reward'][-1]: 4.4f} (init={logs['reward'][0]: 4.4f})"
    )
    logs["step_count"].append(tensordict_data["step_count"].max().item())
    stepcount_str = f"step count (max): {logs['step_count'][-1]}"
    logs["lr"].append(optim.param_groups[0]["lr"])
    lr_str = f"lr policy: {logs['lr'][-1]: 4.4f}"
    if i % 10 == 0:
        # We evaluate the policy once every 10 batches of data.
        # Evaluation is rather simple: execute the policy without exploration
        # (take the expected value of the action distribution) for a given
        # number of steps (1000, which is our ``env`` horizon).
        # The ``rollout`` method of the ``env`` can take a policy as argument:
        # it will then execute this policy at each step.
        with set_exploration_type(ExplorationType.DETERMINISTIC), torch.no_grad():
            # execute a rollout with the trained policy
            eval_rollout = env.rollout(1000, policy_module)
            logs["eval reward"].append(eval_rollout["next", "reward"].mean().item())
            logs["eval reward (sum)"].append(
                eval_rollout["next", "reward"].sum().item()
            )
            logs["eval step_count"].append(eval_rollout["step_count"].max().item())
            logs["action"].append((eval_rollout["action"]).mean().item())
            eval_str = (
                f"eval cumulative reward: {logs['eval reward (sum)'][-1]: 4.4f} "
                f"(init: {logs['eval reward (sum)'][0]: 4.4f}), "
            )
            del eval_rollout
    pbar.set_description(", ".join([eval_str, cum_reward_str, stepcount_str, lr_str]))

    # We're also using a learning rate scheduler. Like the gradient clipping,
    # this is a nice-to-have but nothing necessary for PPO to work.
    scheduler.step()
</code></pre>
<p>My env step looks something like this:</p>
<pre><code class="lang-auto">sinr, blerc = tensordict["sinr"], tensordict["blerc"]
action = (tensordict["action"]).round().int().item()
# print(action)
result = get_bler_from_file([action], [sinr.item()])[0]
# print(result)
new_blerc = torch.tensor(result, dtype=torch.float32, device="cuda", requires_grad=True)
new_blera = new_blerc - blerc
spectral_efficiency = torch.tensor(float(se_table[0][action]*se_table[1][action]/1024), dtype=torch.float32)
# print("Action, spectral efficiency: ", action, " ", spectral_efficiency)
percentage_transmitted_blocks = (torch.tensor(float(1), dtype=torch.float32)-new_blerc)
percentage_spectral_efficiency = percentage_transmitted_blocks*spectral_efficiency
reward = percentage_spectral_efficiency
new_sinr = sinr + 0.1
if new_sinr.item() &gt; 20.5:
    new_sinr = new_sinr - 26.5
</code></pre>
<p>SINR is bound -6.5, 20.5</p>
<p>The NN always seems to converge to action = 1, no matter if reward is negative or positive. If negative, it will reach ~(-5.5) and if positive, it will reach 0. (As max reward can be 5.5)</p>
<p>Why does my network converge to 0 reward? I have tried increasing or decreasing lr or entropy to no avail, having more training time or more cells but nothing seems to help.</p>
<p>Results after training for 100k frames:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/4/f/4f49f318b656109cd04ed572c9cde3e2ee73c556.png" data-download-href="https://discuss.pytorch.org/uploads/default/4f49f318b656109cd04ed572c9cde3e2ee73c556" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/4/f/4f49f318b656109cd04ed572c9cde3e2ee73c556_2_517x500.png" alt="image" data-base62-sha1="bjq9HmWIT1cV4Pxfq3LV6yNwyJo" width="517" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/4/f/4f49f318b656109cd04ed572c9cde3e2ee73c556_2_517x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/4/f/4f49f318b656109cd04ed572c9cde3e2ee73c556_2_775x750.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/4/f/4f49f318b656109cd04ed572c9cde3e2ee73c556.png 2x" data-dominant-color="F2F6F8"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">848×820 72.8 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>