Potential Bug in `env_step()`

[modanesh]

In the env_step() method, there's a discrepancy between:

• The action used for the hidden state update (line where latent, _ = self.tssm(...) is called)
• The action actually executed in the environment (which can be different due to random pre-filling)

As a demonstration:

# Generate action from policy
action = self.policy_network(feat).sample().cpu()

# Update hidden state with this action
latent["hidden"] = self.tssm.slice_hidden(latent["hidden"])
hidden = (latent, action)

# But then potentially override the action for random pre-filling
if (self.replay_buffer.num_steps < self.config.pre_fill_steps) and self.config.random_pre_fill_steps:
    action = self.env.sample()  # This is a DIFFERENT action

# Execute the (potentially different) action
obs = self.env.step(action.argmax(dim=-1) if self.config.policy_discrete else action)

And then the new randomly sampled action is stored in the replay buffer.

This creates an inconsistency where:

• The hidden state (latent) contains information computed using the policy's action
• The replay buffer stores the actually executed action (which might be random)
• Future TSSM computations will use the hidden state that was computed with a different action than what was actually executed

This breaks the causal consistency of the world model, as the hidden state representation doesn't accurately reflect the true state transition that occurred in the environment.

Any clarification is appreciated!

[burchim]

Thank you very much for your feedback!

Indeed, this unintended behavior will happen during pre-filling. This may cause some confusion for the world model because of the incorrect state representation during the first few environment steps (just after random pre-filling). Fortunately, this only happens at the start of training when the world model is not trained yet.

Thanks for reporting this bug. I will update the code to store the correct action in the hidden state as follows:

# Generate action from policy
action = self.policy_network(feat).sample().cpu()

# Potentially override the action for random pre-filling
if (self.replay_buffer.num_steps < self.config.pre_fill_steps) and self.config.random_pre_fill_steps:
    action = self.env.sample()  # This is a DIFFERENT action

# Execute the (potentially different) action
obs = self.env.step(action.argmax(dim=-1) if self.config.policy_discrete else action)

# Update hidden state with the correct action
latent["hidden"] = self.tssm.slice_hidden(latent["hidden"])
hidden = (latent, action)

Best,
Maxime

[modanesh]

Thanks for your quick response!
I'm wondering how critical this bug can be in terms of hurting the performance/convergence of the world model.

[burchim]

The bug indeed may have an impact on the policy for the first few environment steps, but its impact on the overall training should be negligible.

The obtained replay buffer trajectories will be in the distribution of expected agent trajectories at this stage of learning. The resulting replay buffer trajectories will also be valid training trajectories.