Skip to content

Provide terminal obs data, add truncation support - #65

Open
stefanfausser wants to merge 15 commits into
edbeeching:mainfrom
stefanfausser:fixes/obs
Open

Provide terminal obs data, add truncation support#65
stefanfausser wants to merge 15 commits into
edbeeching:mainfrom
stefanfausser:fixes/obs

Conversation

@stefanfausser

@stefanfausser stefanfausser commented May 2, 2026

Copy link
Copy Markdown
Contributor

Update: What this PR does

When send_terminal_obs_info of the sync node is set to true (default is false), then the terminal observations are set in info[idx]["terminal_observation"][key] either on reaching a terminal state or on truncation.

Truncation happens automatically when an AI-controller reaches a timeout (see _physics_process), but this can be overwritten by a user when necessary. Further, truncation can be triggered manually by setting controller.truncated = true, followed by controller.needs_reset = true. On truncation, info[idx]["TimeLimit.truncated"] = True is set.

How the PR started (not relevant anymore)

Only read the following when interested in the development of this PR. Otherwise: Ignore the rest of this comment.

Original title of this PR was: "Fixes: Have done + observation reported in the 'step'-reply to the server aligned, get new actions after a game-reset"

What is the issue with the observations:

  • When a terminal state is observed then in the environments needs_reset and done are both set to true (e.g. see RingPong example where the ball enters the outer ring)
  • However, usually in the same physics frame, _physics_process() of the player is called where needs_reset is checked and then the AI-controller and the game is reset (in RingPong, the ball is randomly set elsewhere)
  • Now when _physics_process() of sync is called afterwards, this causes the problem that the observation is retrieved from the resetted-game for the learning (see _training_process()) instead of the observation when done was set to true (so it is one observation ahead)
  • Furthermore, when repeated actions are utilized, then the observations retrieved in _physics_process() of sync are even more ahead (so the problem is getting worse)
  • However, it is important to have done + observations aligned so the RL algorithm, e.g. PPO, knows at which observations (the terminal states) the maximization of the rewards ends.
  • Furthermore, when repeated actions are utilized, then the previous actions are applied to the initial observations after the game was reset. However, because the game situation of the resetted-game is now very different, new actions should be retrieved instead

How this is fixed:

  • The observation is immediately stored and kept when done is set to true (happens per agent/environment)
  • The training is issued for all agents/environments when at least one of the agents has set done to true, irrespective of the repeat action mechanism
  • Last, new actions are now retrieved for all agents/ environments, both for the cases of inference and training, after the game was reset by at least one of the agents. Future repeated actions (when enabled) will then take these new actions

@stefanfausser stefanfausser changed the title Fix: Have done + observation reported in the 'step'-reply to the server aligned Fixes: Have done + observation reported in the 'step'-reply to the server aligned, get new actions after a game-reset May 3, 2026
@stefanfausser

stefanfausser commented May 15, 2026

Copy link
Copy Markdown
Contributor Author

This fix should now be concise and complete. I tested it successfully with RingPong.

Environments suffering from above described issues as well as environments that work around this issue (e.g. by calling call_deferred("reset_game"), where in reset_game the game-reset happens between the physics frames) work with this fix.

@Ivan-267: I would be happy when this fix would be reviewed. In case you need more tests, please tell me.

@Ivan-267

Ivan-267 commented May 16, 2026

Copy link
Copy Markdown
Collaborator

@Ivan-267: I would be happy when this fix would be reviewed. In case you need more tests, please tell me.

Thanks for the PR. I might not have the time to review this one properly soon, research a bit and check for edge cases, but I wanted to at least let you know I saw the PR.

Some APIs might have changed slightly since the base behavior was implemented by @edbeeching - so I think the author's feedback would be valuable on this.

Some related docs:
https://stable-baselines3.readthedocs.io/en/master/guide/vec_envs.html

  • How SB3 handles VecEnvs, I think our primary focus for now

https://farama.org/Vector-Autoreset-Mode

  • How new Gymnasium VecEnv handles this (for future consideration, there are multiple options, one most compatible with SB3's solution above if we want to make it be a Gymnasium VecEnv or work on a compatible wrapper)

On action repeat:
Action repeat does complicate this as you correctly note. After a single AIController's episode is auto-restarted by the env, the next "RL step" can take place after n physics steps elapsed in the new episode.

In multiple example envs this was not very harmful, it just slightly randomized the starting state each episode (unless the action repeat is large enough to cause issues).

In some envs where this is harmful (e.g. repeating the old action leads the agent to move to a wall), one option is to zero the repeated action for that agent, and wait for the next set_action before the agent moves.

In envs where an action right after the env resets is needed, it might be possible to pause the env associated with that AIController until the first action is received (set_action called). This still allows the other AIController/agents to "repeat" their actions for action_repeat steps for consistency.

The easiest solution/implementation might depend on the environment. All of these options are compromises with some downsides, but they might be able to resolve some of the cases.

Ideally we would get a new action for each env from the server right after the restart, but as this is a vectorized env, when stepping with SB3 we get an action for all envs at once, and manual/async stepping is not directly supported, so it's worth considering workarounds when action repeat is used.

@stefanfausser

stefanfausser commented May 16, 2026

Copy link
Copy Markdown
Contributor Author

Thank you @Ivan-267, for pointing out that godot-rl-agents is using the VecEnvAPI and not the (old) Gym API. I should have known but was not aware of this.

However, I am not sure that the auto-reset mechanism does fix all of what I have seen.

Let me elaborate more on what I have seen and debugged. In the RingPong example, as soon as the ball (a RigidBody3D) is entering the ring (an Area3D), following function is called (see game.gd):

func _on_ring_body_entered(body):
	$Player.game_over()

Now in game_over(), both done and needs_reset of the ai_controller are set to true (see player.gd):

func game_over():
	ai_controller.done = true # Terminal state
	ai_controller.needs_reset = true # Issue a reset as soon as possible
	print("game_over() frames: ", Engine.get_physics_frames())
	print("game_over() at n_steps: ", ai_controller.n_steps)

The last two lines were added by me for debug purposes.

Afterwards _physics_process() of the player is called:

func _physics_process(delta):
	print("player _physics_process() frames: ", Engine.get_physics_frames())
	print("ai_controller.n_steps: ", ai_controller.n_steps)
	if ai_controller.needs_reset:
		ai_controller.reset() # set steps to zero (start with a new episode)
		ball.reset()
		return

The first two prints were added by me as well for debug purposes. However, those prints showed that game_over() and then _physics_process() of the player is called in the same physics frame.

Why this is problematic: That means that any later call of get_obs() / ai-controller will not return the obs of the terminal state but the obs of the resetted-game.

Then in the same physics frame follows the _physics_process() of the sync node. There, _training_process() is called when training is configured (excerpt of sync.gd):

		var obs = _get_obs_from_agents(agents_training)
		var info = _get_info_from_agents(agents_training)
		# some irrelevant code in-between, removed for demonstration purposes
		if need_to_send_obs:
			need_to_send_obs = false
			var reward = _get_reward_from_agents()
			var done = _get_done_from_agents()

			var reply = {"type": "step", "obs": obs, "reward": reward, "done": done, "info": info}
			_send_dict_as_json_message(reply)

So within this code snippet I would have expected 'obs'-terminal-state representations of the envs where 'done' is true. But because _get_obs_from_agents() is calling get_obs() of the respective ai-controllers, this is returning 'obs'-resetted-game representations instead.

This information is then sent to the Python server and end up in step_recv() of a GodotEnv object (see godot_env.py):

    def step_recv(self):
        """
        Receive the step response from the Godot environment.

        Returns:
            tuple: Tuple containing observation, reward, done flag, termination flag, and info.
        """
        response = self._get_json_dict()
        response["obs"] = self._process_obs(response["obs"])

        # Kept for backward compatibility if the plugin doesn't send info.
        default_info = [{}] * len(response["done"])

        return (
            response["obs"],
            response["reward"],
            np.array(response["done"]).tolist(),
            np.array(response["done"]).tolist(),  # TODO update API to term, trunc
            response.get("info", default_info),
        )

So in the end, response["obs"] and response["done"] does also not fit as described above. And this is the case even without using action repeats. With action repeats the problem is getting worse.

Before moving on with my proposal (e.g. closing or changing it) I would like you, @Ivan-267 or @edbeeching, to confirm this. Maybe I am overlooking something?

If this is truly a problem then it is not limited to RingPong Godot example only as other examples behave similar (the game reset is issued in the same physics frame after done has been set to true)

@Ivan-267

Ivan-267 commented May 16, 2026

Copy link
Copy Markdown
Collaborator

Yes, let's keep the PR until more discussion and/or a closer review by someone. I also have to learn a bit more about this, as I mentioned I won't be able to do a full review soon:

Env specification/terminal obs:
I'm also not claiming an exact specification for the original implementation, I called it a vectorized env because it supports multiple env instances, while I think the old Gym API was for single envs with manual restarting.

In context of SB3 at least with the SB3 wrapper, it is a custom VecEnv, just not fully implemented (i.e. we don't have truncation yet). There is also support for other frameworks, and we should ideally double check how they handle this as well.

So, for SB3, the documented behavior (I previously linked the docs):

When using vectorized environments, the environments are automatically reset at the end of each episode. Thus, the observation returned for the i-th environment when done[i] is true will in fact be the first observation of the next episode, not the last observation of the episode that has just terminated. You can access the “real” final observation of the terminated episode—that is, the one that accompanied the done event provided by the underlying environment—using the terminal_observation keys in the info dicts returned by the VecEnv.

In Gymnasium, this is not the default mode, there are two other modes in the linked doc, this one just seems the most similar one:

Same-Step Mode

If a sub-environments terminated, in the same step call, it is reset, beware that some vector wrappers do not support this mode and the step’s observation can be the reset’s observation with the terminated observation being stored in info["final_obs"]. This makes it is a simplistic approach for training algorithms if value errors with truncation are skipped. See this, for details.

If I understand it correctly, strictly considering the obs (not info dictionary), the terminal obs is not always returned in this mode, especially under the SB3's definition.

Also if I understood the doc right (it needs a bit more research from my side), it becomes a larger issue with truncation. When we work on adding truncation (I considered implementing this a while ago and we can at some point, so it's worth reading up on it), the terminal obs must be included in the info dictionary, and then it's up to the training code to extract the terminal obs from the info dictionary for bootstrapping.

Here's what SB3's on policy algorithm implementation does when truncated and when terminal obs is included in the info: https://github.com/DLR-RM/stable-baselines3/blob/10dda8678759b1b42b3afd9baf442addc19aabcd/stable_baselines3/common/on_policy_algorithm.py#L236-L245

Action repeat
I think all of the approaches we've mentioned to solving issues caused by this have some compromise, and some envs might be more sensitive to one approach than another.

E.g., per env:

  • Is it an issue if the agent sometimes starts in a different state caused by a previous action repeating
  • Is a new action needed right after the env resets
  • Is it complicated to pause the specific env and wait for the next step
  • What happens if the action does not always last for action_repeat steps but the frequency varies.

Different envs might have different issues and some solutions might work better than others.

@stefanfausser

Copy link
Copy Markdown
Contributor Author

So, for SB3, the documented behavior (I previously linked the docs):

When using vectorized environments, the environments are automatically reset at the end of each episode. Thus, the observation returned for the i-th environment when done[i] is true will in fact be the first observation of the next episode, not the last observation of the episode that has just terminated. You can access the “real” final observation of the terminated episode—that is, the one that accompanied the done event provided by the underlying environment—using the terminal_observation keys in the info dicts returned by the VecEnv.

Thank you again, @Ivan-267, for taking your time to have a look at this PR. Actually after you showed to me the documentation I have read above passage. Furthermore, I have seen a note below:

"When defining a custom VecEnv (for instance, using gym3 ProcgenEnv), you should provide terminal_observation keys in the info dicts returned by the VecEnv (cf. note above)."

However, StableBaselinesGodotEnv(VecEnv), see stable_baselines_wrapper.py / SB3 wrapper, does not provide the terminal_observation keys in info dicts. The reason is that the godot-rl-plugin on the Godot side does not capture the correct terminal-state-obs (and it cannot because it is immediately overwritten by the resetted-game in the same physics frame) and does not set this in the info which is then sent to the SB3 wrapper.

So I now hope that I have described what I feel is an issue (maybe it isn't) clearly enough.

@stefanfausser

stefanfausser commented May 17, 2026

Copy link
Copy Markdown
Contributor Author

I am not a friend of LLMs (except for some boring tasks and sometimes to challenge myself) but I "gambled" around with Claude. Given that StableBaselinesGodotEnv is derived from VecEnv (so it is a custom env) and my question "must I set info["terminal_observation"]?"

Claude said:

"Great question! info["terminal_observation"] must be set when an episode ends (done=True), and it should contain the actual observation at the moment the episode terminated.

Why It's Important

When an episode ends, VecEnv automatically resets the environment and returns a fresh initial observation in the next step. However, the learning algorithm needs access to the true final observation to properly compute value estimates and advantages (especially important for algorithms like PPO and A2C that use bootstrapping).

Without terminal_observation, the algorithm might incorrectly use the reset observation for value calculations, which would be wrong."

I "interacted" a bit more with Claude and asked what would be the issue with PPO and got:

"The Actual Problem

When done=True:

  • Your environment's step() returns the reset observation (the initial observation of the new episode), not the terminal observation
  • So PPO receives: observation_t, action_t, reward_t, reset_observation, done=True
  • PPO has lost the actual terminal observation — it only has the fresh initial observation from the reset

Why info["terminal_observation"] Matters

PPO needs the actual terminal observation for proper value bootstrapping. When computing the advantage:

advantage = reward + gamma * V(next_observation) * (1 - done)

When done=True, this becomes:

advantage = reward + 0  (because done masks out the bootstrap)

But PPO's value network still needs to process something. It uses info["terminal_observation"] to get the true final state and compute the correct value estimate.
So the Correct Answer

Yes, a custom VecEnv absolutely must set info["terminal_observation"] when done=True, because:

  • The observation returned in that step is the reset observation, not the terminal one
  • PPO needs the actual terminal observation to compute proper value estimates
  • Without it, PPO would use the value of the reset state (wrong episode) for bootstrapping"

At least this would match what I was on my mind. But I (and Claude) still could be wrong...

@stefanfausser

Copy link
Copy Markdown
Contributor Author

Given all infos I now have, I posted another commit. The overall changes to main repo are now minuscule but I think important. As I would like to keep the good discussions in this PR, I am not editing / changing my original ideas in my first posts. Instead I describe here what the fix does:

  • The observation is immediately stored and kept when done is set to true (happens per agent/environment)
  • The training is issued for all agents/environments when at least one of the agents has set done to true, irrespective of the repeat action mechanism
  • In the training, info["terminal_observation"] is set to the observation of the terminal state so godot-rl-agent examples, e.g. SB3, can find the actual observation when the episode ended (done was set to true)

Now I am patient and will keep this PR without any additions until this has been confirmed / reviewed thoroughly.

@Ivan-267

Ivan-267 commented May 18, 2026

Copy link
Copy Markdown
Collaborator

Thanks for the update.

For now I'll also leave my current understanding, not guaranteed to be correct, as I haven't researched it fully yet.

Yes, based on the SB3 docs that is needed to implement the VecEnv completely.

For PPO training specifically:
I think it's used during bootstrapping on truncation in the training code.

Termination:
We currently treat done as terminated, and truncation is not implemented yet (but we've previously discussed adding it, it's worth adding).

In case of an episode termination, there is no action to take from the terminal state,
and there are no future rewards to receive for that episode. So the predicted future returns are also 0 (as claude mentioned).

Also, let's look at the Godot env side, if it sends the next episode initial obs with done = true.

Let's take a very simple environment example, only 1 step. There's a starting state 0 (middle position),
goal can be left or right (-1 or 1 goal state), and there are two actions (move left or move right), goal pos is fixed to simplify (always 1)

Step by step:

  1. The state is 0, goal is 1 (right), env sends: obs = 0
  2. The server replies: action = 1
  3. The env takes the action, sets reward = 1, done = true (only one step is allowed in this env), resets the env, state 0 again
  4. The env sends to server: obs = 0 (next episode started), reward = 1 (result of previous action), done = true (result of previous action correction: episode terminated after the previous action)

Now, because the server has obs = 0, it can send the next action to take, and the cycle can repeat.

We didn't send the terminal obs there at all.
But without any other changes, if we sent the terminal obs instead of the next episode obs,
the server would send an action based on the end position (-1 or 1), however we only want an action based on the starting state (0).
Also when we do inference later, we'll only act on the initial state, not the terminal state, so at least the actor network doesn't need to see the terminal state (in case of the standard networks we used with SB3 PPO for examples).

The one step example doesn't really lend itself to truncation (will mention it more below), but it focuses on why in this loop we send the next episode obs together with done.

As mentioned with the Gymnasium API, there might be different ways to handle the terminal obs and reset,
but for now we're focusing purely on the SB3 VecEnv and the GDRL implementation.

Truncation:
This can be used when we want to terminate an episode early for training (e.g. reset on timeout),
but for the end task/inference, we don't want to terminate the episode on timeout.

If we always want to terminate it after the set timesteps (the task has a time limit),
we can use termination instead, and we can add an observation such as n_steps / float(reset_after).

However, in the previous case, truncation is used. At that point, future returns are predicted from the terminal state, as it's terminal only during training, but not for the final task/inference.

Here, since we haven't provided the terminal state, it cannot bootstrap from it in case of truncation.

In conclusion:
I think this shows why PPO training can proceed normally in the example environments using the current implementation without providing the terminal obs.
But, for a full vecenv implementation and supporting truncation (and perhaps some env wrappers too), we should add it.

It's been interesting to take a bit of a deeper look at this, although I don't have the time to check it fully.
Because of that, I also haven't yet commented on the implementation you provided much, but focused on the docs
and agreeing on the specs first (I do see you did some updates, but haven't checked in depth yet).

I think it should just be a matter of adding the terminal obs to the info per env (your latest changes address this) and truncation support (also in info), at least for SB3 support. We also have some docs written related to adding info for success rate reporting here, it might need to be updated so the info is not overwritten.

The action repeat is a separate issue, as discussed previously.

P.S.
As for the implementation, I'm also considering whether we should make sending this info optional, as some frameworks might not use the info as sent (I haven't checked others yet), and a very specific edge case might involve episodes with only a few steps (or single step episodes), if processing the obs is expensive, and if not using truncation or not using the terminal obs in the algorithm implementation. In the single step episode case, we are doubling the sent obs data. A switch would be relatively simple to implement, but we can still consider whether it's needed until the final implementation.

@stefanfausser

stefanfausser commented May 18, 2026

Copy link
Copy Markdown
Contributor Author

In conclusion: I think this shows why PPO training can proceed normally in the example environments using the current implementation without providing the terminal obs. But, for a full vecenv implementation and supporting truncation (and perhaps some env wrappers too), we should add it.

It's been interesting to take a bit of a deeper look at this, although I don't have the time to check it fully. Because of that, I also haven't yet commented on the implementation you provided much, but focused on the docs and agreeing on the specs first (I do see you did some updates, but haven't checked in depth yet).

I think it should just be a matter of adding the terminal obs to the info per env (your latest changes address this) and truncation support (also in info), at least for SB3 support. We also have some docs written related to adding info for success rate reporting here, it might need to be updated so the info is not overwritten.

Although I started differently (I originally thought that the obs representation needs to be that of a terminal state and not of the next state / resetted-game) I converged eventually to the same conclusion.

In addition to this, I found how SB3 PPO is using this info. Excerpt of collect_rollouts (see on_policy_algorithm.py):

 # Handle timeout by bootstrapping with value function
            # see GitHub issue #633
            for idx, done in enumerate(dones):
                if (
                    done
                    and infos[idx].get("terminal_observation") is not None
                    and infos[idx].get("TimeLimit.truncated", False)
                ):
                    terminal_obs = self.policy.obs_to_tensor(infos[idx]["terminal_observation"])[0]
                    with th.no_grad():
                        terminal_value = self.policy.predict_values(terminal_obs)[0]  # type: ignore[arg-type]
                    rewards[idx] += self.gamma * terminal_value

And PPO, which is an On-Policy algorithm is derived from this. So the original idea of this PR is invalid (spoken openly it is a mess except for the good discussions we had). However, this PR could still be useful by adding support of SB3-style truncation instead. The changes to this would be minimal.

How I would implement it:

  • Add "SB3 truncation support" as an export variable in sync.gd that is disabled (false) by default
  • With "SB3 truncation support" enabled: When the 2D or 3D AI controller reaches a timeout, then info["terminal_observation"] = obs of the terminal state, info["TimeLimit.truncated"] = True and done = True. Exactly this combination would allow the On-Policy / PPO code to set the bootstrap estimation correctly (see code snippet above)
  • I would revert the forced learning on reaching a terminal state as this is not necessarily needed and would only affect the behaviour when it comes to repeated actions / frameskipping. I still think that this might be an issue but this could be tackled env-wise or by another PR
  • Optionally, I would set info["terminal_observation"] not only for truncated states but also on terminal states so other/custom wrappers could access it. This would also be configurable via an exported variable in sync.gd with is false by default

I will only move on if this is fine for you (@Ivan-267 and @edbeeching ).

@Ivan-267

Ivan-267 commented May 18, 2026

Copy link
Copy Markdown
Collaborator

Yes, I also checked the same:

Here's what SB3's on policy algorithm implementation does when truncated and when terminal obs is included in the info: https://github.com/DLR-RM/stable-baselines3/blob/10dda8678759b1b42b3afd9baf442addc19aabcd/stable_baselines3/common/on_policy_algorithm.py#L236-L245

We reached the same conclusions, the code confirms it's used with truncation. The base idea from your PR (to provide the terminal obs data) is useful, we just iterated on the specs so that the final implementation should be compatible.

The only question left is the implementation, we can iterate on that. I grouped some of your points together or out of order for commenting. All of my comments below are open for feedback, consideration, and iteration. But the main points are in agreement, it's mostly the details.

How I would implement it:
Add "SB3 truncation support" as an export variable in sync.gd that is disabled by default
Optionally, I would set info["terminal_observation"] not only for truncated states but also on terminal states so other/custom wrappers could access it. This would also be configurable via an exported variable in sync.gd with is false by default

Yes, a switch like this is useful. A small thing to consider is the name. Based on your other point about always sending the terminal obs, I'm considering something like send_terminal_observation, or more clear: send_terminal_obs_info instead as one possible name (with a tooltip explaining what it's used for, e.g. SB3 and truncation, other wrappers might also use it so we shouldn't bind the name to SB3).

Then, if this is enabled, truncation can be used. If not, an assert could be raised (or a warning printed) if truncated = true is set that informs the user to enable send_terminal_obs_info. This one switch should be enough to solve the cases I previously mentioned.

I agree about setting it to false as default, both for compatibility (it keeps the previous behavior as default), and as not all frameworks/wrappers might use it.

I also agree on sending the terminal obs info on done regardless of whether the episode is truncated (as long as send_terminal_obs_info is true, as your second point suggests and it is more in line with the SB3 docs).

With "SB3 truncation support" enabled: When the 2D or 3D AI controller reaches a timeout, then info["terminal_observation"] is set with the terminal-state-obs, info["TimeLimit.truncated"] = True and done = True. Exactly this combination would allow the On-Policy / PPO code to update the bootstrap estimation correctly (see code snippet above)

Yes, but in addition to using it the built-in timeout code in AIController, I think it we can have truncated as another flag similar to done, to be used in combination when needed.

Then, when making an env, if we want to choose when to truncate, the usage is for example:

# Some game code here
if game_over:
    ai_controller.done = true
    ai_controller.reset()    
    reset_game()

if game_over_on_timeout:
    ai_controller.done = true
    ai_controller.truncated = true
    ai_controller.reset()    
    reset_game()

(just an example with some duplicated code to highlight the difference, it can be used in the extended ai_controller script, or with needs_reset like in the tutorial, etc.)

Both the done and truncated flags get reset to false after being read by the sync node.

Then, as you proposed, we can also add done = true and truncated = true to the AIController's built-in timeout reset code as it's the natural place to use truncation, but it can also be be overridden (by overriding the _physics_processing method in extended ai_controller) and used manually based on the specific env.

E.g. in the base AIController script, you could invoke reading/storing the terminal obs when done and allow_truncation is set, regardless of whether truncation is set or not. Then it can be accessed on the Python side.

I would revert the forced learning on reaching a terminal state as this is not necessarily needed and would only affect the behaviour when it comes to repeated actions / frameskipping. I still think that this might be an issue but this could be tackled env-wise or by another PR

Agreed, keep the previous action repeat behavior, and the action repeat solution can be discussed separately.

All of this is also open to more feedback by @edbeeching and/or another contributor. I might be able to give more feedback later, but as mentioned before, not in the next few days.

@Ivan-267

Copy link
Copy Markdown
Collaborator

As another quick iteration on my previous proposed changes, truncated can automatically set done to true, so the user only has to set truncated to true. Optional, but if it doesn't cause any issues, it will be slightly easier to use.

@stefanfausser

Copy link
Copy Markdown
Contributor Author

Thank you, @Ivan-267 , for your kind support to this PR. This is highly appreciated.

I have implemented almost all of above and found out during debugging that info["terminal_observation"]["obs"] must be set to the observation of the terminal state (and not just info["terminal_observation"]) so PPO (or other On-Policy algorithms in SB3 lib) can use it. When I have tested the implementation with varying number of timeout-steps successfully, I will commit my changes to this branch.

…minal_obs_info export variable to configure sending terminal states via info, added truncation support
@stefanfausser

stefanfausser commented May 19, 2026

Copy link
Copy Markdown
Contributor Author

I debugged the implementation from Godot up to the collect_rollouts (see on_policy_algorithm.py) and committed my changes to this branch.

Now, when send_terminal_obs_info of the sync node is set to true (default is false), then the terminal observations are set in info["terminal_observation"][key] either on reaching a terminal state or on truncation.

Truncation happens automatically when an AI-controller reaches a timeout (see _physics_process), but this could be overwritten by a user when necessary. Further, truncation can be triggered manually by setting controller.truncated = true, followed by controller.needs_reset = true.

Last, the forced learning on reaching a terminal state has been removed.

I did not implement to check if truncation is triggered manually when send_terminal_obs_info is false. The reason is that the sync node knows about the AI-controllers but an AI-controller currently does not know any settings / states of the sync node.

Test results with RingPong

The results are... interesting. The ep_len_mean and ep_rew_mean values are worse with truncation than without truncation. However, I think that this is to be expected as the end of an episode is now when the ball leaves the ring (terminal state) and in addition to this when an agent has done more than 1000 steps.

So the episode lengths are now limited with truncation enabled. Actually, with a maximum of 1000 steps and an action repeat of 8, the ep_len_mean could be maximum 1000 / 8 = 125 (I have seen this value during a test). With truncation disabled, the episode lengths are theoretically not limited, as the ball could never leave the ring.

However, from a visual inspection, the performance of RingPong with truncation is better to me as the agents actually learn quicker to rotate the paddle where it should be. Maybe because the agents have seen more often different starting positions of the ball.

So in case of RingPing comparing ep_len_mean / ep_rew_mean with or without truncation is not fair. Maybe other metrics, e.g. average episode reward per 100 steps are more fair.

Btw. also different values for action repeat have an impact on ep_len_mean (higher values of action repeat means lower values for ep_len_mean).

In SB3, custom evaluations via a callback could be added. While I am still thinking about the evaluation part, I feel that this PR that is adding terminal observation support to the Godot side is now ready for review.

@stefanfausser stefanfausser changed the title Fixes: Have done + observation reported in the 'step'-reply to the server aligned, get new actions after a game-reset Provide terminal obs data, add truncation support May 19, 2026
@stefanfausser

Copy link
Copy Markdown
Contributor Author

I changed the title to a more meaningful title. However, for archiving reasons, I added the original title to my first comment to this PR. So it can be seen how this developed.

@Ivan-267

Ivan-267 commented May 20, 2026

Copy link
Copy Markdown
Collaborator

Thanks, that sounds good. Difference in ep length and reward is definitely expected if the env didn't terminate on timeout before. One interesting test could be to try with done = true instead (even though not intended for timeout without adding an obs with normalized timestep progress as mentioned before), then ep lenght should match the truncated one, and reward could be compared directly. Still, it's a relatively quick to train env and I'm not sure how large the effect would be.

I still haven't considered a perfect test case to compare the training impact of this vs termination in an env, but if it functionally works (sb3 receives the data of the terminal obs and uses it), then it should match the spec.

I won't be able to make the final checks soon as mentioned on this or other PRs (I'm aware the raycast debug one should be checked too, it's a useful thing to have), but it's good to know this one is ready for review now.

@stefanfausser

stefanfausser commented May 20, 2026

Copy link
Copy Markdown
Contributor Author

I thought a bit about a fair comparison and came up with ep_ign_trunc_rew_mean. While ep_rew_mean is considering terminated or truncated as an episode end, ep_ign_trunc_rew_mean is considering terminated only as an episode end.

Results for RingPong (standard values except for action_repeat = 4) with send_terminal_obs_info = false and done = false on truncation:

Bildschirmfoto vom 2026-05-20 22-18-25

ep_ign_trunc_rew_mean is visible on the left figure and ep_rew_mean on the right. As with send_terminal_obs_info = false, done is not set to true on a truncation (this is the current behaviour of godot-rl-agents-plugin in the main repo), both values are exactly the same. This is also shows that my custom metric ep_ign_trunc_rew_mean works as expected.

Results for RingPong (standard values except for action_repeat = 4) with send_terminal_obs_info = true and done = true on truncation:

Bildschirmfoto vom 2026-05-20 22-03-16

ep_ign_trunc_rew_mean is visible on the left figure and ep_rew_mean on the right. As with send_terminal_obs_info = true, done is set to true on a truncation (this would be the new behaviour with this PR), the ep_ign_trunc_rew_mean becomes larger than ep_rew_mean. Interestingly, ep_ign_trunc_rew_mean becomes even larger as above (but of course this should be repeated several times to confirm this). This also confirms my visual inspection: The agent is now learning faster to move the paddle to where the ball is. From a theoretical standpoint, the results with proper truncation should be better with proper truncation as the bootstrap estimations for the next state are now correct.

Btw. I had to add a custom checkpoint and the calculations for ep_ign_trunc_rew into stable_baselines3_example.py. If this PR is eventually accepted, I will open up another PR on interest, where I post the changes to stable_baselines3_example.py so this new metric becomes available in main repo (although ep_ign_trunc_rew is a long name for a metric. Any suggestions for a better name?)

@Ivan-267

Ivan-267 commented May 21, 2026

Copy link
Copy Markdown
Collaborator

Just a quick feedback on this, thanks for adding the results. I agree run to run variance can have an effect.

To clarify something about the ep_ign_trunc_rew_mean (it's possible I overlooked something from reading the description, feel free to let me know), looking at the second case results where the reward is much larger than ep_rew_mean (more on why I didn't fully understand that below), is this a measurement during training (with episodes truncating after n_steps), or some evaluation with truncation disabled (so episodes never reset)?

I assume during training the episodes are always truncated after n steps in this setup. If the episodes are always truncated after n steps, and if the ep_ign_trunc_rew_mean stat is taken during training, I assume the maximum episodic reward should not be much larger when terminated (including only game_over episodes in this env) vs truncated if the ep reward is roughly increasing with episode length. That is, the agent never experiences an episode long enough during training to collect such rewards. We can try to predict (as the critic does) the future returns that might have happened if the episode lasted longer, but we cannot measure them while truncating and resetting the episode (unless we also run episodes without ep resets on truncation in parallel only in evaluation mode, but that complicates the implementation).

I would recommend this check also (optionally, I can also try this on an env when I have more time to test this), as previously described, with an addition to make it 3 test cases:

  1. Ending the episode on n_steps > reset_after by setting truncated = true (sb3 receives done = true, truncated and terminal obs info)
  2. Ending the episode on n_steps > reset_after by setting done = true (sb3 only receives done = true and terminal obs info)
  3. Same as 2 but with adding obs.append(n_steps / float(reset_after))

Those would be 3 possible scenarios to benchmark against each other, and the ep length/reward progress during training can be directly compared.

When making envs, I often use done = true to signal an episode end even on truncation, usually with the added obs, so it's a practically relevant comparison.

@stefanfausser

stefanfausser commented May 21, 2026

Copy link
Copy Markdown
Contributor Author

I would recommend this check also (optionally, I can also try this on an env when I have more time to test this), as previously described, with an addition to make it 3 test cases:

1. Ending the episode on `n_steps > reset_after` by setting `truncated = true` (sb3 receives `done = true`, truncated and terminal obs info)

2. Ending the episode on `n_steps > reset_after` by setting `done = true` (sb3 only receives `done = true` and terminal obs info)

3. Same as 2 but with adding `obs.append(n_steps / float(reset_after))`

Those would be 3 possible scenarios to benchmark against each other, and the ep length/reward progress during training can be directly compared.

I'll look into this. Currently, the second figure in my post above matches your first scenario. The first figure in my post above is under the non-working default-truncation condition where just n_steps is reset on n_steps > reset_after but neither done nor truncated is set to true (but this is what current main repo does without any changes). So I just brought up the first figure so you can see that the calculation of ep_ign_trunc_rew_mean is correct because in this scenario it must be identical to ep_rew_mean. And a proper truncation support improves the performance compared to this. I hope this makes sense.

I`ll try the other two scenarios you mentioned above later.

PS: Both figures are under training

@Ivan-267

Copy link
Copy Markdown
Collaborator

I would recommend this check also (optionally, I can also try this on an env when I have more time to test this), as previously described, with an addition to make it 3 test cases:

1. Ending the episode on `n_steps > reset_after` by setting `truncated = true` (sb3 receives `done = true`, truncated and terminal obs info)

2. Ending the episode on `n_steps > reset_after` by setting `done = true` (sb3 only receives `done = true` and terminal obs info)

3. Same as 2 but with adding `obs.append(n_steps / float(reset_after))`

Those would be 3 possible scenarios to benchmark against each other, and the ep length/reward progress during training can be directly compared.

I'll look into this. Currently, the second figure in my post above matches your first scenario. The first figure in my post above is under the non-working default-truncation condition where just n_steps is reset on n_steps > reset_after but neither done nor truncated is set to true (but this is what current main repo does without any changes). So I just brought up the first figure so you can see that the calculation of ep_ign_trunc_rew_mean is correct because in this scenario it must be identical to ep_rew_mean. And a proper truncation support improves the performance compared to this. I hope this makes sense.

I`ll try the other two scenarios you mentioned above later.

PS: Both figures are under training

Thanks, I understood why in the first case they were equal. My comments mostly referred to the second figure:

Results for RingPong (standard values except for action_repeat = 4) with send_terminal_obs_info = true and done = true on truncation:
Bildschirmfoto vom 2026-05-20 22-03-16

Here, ep_ign_trunc_rew is larger than ep_rew_mean in the same image (from the same training run I suppose), that part is what my comments were about. I also assumed based on description that both are intended to show the average reward of an episode during training, and that the first one covers only episodes that terminated with done, and the second one episodes that truncated. I also made an assumption that longer episode more or less means a larger reward, and under those assumptions, there isn't a longer episode available to the agent to produce the larger reward during training itself. There are a few possible explanations, I just didn't fully understand the implementation so I wanted to check.

@stefanfausser

stefanfausser commented May 21, 2026

Copy link
Copy Markdown
Contributor Author

Here, ep_ign_trunc_rew is larger than ep_rew_mean in the same image (from the same training run I suppose), that part is what my comments were about. I also assumed based on description that both are intended to show the average reward of an episode during training, and that the first one covers only episodes that terminated with done, and the second one episodes that truncated. I also made an assumption that longer episode more or less means a larger reward, and under those assumptions, there isn't a longer episode available to the agent to produce the larger reward during training itself. There are a few possible explanations, I just didn't fully understand the implementation so I wanted to check.

Yes, all values seen in the same image are from the same training run. And yes, the assumption that longer episodes more or less means larger rewards also holds for the RingPong environment.

And this is also the reason why ep_ign_trunc_rew_mean is larger than ep_rew_mean:

  • ep_rew_mean: Original metric, unchanged by me. It is the average reward over the last 100 collected episodes while an episode ends either on terminated or truncated
  • ep_ign_trunc_rew_mean: Metric added by me (via a custom checkpoint in stable_baselines3_example.py). It is the average reward over the last 100 collected episodes while an episode ends on terminated only. So within this checkpoint in stable_baselines3_example.py, the episode is considered to be terminated when done == True and info["TimeLimit.truncated"] == False (or the key "TimeLimit.truncated" is missing) only. So the episodes considered here are longer because truncated is basically ignored ("ign").

In addition to this, I now added ep_ign_trunc_len_mean(episode ends on terminated only) to my local SB3 code and ran the same scenario again (in your list it is the first, i.e. 1. Ending the episode on n_steps > reset_after by setting truncated = true (sb3 receives done = true, truncated and terminal obs info)):

Bildschirmfoto vom 2026-05-21 15-01-41

So in above image, you can see that ep_ign_trunc_rew_mean is larger than ep_rew_mean because ep_ign_trunc_len_mean is larger than ep_len_mean.

I know I shouldn't have brought up new metrics to evaluate the impact / performance of this PR and I am sorry for any confusion caught.

@stefanfausser

Copy link
Copy Markdown
Contributor Author

3. Same as 2 but with adding obs.append(n_steps / float(reset_after))

Bildschirmfoto vom 2026-05-21 17-45-07

I added the obs to the AI-controller (controller.gd) of RingPong this way:

		var obs = [ball_pos.x, ball_pos.z, ball_vel.x/10.0, ball_vel.z/10.0, n_steps / float(reset_after)]

So instead of 4 input neurons, the neural network has 5 with scenario 3.

The results of scenario 3 (see above) are similar to the results of scenario 1. Just maybe, the agent is getting a bit "depressed" with too many steps.

But overall ep_rew_mean is approaching about 60 in all three scenarios.

@stefanfausser

stefanfausser commented May 21, 2026

Copy link
Copy Markdown
Contributor Author

I'll take a close look when I get more time (it could be a few weeks or more).

Thx, appreciated. Now that the results of the test runs are here and they at least do not show anything out of order, I think that the PR submission should be ready for review. So this PR certainly can wait a few weeks :)

@stefanfausser

Copy link
Copy Markdown
Contributor Author

It seems I accidentally deleted the results for scenario 2, so I add just the screenshot again.

2. Ending the episode on n_steps > reset_after by setting done = true (sb3 only receives done = true and terminal obs info)

Bildschirmfoto vom 2026-05-21 16-06-13

@stefanfausser

stefanfausser commented May 21, 2026

Copy link
Copy Markdown
Contributor Author

As I didn't have the time to take a close look, the below could be wrong: If the code is including rewards and counting steps between episode terminations and ignoring truncation, then one "episode" can actually contain more than one truncated episode and in-game reset.

It would explain why the reported stat is larger than the SB3 episodic rew/mean, but it's not a true measure of an episodic reward. The name did say it ignores truncation rather than records only episodes that terminated with done (which should not show a significantly larger reward than truncated episodes in this env), but I don't think we should consider those episodes, because it potentially contains rewards from multiple episodes. Of course, only if the above is true.

You are correct. But ep_rew_mean does the same for scenario 0.

(Scenario 0 is what is currently happening on godot-rl-agents-plugin main: Ending the episode on n_steps > reset_after without setting done, so done = false. For this case, a neural network learns that a very late state is connected with an initial state but it isn't. And this scenario 0 is reflected in my very first figure now visible way above.)

So it was my intention to implement ep_ign_trunc_rew_mean exactly this way to be able to compare any other scenarios, e.g. scenario 2, with scenario 0. I absolutely agree with you, that this is not a true measure of an episodic reward. And now I understand how I could have implemented it instead.

However, the results for scenario 1 to 3 also include the standard ep_rew_mean (see above). And I did not see any (large) differences in ep_rew_mean so I guess the test still somehow succeeded.

@Ivan-267

Ivan-267 commented May 21, 2026

Copy link
Copy Markdown
Collaborator

Thanks for the full explanation, you're correct it is useful to compare to the base implementation. After this PR, it likely wouldn't be needed (except for such tests, and at least in the current form, we can of course still consider it and alternatives) as envs will either truncate or terminate. I also understand the usefulness of measuring real non-truncated performance, but that's a different matter, and can lead to unbounded length evaluations on infinite horizon episodes without truncation. It would also be possible to measure episodes that terminated only and reset the accumulated rewards on either truncated or terminated with a small change, those however cannot outperform the truncated episodes significantly in cases such as this env and wouldn't show the maximum potential performance either (that one is tricky to measure with truncation).

I've also heard of average reward formulations as alternative for episodic possibly used for infinite horizon tasks (I mean for training algorithms, not just stat reporting), but haven't looked into that in any detail and it's not related to this PR. Truncation can also be useful for infinite horizon tasks (some envs might never terminate episodes).

It sounds good now. Thanks for those tests and I'll take a further look at the implementation when possible.

@stefanfausser

stefanfausser commented May 21, 2026

Copy link
Copy Markdown
Contributor Author

I've also heard of average reward formulations as alternative for episodic possibly used for infinite horizon tasks (I mean for training algorithms, not just stat reporting), but haven't looked into that in any detail and it's not related to this PR. Truncation can also be useful for infinite horizon tasks (some envs might never terminate episodes).

This is certainly worth investigating (outside this PR as you said). However, when it comes to stat reporting, I think that step_rew_mean (reward per step averaged over the last X collected steps) might be an interesting alternative to ep_rew_mean. However, as this PR is already filling a book with its comments (but I am not ungrateful as I got a better understanding of godot-rl-agents + plugin this way) I feel that I better propose this metric in another PR sometime later.

Edit: I implemented step_rew_mean (reward per step averaged over the last 1000 steps) and got more than 0.3 in several runs of scenario 1 (which this PR implements when send_terminal_info_obs = true by setting done = true, info[idx]["TimeLimit.truncated"] = True and info[idx]["terminal_observation"][key] = terminal-obs on a truncation) and never got more than 0.25 in several runs of scenario 2. So this metric shows a small difference between the two. However, I will only post the results when this is absolutely necessary for this PR because of the sheer amount of comments we already have.

It sounds good now. Thanks for those tests and I'll take a further look at the implementation when possible.

Thanks and take your time 👍

Comment thread addons/godot_rl_agents/controller/ai_controller_2d.gd Outdated
Comment thread addons/godot_rl_agents/sync.gd Outdated
Comment thread addons/godot_rl_agents/sync.gd Outdated
stefanfausser and others added 3 commits June 7, 2026 14:55
Co-authored-by: Ivan-267 <61947090+Ivan-267@users.noreply.github.com>
…t extracting obs from dictionary and rebuilding the dictionary
@stefanfausser

stefanfausser commented Jun 7, 2026

Copy link
Copy Markdown
Contributor Author

Thank you for the review @Ivan-267. Hopefully all should be resolved by now.

As a reminder, also for myself, when (if) this PR is being accepted then the documentation should also be updated in order to explain why and how truncation and / or terminal observations can be utilized. However, the documentation of godot_rl_agents_plugin is unfortunately in the godot_rl_agents repo (and as far as I know this PR cannot immediately propose documentation changes to another repo so it must be done in another PR in the godot_rl_agents repo).

Luckily, the changes proposed in this PR will by default not change the behaviour of godot-rl-agent-plugin. They will only apply when send_terminal_obs_info of the sync node is set to true, so the documentation can wait.

@Ivan-267

Ivan-267 commented Jun 10, 2026

Copy link
Copy Markdown
Collaborator

I made a small PR to your branch (stefanfausser#1). Not well tested yet, will add some results when I check it a bit.

I wanted to resolve an error I got (more on the error below), but also made some changes such as:

  • Terminal obs will be recorded only if send terminal obs info is true:
    A small optimization.

  • Terminal obs will be stored only for training agents:
    Another small optimization. I added it after running an edge case test with episode timeouts at every few steps and noticed a performance difference in onnx inference. There's no requirement to store terminal obs for onnx inference, but if later needed, we can enable it easily.

  • Stored terminal obs are only checked from sync node when send terminal obs info is true:
    For full backward compatibility.

  • Tiny change of the asserts:
    Using is_empty() check as Godot supports it and the intention is directly readable.
    However, if this breaks your original intention or functionality in any way, you can set it back.

Error that this resolves:

I ran a short test on SimpleReachGoal (gdscript) and ran into some errors even when Send Terminal Obs Info is off.

The cause:

  1. The env sets done = true initially (during an initial env reset that happens in the ready method).
  2. SB3 training will call the reset method initially. After a reset, the Godot env sends only obs and info, but not dones (as no action was yet received or taken). Godot will not read or reset dones on that step, but _get_info_from_agents will read and erase the stored terminal obs. So done stays true (until an action step where it's read and sent), but terminal obs is erased, causing an assertion error.

It's worth noting that setting done to true initially is not necessarily useful here and the env could be modified to not set done initially, but it serves as a nice edge case to resolve. Potentially it would happen in other example envs and/or edge cases with resetting the env from Python.

My first idea was to check for and send terminal obs only during action steps (ignore reset steps), but the above alternative covers both this and the edge case of setting done before sync is ready.

Asserts:

I think you've included the raycast PR in the assert message commit. You can consider my PR after an undo of that commit (easier than removing the raycast changes after my commit), then re-implement the assert messages on top.

@stefanfausser

Copy link
Copy Markdown
Contributor Author

I made a small PR to your branch (stefanfausser#1). Not well tested yet, will add some results when I check it a bit.

Thank you @Ivan-267 for the great additions and the compatibility fix 👍 I tested your PR and have found no additional issues. However, I have two minor suggestions / improvements, see stefanfausser#1

I think you've included the raycast PR in the assert message commit. You can consider my PR after an undo of that commit (easier than removing the raycast changes after my commit), then re-implement the assert messages on top.

Right, I did commit something completely unrelated to this PR :( Thank you for noticing, I have undone that commit. After you have had a look at my two minor suggestions, I will accept your PR so this will then appear in this PR (this sounds so illogical but that is actually the way it is)

@Ivan-267

Copy link
Copy Markdown
Collaborator

Right, I did commit something completely unrelated to this PR :( Thank you for noticing, I have undone that commit. After you have had a look at my two minor suggestions, I will accept your PR so this will then appear in this PR (this sounds so illogical but that is actually the way it is)

While we could theoretically push changes to the same branch, I think the PR approach is much better in this case as for example both of us might be working on some changes at the same time, you might have been working on some changes while I was working on others.

@stefanfausser

Copy link
Copy Markdown
Contributor Author

You can consider my PR after an undo of that commit (easier than removing the raycast changes after my commit), then re-implement the assert messages on top.

Your PR was merged and I re-implemented the assert messages on top and committed it

@Ivan-267

Ivan-267 commented Jun 10, 2026

Copy link
Copy Markdown
Collaborator

Quick tests:
GDRL installed from the main repository. New main repository SB3 training script used, with LR set to 0.0001. Training from Godot editor with manual stopping.

Rollout plot colors (smoothing 0.9):
~Yellow: terminated on timeout
~Green: truncated on timeout

BallChase

reset_after set to 500 so truncation/termination will happen more frequently.

Extended AIController modified to:

func _physics_process(_delta):
	n_steps += 1
	if n_steps > reset_after:
		needs_reset = true
		#done = true
		truncated = true

	if needs_reset:
		_player.game_over()
		#needs_reset = false
		#reset()

Note

Done is uncommented and truncated is commented for testing termination

image image

Similar results on reward (perhaps slightly faster improvement for truncated, but run to run variance can have an effect here), a more noticeable difference on explained variance, which was better for truncated runs.

SimpleReachGoal (gdscript)

Extended AIController modified to:

func _physics_process(delta):
	n_steps += 1

	if n_steps > reset_after:
		truncated = true
		player.game_scene_manager.reset()

Note

truncated = true is commented when testing termination, as reset sets done = true.

And reset after set to 150 for more frequent termination/truncation.

image image

Similar results.

More complex envs and further testing might show more difference.

@stefanfausser Thanks for your work on this. I think we resolved all of the major points and it looks ready, only the docs remain, which can be written as another PR on the main repository.

@edbeeching This one looks good now, I'll approve, feel free to take a final look and give a go ahead and we can merge.

@Ivan-267
Ivan-267 requested a review from edbeeching June 10, 2026 21:03
@stefanfausser

stefanfausser commented Jun 10, 2026

Copy link
Copy Markdown
Contributor Author

Thanks @Ivan-267, I highly appreciated your review, the joint work and the testing

@stefanfausser

Copy link
Copy Markdown
Contributor Author

I am looking forward for a review @edbeeching . Please let me know If you need additional tests

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants