Integrate Visual Execution Model (VEM): planning priors, eval harness, docs, and repo index#6
Conversation
|
CodeAnt AI is reviewing your PR. Thanks for using CodeAnt! 🎉We're free for open-source projects. if you're enjoying it, help us grow by sharing. Share on X · |
|
Warning Rate limit exceeded
You’ve run out of usage credits. Purchase more in the billing tab. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (15)
✨ Finishing Touches🧪 Generate unit tests (beta)
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 |
codeant-ai
Bot
added
the
size:XXL
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5ef3b40771
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| pt, ph, pw = cfg["encoder"]["patch_size"] | ||
| seq_len = (t // pt) * (h // ph) * (w // pw) | ||
| num_mask = max(1, seq_len // 4) | ||
| mask = torch.randperm(seq_len)[:num_mask] |
There was a problem hiding this comment.
The smoke test creates mask as index positions (torch.randperm(...)) instead of a boolean patch mask, but the training path feeds boolean masks from dataset/video_dataset.py. In apply_mask, ~mask on an integer tensor is interpreted as bitwise inversion, so the “visible” patches are selected from unintended negative indices and the test no longer validates real masking behavior. This can let integration checks pass even when the boolean-mask path is broken.
Useful? React with 👍 / 👎.
| model.eval() | ||
|
|
||
| bsz = 1 | ||
| t = cfg["encoder"].get("max_t", 8) |
There was a problem hiding this comment.
The integration clip uses t = max_t frames, but the encoder config defines max_t in patch units and training uses clip_len = max_t * patch_size[0]. With the current shorter clip, sequence length is reduced and a correct boolean mask of that length causes an index mismatch against rotary embeddings sized for max_t, so the smoke test avoids (rather than exercises) the production forward-path assumptions.
Useful? React with 👍 / 👎.
![codeant-ai[bot]](https://avatars.githubusercontent.com/in/646884?s=60&v=4){kind=small}
| pooled_next_state = next_state.mean(dim=1) if next_state.ndim > 2 else next_state | ||
| policy_query = self.model.policy_query_head(pooled_next_state) |
There was a problem hiding this comment.
Suggestion: _imagine_future now unconditionally calls self.model.policy_query_head, but MCTS is documented to require only predictor and value_head. Any existing model implementation without policy_query_head will now crash with AttributeError during expansion. Keep this path backward-compatible by guarding the call (as done in _expand) and returning a nullable/empty policy query when the head is absent. [api mismatch]
Severity Level: Major ⚠️
- ⚠️ Latent MCTS planner unusable with models sans policy head.
- ⚠️ Contract in MCTS docstring no longer matches implementation.Steps of Reproduction ✅
1. Inspect the MCTS contract in `models/vjepa/planning.py:120-137`, which documents
`model: the VJEPA model (must have predictor and value_head)` and does not mention any
`policy_query_head` requirement.
2. Observe `_imagine_future` in `models/vjepa/planning.py:20-51`, which unconditionally
executes `policy_query = self.model.policy_query_head(pooled_next_state)` at lines 47-49,
without checking for the attribute.
3. Instantiate `MCTS` with any model object that follows the documented minimum contract
(has `predictor.physics_engine` and `value_head` but no `policy_query_head`), then call
`plan()` as implemented in `models/vjepa/planning.py:10-60`; during expansion, `_expand()`
at lines 83-135 calls `_imagine_future(node.state, action)` and hits the unconditional
`policy_query_head` access.
4. At runtime, when `_imagine_future` executes for such a model, Python raises
`AttributeError: '<Model>' object has no attribute 'policy_query_head'` at
`models/vjepa/planning.py:47-49`, demonstrating that the planner now requires an
undocumented head and breaks backward compatibility, unlike `_expand`, which correctly
guards on `hasattr(self.model, "policy_query_head")` at lines 112-117.Fix in Cursor | Fix in VSCode Claude
(Use Cmd/Ctrl + Click for best experience)
Prompt for AI Agent 🤖
This is a comment left during a code review.
**Path:** models/vjepa/planning.py
**Line:** 187:188
**Comment:**
*Api Mismatch: `_imagine_future` now unconditionally calls `self.model.policy_query_head`, but `MCTS` is documented to require only `predictor` and `value_head`. Any existing model implementation without `policy_query_head` will now crash with `AttributeError` during expansion. Keep this path backward-compatible by guarding the call (as done in `_expand`) and returning a nullable/empty policy query when the head is absent.
Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix
| pt, ph, pw = cfg["encoder"]["patch_size"] | ||
| seq_len = (t // pt) * (h // ph) * (w // pw) | ||
| num_mask = max(1, seq_len // 4) | ||
| mask = torch.randperm(seq_len)[:num_mask] |
There was a problem hiding this comment.
Suggestion: mask is built as an index tensor, but the model's apply_mask path expects a boolean mask (True = masked). Passing indices here causes invalid boolean/integer inversion behavior in masking (~mask) and can produce incorrect indexing or runtime failures. Build a boolean mask tensor of length seq_len instead of a list of selected indices. [type error]
Severity Level: Major ⚠️
- ❌ Integration script misuses masking, not exercising boolean mask path.
- ⚠️ Potential future failures if batch size becomes greater one.Steps of Reproduction ✅
1. Run the integration harness via `python check_integrations.py` which executes `main()`
defined in `check_integrations.py:13-56`.
2. Inside `main()`, a random mask tensor is created at `check_integrations.py:29-32` as
`mask = torch.randperm(seq_len)[:num_mask]`, yielding a 1D LongTensor of indices with
shape `(num_mask,)`, not a boolean mask of length `seq_len`.
3. The `batch` dict is constructed at `check_integrations.py:34-39` with this index-valued
`mask` and passed into `VisualExecutionModel` from `models/vjepa/vjepa_model.py` via the
call `out = model(batch)` at `check_integrations.py:41`.
4. In `VisualExecutionModel.forward` (`models/vjepa/vjepa_model.py:86-101`), `mask =
batch["mask"]` is forwarded to `apply_mask(all_latents, mask)` and
`apply_mask(target_latents, mask)`; `apply_mask` is implemented in
`models/vjepa/utils.py:21-44` and documented to expect a boolean mask (`get_block_mask()`
at `utils.py:3-19` returns `dtype=torch.bool` and `apply_mask` comments at
`utils.py:23-25` state `mask: (seq_len,) or (bs, seq_len)` with boolean semantics).
5. Because `mask` is a LongTensor of indices, `apply_mask` treats it incorrectly: it uses
`~mask` at `utils.py:32` and `utils.py:41` intending boolean inversion, but bitwise
negation on integer indices produces negative values, so `x[i, ~mask[i]]` and `x[i,
mask[i]]` index using arbitrary integer positions instead of selecting visible vs masked
patches, breaking the masking logic used by the model in this integration check.Fix in Cursor | Fix in VSCode Claude
(Use Cmd/Ctrl + Click for best experience)
Prompt for AI Agent 🤖
This is a comment left during a code review.
**Path:** check_integrations.py
**Line:** 32:32
**Comment:**
*Type Error: `mask` is built as an index tensor, but the model's `apply_mask` path expects a boolean mask (`True` = masked). Passing indices here causes invalid boolean/integer inversion behavior in masking (`~mask`) and can produce incorrect indexing or runtime failures. Build a boolean mask tensor of length `seq_len` instead of a list of selected indices.
Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix| available_actions = torch.randn(num_actions, 128, device=device) | ||
| pooled = traj_a[-1].mean(dim=1) | ||
| query = model.policy_query_head(pooled).squeeze(0) | ||
| logits = torch.matmul(available_actions, query) | ||
| probs = torch.softmax(logits, dim=0) | ||
| max_prior = float(probs.max().item()) |
There was a problem hiding this comment.
Suggestion: The script accepts --num-actions without enforcing it to be positive, but later computes probs.max(). When num_actions is 0, available_actions and probs are empty and probs.max() raises a runtime error. Reject non-positive values (or handle the empty-action case) before computing priors. [logic error]
Severity Level: Major ⚠️
- ❌ World-model eval crashes when --num-actions is zero.
- ⚠️ Evaluation harness unusable for configurations with no candidate actions.Steps of Reproduction ✅
1. Invoke the world-model evaluation harness with zero actions, e.g. `python
evaluate_world_model.py --num-actions 0`, which runs `main()` in
`evaluate_world_model.py:116-166`.
2. Argument parsing at `evaluate_world_model.py:117-123` stores `args.num_actions = 0`;
there is no validation enforcing this to be positive before `args` is used.
3. `main()` calls `evaluate_metrics(model=model, device=device,
rollout_steps=args.rollout_steps, num_actions=args.num_actions)` at
`evaluate_world_model.py:137-142`, passing `num_actions=0` into `evaluate_metrics` defined
at `evaluate_world_model.py:76-113`.
4. Inside `evaluate_metrics`, `available_actions = torch.randn(num_actions, 128,
device=device)` at line 100 creates an empty tensor of shape `(0, 128)`, leading to
`logits = torch.matmul(available_actions, query)` at line 103 and `probs =
torch.softmax(logits, dim=0)` at line 104 also being empty; the subsequent reduction
`probs.max()` at line 105 is applied to an empty tensor and raises a runtime
`RuntimeError` (max reduction on zero elements), causing the evaluation script to crash
before writing the manifest or printing metrics.Fix in Cursor | Fix in VSCode Claude
(Use Cmd/Ctrl + Click for best experience)
Prompt for AI Agent 🤖
This is a comment left during a code review.
**Path:** evaluate_world_model.py
**Line:** 100:105
**Comment:**
*Logic Error: The script accepts `--num-actions` without enforcing it to be positive, but later computes `probs.max()`. When `num_actions` is `0`, `available_actions` and `probs` are empty and `probs.max()` raises a runtime error. Reject non-positive values (or handle the empty-action case) before computing priors.
Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix
|
CodeAnt AI finished reviewing your PR. |
1 participant
User description
Motivation
Description
CODE_ADDRESS_INDEX.md,AUDIT_RECHECK.md, andRECTIFICATION_STATUS.mdto record line-level mapping, recheck steps, and rectification summary.evaluate_world_model.py,evaluate_perception.py, andcheck_integrations.pyfor smoke metrics and wiring checks.policy_query_headtomodels/vjepa/vjepa_model.py, updatedmodels/vjepa/planning.pyto return/use policy-query vectors and compute action priors by similarity, and preserved backward-compatibleVisualExecutionModelalias.vjepa_train.pynow accepts--config, detectsffmpegbefore generating synthetic videos, usestraining.epochsdefaulting to 100, and improved data-directory handling.models/adaptive_depth.pynow constructs returned carry usingtype(carry)instead of a concrete class,models/topological.pyfixes shapes and pooling logic for Betti estimates, and other wiring updates acrossmodels/*referenced by docs and the index.Testing
python -m compileall -q .completed successfully with no compile errors.python check_integrations.pywhich performed a forward pass and a small MCTS plan and printed an "Integration check passed" message.python evaluate_world_model.py --config config/vjepa_micro.yaml --seed 42which produced and saved JSON metrics (rollout drift, trajectory divergence, action-prior stats).python evaluate_perception.py --config config/vjepa_micro.yaml --seed 42which produced and saved latent-consistency metrics; all scripted smoke checks completed without runtime failures.Codex Task
CodeAnt-AI Description
Add VEM smoke checks, policy-guided planning, and repo recovery notes
What Changed
ffmpegis missing or when epoch count is set in config.Impact
✅ Faster model smoke checks✅ Clearer planning action selection✅ Fewer training setup failures💡 Usage Guide
Checking Your Pull Request
Every time you make a pull request, our system automatically looks through it. We check for security issues, mistakes in how you're setting up your infrastructure, and common code problems. We do this to make sure your changes are solid and won't cause any trouble later.
Talking to CodeAnt AI
Got a question or need a hand with something in your pull request? You can easily get in touch with CodeAnt AI right here. Just type the following in a comment on your pull request, and replace "Your question here" with whatever you want to ask:
This lets you have a chat with CodeAnt AI about your pull request, making it easier to understand and improve your code.
Example
Preserve Org Learnings with CodeAnt
You can record team preferences so CodeAnt AI applies them in future reviews. Reply directly to the specific CodeAnt AI suggestion (in the same thread) and replace "Your feedback here" with your input:
This helps CodeAnt AI learn and adapt to your team's coding style and standards.
Example
Retrigger review
Ask CodeAnt AI to review the PR again, by typing:
Check Your Repository Health
To analyze the health of your code repository, visit our dashboard at https://app.codeant.ai. This tool helps you identify potential issues and areas for improvement in your codebase, ensuring your repository maintains high standards of code health.