Use clearer function

Author: alok

pg/main.py

@@ -69,12 +69,12 @@ def G(rewards, start=0, end=None):
  ]
 
  states = torch.stack(states)
- state_values = critic(states).reshape(-1)
+ state_values = critic(states).flatten()
 
  cumulative_returns = tensor(cumulative_returns)
  Adv = cumulative_returns - state_values
 
- log_probs = torch.stack(log_probs).reshape(-1)
+ log_probs = torch.stack(log_probs).flatten()
 
  loss = -(Adv @ log_probs) / len(rewards)
  if episode > 500 and loss.item() < -1000:

ppo/ppo.py

@@ -122,7 +122,7 @@ def train(model, old_model, data) -> float:
  [sum(discounted_rewards[t:]) for t, _ in enumerate(discounted_rewards)]
  )
 
- state_values = model.vf(states).reshape(-1)
+ state_values = model.vf(states).flatten()
 
  adv = cumulative_returns - state_values
  vf_loss = F.mse_loss(state_values, cumulative_returns)