Add compute_episodic_return_on_last_step

[QuantuMope]

Currently, MC returns only get computed when encountering a discount of 0. This makes it so that MC returns cannot be computed if we do not terminate on success. This PR adds a flag compute_episodic_return_on_last_step to instead compute returns upon step_type.LAST instead.

[emailweixu]

Currently, MC returns only get computed when encountering a discount of 0. This makes it so that MC returns cannot be computed if we do not terminate on success. This PR adds a flag compute_episodic_return_on_last_step to instead compute returns upon step_type.LAST instead.

Why don't we set discount to 0 for LAST step?

[Haichao-Zhang]

This makes it so that MC returns cannot be computed if we do not terminate on success

1. keep adding more flags might not be ideal and might be a bit confusing to users?
2. This makes it so that MC returns cannot be computed if we do not terminate on success --> not sure what kind of use case you are considering, but typically, upon task success, we set step type as LAST + discount as 0?

[QuantuMope]

Currently, MC returns only get computed when encountering a discount of 0. This makes it so that MC returns cannot be computed if we do not terminate on success. This PR adds a flag compute_episodic_return_on_last_step to instead compute returns upon step_type.LAST instead.

Why don't we set discount to 0 for LAST step?

This is for RL setups where we don't terminate on success and only terminate upon timeout.
For EmbodiedGen resets, batch resetting all environments at the same step is often more efficient and sometimes the only possible strategy (e.g., if lighting randomization is turned on, cannot reset individual envs on GPU).

[QuantuMope]

This makes it so that MC returns cannot be computed if we do not terminate on success

1. keep adding more flags might not be ideal and might be a bit confusing to users?
2. This makes it so that MC returns cannot be computed if we do not terminate on success --> not sure what kind of use case you are considering, but typically, upon task success, we set step type as LAST + discount as 0?

1. Would making this the default behavior be better?
2. Most of our OpenVLAPPO, PPOFlow, SACFlow RL training uses timeout termination only. This would model an infinite horizon RL formulation.

[emailweixu]

Why can't we set discount to 0 at the time of setting step type to LAST?

[QuantuMope]

Why can't we set discount to 0 at the time of setting step type to LAST?

Accidentally editted the question comment instead of replying. Simply fixing that.

I think in practice we could, but it would be less correct to do so? We're essentially modeling an infinite time horizon problem. Applying a discount of 0 at time out when there were numerous success idle steps prior may cause confusion to the model?

[QuantuMope]

After some investigating, this flag is necessary to maintain correctness during TD bootstrapping for the infinite horizon RL formulation.

There is an important caveat where if compute_episodic_return_on_last_step==True, the computed MC return can be heavily biased. Added an explanation of why in commit 35a6b6d.