Explicit distinction between done and truncated needed in PPOTransition?

[emergenz]

Above question. Reason being that currently, the 'Stoix API' is not consistent across environments. To give an example. When using a navix environment:

• timestep.last() returns True when the transition is either a truncation-transition (i.e. the episode ended due to max_steps), or a termination-transition (i.e. the episode ended due to achieving the goal).
• PPOTransition.done is True when the transition is a termination-transition
• PPOTransition.truncated is True when the transition is a truncation-transition

I think this is intended behaviour from the Stoix-side. However, when using a gymnax environment, the behaviour changes to the following:

• PPOTransition.done is True when the transition is a truncation- OR termination-transition (so the same as timestep.last(), this becomes evident when looking at GymnaxWrapper, where discount and step_type are assigned using the same underlying logic:
  

  Stoix/stoix/wrappers/gymnax.py

  Lines 89 to 90 in 11dca7d

  	discount=jnp.array(1.0 - done, dtype=float),
  	step_type=jax.lax.select(done, StepType.LAST, StepType.MID),

  )
• PPOTransition.truncated is never True. Let's look at the code to understand why:
  

  Stoix/stoix/systems/ppo/anakin/ff_ppo.py

  Line 95 in 11dca7d

  truncated = (timestep.last() & (timestep.discount != 0.0)).reshape(-1)

  
  To simplify the logic, I am going to replace the variables with descriptive names in pseudocode:

truncated = ((truncation_transition_OR_termination_transition) & (timestep.discount != 0.0)).reshape(-1) 

truncated = ((truncation_transition_OR_termination_transition) & (timestep.discount != False)).reshape(-1) 

truncated = ((truncation_transition_OR_termination_transition) & (timestep.discount == True)).reshape(-1) 

truncated = ((truncation_transition_OR_termination_transition) & (timestep.discount)).reshape(-1) 

truncated = ((truncation_transition_OR_termination_transition) & (not truncation_transition_OR_termination_transition)).reshape(-1) 

truncated = False.reshape(-1) 

I assume that this is not intended behaviour in Stoix. The issue is that the GymnaxWrapper assigns step_type and discount using the same underlying logic, and I assume that the reason you did this is because gymnax does not separately expose truncation and termination in its API (https://github.com/RobertTLange/gymnax/blob/aef77d5c642ea48b95f34c51d05b8417d9450e15/gymnax/environments/classic_control/acrobot.py#L152).

As far as I can tell, in ff_ppo.py, truncations and dones are only used in the GAE calculation, but since the GAE calculation does not care whether episodes where truncated or terminated, we could remove truncation_flags entirely, and have the termination/truncation information entirely in discount_t. Then, we could discard PPOTransition.truncated and change the intended information of PPOTransition.done from termination-transitions to truncation-or-termination-transitions. For this, we would need to adapt wrappers like NavixWrapper to adhere to the same logic, but it shouldn't be too much work.

One caveat here is that I exclusively worked with ff_ppo.py, so I am not sure whether other algorithms do treat truncation and termination differently. Please let me know whether the do. If not, then what do you think of removing the distinction between truncated and terminated Stoix? I could work on this and open a draft PR for the ff_ppo case.

[EdanToledo]

Hi, thanks for this detailed breakdown. You're correct in identifying that the current inconsistency in the Stoix API regarding truncation and termination flags can lead to confusion. I've been aiming to be more explicit about differentiating termination from truncation flags precisely because their distinction significantly impacts GAE calculations.

To clarify, truncation vs. termination does directly affect bootstrapping logic which in turn affects GAE:

• Termination implies the episode reached a definitive end (goal achieved), meaning the value used for bootstrapping at that timestep should be set to zero.
• Truncation indicates the episode ended due to hitting the maximum number of allowed steps, in which case the value at that state should still be utilized for bootstrapping.

Here's a simplified example illustrating why this distinction matters for GAE:

Suppose you have a sequence of values: [1, 2, 3, 4, 5, 6]

• With termination flags [0, 0, 0, 0, 1, 0], the value at position 5 (which is '5') must be zeroed out for bootstrapping purposes, clearly marking that position 6 ('6') is the start of a new episode and not relevant for previous GAE calculations.
• With truncation flags [0, 0, 0, 0, 1, 0], we retain the value '5' for bootstrapping, and position 6 ('6') remains explicitly recognized as a new episode start and therefore excluded from prior episode calculations.

Given your observations:

• It's clear that the current GymnaxWrapper logic conflates termination and truncation because Gymnax itself does not explicitly expose this differentiation.
• Your suggestion to consolidate termination/truncation logic entirely into discount_t and potentially remove the explicit truncated field from PPOTransition could indeed simplify the implementation but for some reason i remember there was an issue with this which is why i introduced this but i could just be misremembering.

I'd love to for us to think more about this and try consolidate and simplify the termination/truncation API but we have to be 100% we are not breaking other calculations.