Closed Armandpl closed 5 months ago
I am confused about how to align the advantage with the actions.
Since we predict the reward and compute the lambda value and the advantage from the world model state ht+zt
and since ht+zt
is the result of the action at-1
I think the advantage at t
should be used to push the log prob of the action that is responsible for this advantage, and that's at-1
. That's how the code should look like:
policy = actor(
sg(hts[:, :-2]),
sg(zts[:, :-2]),
)
logpi = policy.log_prob(sg(ats[:, :-2]).squeeze(-1))
actor_loss = -logpi * sg(advantage[:, 1:].squeeze(-1)) # offset the advantage by one
actor_entropy = policy.entropy()
actor_loss -= ACTOR_ENTROPY * actor_entropy
actor_loss = actor_loss * sg(traj_weight[:, :-2])
actor_loss = actor_loss.mean()
But doing this the training collapses (wandb run):
However, if I use the advantage at
t
to push the log prob of at
:
policy = actor(
sg(hts[:, :-1]),
sg(zts[:, :-1]),
)
logpi = policy.log_prob(sg(ats[:, :-1]).squeeze(-1))
actor_loss = -logpi * sg(advantage.squeeze(-1))
actor_entropy = policy.entropy()
actor_loss -= ACTOR_ENTROPY * actor_entropy
actor_loss = actor_loss * sg(traj_weight[:, :-1])
actor_loss = actor_loss.mean()
It now works (wandb run):
ok so I trained on pong overnight, for 150k step with a training ratio of 1024. It learns some stuff but it seems more unstable than the official scores and reaches a way lower return (-10 instead of 20). Why is that?
todo:
move from batch first to time first to make the code slightly more readablenot worth the effort nowneed a quick e2e training test so that I make sure I don't break stuff when refactoringremove debug scripts-> wait a little bit longer just in case, wait til we've trained the robot, also remove test envs if neededmemap the replay buffer?-> 100k atari steps fits on 64GB RAM, ok for nowadd back first obs to zt-> eh not super important i feel we can lose one obs?? investigate later if neededadd layer norm to CNN?-> not super important we mostly want to train on the pendulum so