Skip to content

local_ego_id changes for bsz > 1 #18

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
m2kulkarni wants to merge 1 commit into
base: main
Choose a base branch
from
Open

local_ego_id changes for bsz > 1 #18

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions .gitignore
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -184,3 +184,6 @@ pufferlib/resources/drive/output_agent.gif pufferlib/resources/drive/output.gif
artifacts/
# Local drive renders
pufferlib/resources/drive/output*.mp4

# Local TODO tracking
TODO.md
14 changes: 10 additions & 4 deletions pufferlib/pufferl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -302,17 +302,23 @@ def evaluate(self):
total_agents = len(o)
num_agents_per_env = total_agents // batch_size

# Convert global ego_ids to local (per-environment) indices
# ego_ids contains ALL ego agents across ALL environments, not just this batch
# Count how many ego IDs belong to the first environment (IDs < num_agents_per_env)
num_ego_per_env = sum(1 for eid in ego_ids if eid < num_agents_per_env)
local_ego_ids = [eid % num_agents_per_env for eid in ego_ids[:num_ego_per_env]]
Comment on lines +308 to +309
Copy link

@greptile-apps greptile-apps bot Jan 31, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

assumes all environments have the same number of ego agents at identical local indices. If env 0 has ego at local index [0, 1] and env 1 has ego at local index [2, 3], this will fail and only select indices [0, 1] from both environments. verify this assumption holds in your use case, or consider using the pre-computed local_ego_ids from drive.py instead

Prompt To Fix With AI 
This is a comment left during a code review.
Path: pufferlib/pufferl.py
Line: 308:309

Comment:
assumes all environments have the same number of ego agents at identical local indices. If env 0 has ego at local index [0, 1] and env 1 has ego at local index [2, 3], this will fail and only select indices [0, 1] from both environments. verify this assumption holds in your use case, or consider using the pre-computed `local_ego_ids` from `drive.py` instead

How can I resolve this? If you propose a fix, please make it concise.


original_shape = o.shape

o = o.reshape(batch_size, num_agents_per_env, *original_shape[1:])
r = r.reshape(batch_size, num_agents_per_env)
d = d.reshape(batch_size, num_agents_per_env)
t = t.reshape(batch_size, num_agents_per_env)

o = o[:, ego_ids].reshape(batch_size * len(ego_ids), *original_shape[1:])
r = r[:, ego_ids].flatten()
d = d[:, ego_ids].flatten()
t = t[:, ego_ids].flatten()
o = o[:, local_ego_ids].reshape(batch_size * num_ego_per_env, *original_shape[1:])
r = r[:, local_ego_ids].flatten()
d = d[:, local_ego_ids].flatten()
t = t[:, local_ego_ids].flatten()
else:
o = o[ego_ids]
r = r[ego_ids]
Expand Down