Merged
Conversation
wesselvannierop
added
documentation
Contributor
WalkthroughRefactors entropy computation and selection: Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Possibly related PRs
Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 inconclusive)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
📝 Coding Plan
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
tests/test_agent.py (1)
138-142: Assert the grouped-action behavior, not just the shapes.This is the only coverage for
n_possible_actions < w, but it still passes if the implementation returns the right shapes with the wrong number of selected actions or the wrong column expansion inmask.Possible assertion upgrade
agent = selection.GreedyEntropy(n_actions, w // 2, h, w) selected_lines, mask = agent.sample(particles) assert mask.shape == (1, h, w) assert selected_lines.shape == (1, w // 2) + assert ops.count_nonzero(selected_lines[0]) == n_actions + assert ops.all(mask[:, 0] == ops.repeat(selected_lines, agent.stack_n_cols, axis=1))🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/test_agent.py` around lines 138 - 142, Update the test to verify grouped-action semantics from GreedyEntropy.sample: assert that selected_lines contains exactly n_possible_actions entries (len(selected_lines[0]) == n_possible_actions) and that mask reflects grouping of width group_width = w // n_possible_actions by checking for each selected action index a in selected_lines[0] the mask has ones in columns a*group_width:(a+1)*group_width (and zeros elsewhere), and that mask.sum() == n_possible_actions * group_width; use the same names (GreedyEntropy.sample, selected_lines, mask, n_possible_actions, w) to locate and implement these assertions.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@zea/agent/masks.py`:
- Around line 38-42: The result of indices_to_k_hot currently uses ops.any which
always yields a boolean tensor, losing the requested dtype; change
indices_to_k_hot to cast the boolean mask returned by ops.any(...) to the
provided dtype (e.g., use ops.cast(..., dtype)) before returning so callers like
make_line_mask(..., dtype=...) receive the correct numeric type.
In `@zea/agent/selection.py`:
- Around line 201-220: The docstring for reweight_entropies_around_line is
incorrect: update the Returns section to indicate it returns only the updated
entropy vector (not a tuple with the selected line index). Edit the Returns
description to specify the returned value is a Tensor (or same type as
entropy_per_line) of shape (n_possible_actions,) representing the updated
entropies per line, and remove any mention of a selected line index; reference
the function name reweight_entropies_around_line and its parameters
entropy_per_line and line_index to locate the docstring to modify.
- Around line 260-269: The current code derives linewise_action scores by
summing pixelwise entropy (compute_pixelwise_entropy) across columns, which
treats grouped columns as independent and can misrank actions; instead, compute
entropy at the grouped-action level inside sample(): aggregate particle
probabilities/logits across each action's span (self.stack_n_cols) to form a
single probability distribution per action (using particles' logits or summed
probabilities for columns belonging to the same action), then compute the
entropy of that action-level distribution to produce
column_wise_entropy/linewise_entropy; update the logic that builds
column_wise_entropy/linewise_entropy in sample() to aggregate by action
(respecting self.n_possible_actions and self.stack_n_cols) rather than summing
pixelwise_entropy, and keep compute_pixelwise_entropy only for visualization.
---
Nitpick comments:
In `@tests/test_agent.py`:
- Around line 138-142: Update the test to verify grouped-action semantics from
GreedyEntropy.sample: assert that selected_lines contains exactly
n_possible_actions entries (len(selected_lines[0]) == n_possible_actions) and
that mask reflects grouping of width group_width = w // n_possible_actions by
checking for each selected action index a in selected_lines[0] the mask has ones
in columns a*group_width:(a+1)*group_width (and zeros elsewhere), and that
mask.sum() == n_possible_actions * group_width; use the same names
(GreedyEntropy.sample, selected_lines, mask, n_possible_actions, w) to locate
and implement these assertions.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: fe180884-ff24-48ca-a50c-750c7ba10c46
📒 Files selected for processing (3)
tests/test_agent.pyzea/agent/masks.pyzea/agent/selection.py
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (2)
zea/agent/selection.py (2)
201-220:⚠️ Potential issue | 🟡 Minor
reweight_entropies_around_line()docstring still reflects the old helper behavior.This method no longer selects a line or returns a tuple; it only reweights around the provided
line_indexand returns the updated vector. The current Args/Returns text is still describing the old behavior.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@zea/agent/selection.py` around lines 201 - 220, Update the docstring for reweight_entropies_around_line to reflect its current behavior: it does not select a line or return a tuple but instead reweights the provided entropy_per_line vector around the given line_index and returns the updated entropy vector. Modify the Args section to describe entropy_per_line (Tensor of shape (n_possible_actions,)) and line_index (int) as inputs, and change the Returns section to state that the function returns a Tensor of the same shape with reweighted entropies (selected line set to 0, neighbours decreased). Remove or adjust any text that claims the method selects a line or returns (selected_line_index, entropies) and keep the note about torch backend compatibility if still applicable.
260-269:⚠️ Potential issue | 🟠 MajorLine/action scores are still being derived from pixelwise entropy after the log step.
That is not the entropy of the full observed line/action. Once an action spans multiple pixels, summing
compute_pixelwise_entropy()outputs can change the ranking because the Gaussian-mixture entropy needs to be computed on the aggregated action span beforelog, not after. Please keepcompute_pixelwise_entropy()for visualization, and compute action-level entropy directly insidesample().🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@zea/agent/selection.py` around lines 260 - 269, The current code builds action/line scores by summing outputs of compute_pixelwise_entropy (pixelwise_entropy → column_wise_entropy → linewise_entropy), which is incorrect because mixtures must be combined across the action span before taking the log/entropy; instead, keep compute_pixelwise_entropy for visualization only and compute the action-level entropy inside sample() by aggregating the Gaussian-mixture parameters across the span defined by stack_n_cols and n_possible_actions and then computing the entropy of that aggregated mixture (do not sum pixelwise entropies). Locate compute_pixelwise_entropy and replace the linewise_entropy construction in sample() with a call that aggregates mixture parameters over the span (or a new helper that computes entropy_from_span(particles, span_start, span_width)) and use that result for selection while leaving pixelwise_entropy/column_wise_entropy available for viz.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@zea/agent/selection.py`:
- Around line 171-184: The docstring for compute_pixelwise_entropy still refers
to "each line" and is misleading given the implementation returns per-pixel
entropies; update the description and summary to state it computes entropy per
pixel using a Gaussian Mixture Model approximation of the posterior (not
per-line), clarify the input shape (particles: (batch_size, n_particles,
*pixels)) and the returned shape (Tensor of shape (batch_size, *pixels)), and
adjust any mention of "line" to "pixel" throughout the docstring so the text
matches the function behavior.
---
Duplicate comments:
In `@zea/agent/selection.py`:
- Around line 201-220: Update the docstring for reweight_entropies_around_line
to reflect its current behavior: it does not select a line or return a tuple but
instead reweights the provided entropy_per_line vector around the given
line_index and returns the updated entropy vector. Modify the Args section to
describe entropy_per_line (Tensor of shape (n_possible_actions,)) and line_index
(int) as inputs, and change the Returns section to state that the function
returns a Tensor of the same shape with reweighted entropies (selected line set
to 0, neighbours decreased). Remove or adjust any text that claims the method
selects a line or returns (selected_line_index, entropies) and keep the note
about torch backend compatibility if still applicable.
- Around line 260-269: The current code builds action/line scores by summing
outputs of compute_pixelwise_entropy (pixelwise_entropy → column_wise_entropy →
linewise_entropy), which is incorrect because mixtures must be combined across
the action span before taking the log/entropy; instead, keep
compute_pixelwise_entropy for visualization only and compute the action-level
entropy inside sample() by aggregating the Gaussian-mixture parameters across
the span defined by stack_n_cols and n_possible_actions and then computing the
entropy of that aggregated mixture (do not sum pixelwise entropies). Locate
compute_pixelwise_entropy and replace the linewise_entropy construction in
sample() with a call that aggregates mixture parameters over the span (or a new
helper that computes entropy_from_span(particles, span_start, span_width)) and
use that result for selection while leaving
pixelwise_entropy/column_wise_entropy available for viz.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 62a9ccf0-6358-4885-a0cc-4d6a4d51f8b2
📒 Files selected for processing (1)
zea/agent/selection.py
Contributor
There was a problem hiding this comment.
♻️ Duplicate comments (1)
zea/agent/selection.py (1)
218-220:⚠️ Potential issue | 🟡 MinorDocstring return type is still inaccurate.
On Line 219,
reweight_entropies_around_line(...)is documented as returning aTuple, but it returns a single tensor.Suggested doc-only fix
- Returns: - Tuple: The reweighted entropy per line, of shape (n_possible_actions,) + Returns: + Tensor: The reweighted entropy per line, of shape (n_possible_actions,)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@zea/agent/selection.py` around lines 218 - 220, The docstring for reweight_entropies_around_line currently states it returns a Tuple but the function actually returns a single tensor; update the Returns section to list the correct type (e.g., torch.Tensor or Tensor) and describe the shape (n_possible_actions,) and meaning (the reweighted entropy per line) to match the implementation in reweight_entropies_around_line.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@zea/agent/selection.py`:
- Around line 218-220: The docstring for reweight_entropies_around_line
currently states it returns a Tuple but the function actually returns a single
tensor; update the Returns section to list the correct type (e.g., torch.Tensor
or Tensor) and describe the shape (n_possible_actions,) and meaning (the
reweighted entropy per line) to match the implementation in
reweight_entropies_around_line.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 0a03bd4d-60f9-497f-8328-e8b22b1e6ed3
📒 Files selected for processing (2)
zea/agent/masks.pyzea/agent/selection.py
tristan-deep
approved these changes
Mar 13, 2026
Collaborator
tristan-deep
left a comment
tristan-deep
left a comment
There was a problem hiding this comment.
lgtm, 2 questions.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Found some inconsistencies after the updates in #266 in the docstrings and decided to clean it up a bit.
indices_to_k_hotbut this was not used by somezea.agent.selectionclasses, which used their own implementation (which was actually better).compute_pixelwise_entropyis simplified, is a staticmethod now for standalone usage, and can be used on N-D data, which is nice to visualize the Entropy maps for 3D data.Summary by CodeRabbit
Refactor
Tests