offsets shape bug
#72
Conversation
WalkthroughThe pull request introduces modifications to several configuration files, notably Changes
Possibly related PRs
Suggested reviewers
Poem
Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (3)
- configs/model/rodent.yaml (3 hunks)
- configs/stac/demo.yaml (1 hunks)
- output.txt (1 hunks)
🧰 Additional context used
🪛 LanguageTool
output.txt
[grammar] ~181-~181: Make sure that the singular noun after the number ‘11.590774774551392’ is correct.
Context: ...eplace 1 finished in 11.590774774551392 starting q_opt 2 q_opt 2 finished in 0.020217895...(CD_POINT_CD_NN)
[grammar] ~196-~196: This phrase is duplicated. You should probably use “pose optimization” only once.
Context: ...ion finished in 31.25155758857727 Final pose optimization Pose Optimization: Pose Optimization done in 0.5810093879...(PHRASE_REPETITION)
[style] ~213-~213: This adverb was used twice in the sentence. Consider removing one of them or replacing them with a synonym.
Context: ... in run_stac ik_only_data = stac.ik_only(kp_data, offsets) ^^...(ADVERB_REPETITION_PREMIUM)
[uncategorized] ~238-~238: Commas should not be placed before a closing parenthesis. Either move the comma outside of the parentheses, or remove it altogether.
Context: ... Incompatible shapes for broadcasting: (69,) and requested shape (23, 3) -----------...(COMMA_CLOSING_PARENTHESIS)
🔇 Additional comments (8)
configs/stac/demo.yaml (2)
Line range hint
10-12: Verify indentation correction for mujoco sectionThe AI summary mentions that the indentation of
iterationsandls_iterationsunder themujocosection has been corrected. However, this change is not visible in the provided code snippet.Please confirm if:
- The indentation correction was made in a previous commit.
- The AI summary is inconsistent with the actual changes.
You can use the following script to check the git history for recent changes to this file:
This will display the last 5 commits that modified this file, allowing you to verify if the indentation was corrected in a recent commit.
7-7: Confirm the logic change for inverse kinematics processingThe parameter
ik_only: Truehas been replaced withskip_ik_only: False. While this maintains the same effective behavior (IK processing will be performed), the change in naming convention implies a potential logic inversion in the code that uses this parameter.Please ensure that:
- The code that consumes this parameter has been updated to handle
skip_ik_onlyinstead ofik_only.- The logic has been adjusted to accommodate the inverted boolean meaning.
You can use the following script to verify the usage of this parameter in the codebase:
This will help identify any places where the parameter might still be referenced by its old name or where the logic might need to be updated.
configs/model/rodent.yaml (4)
4-4: Clarify the overall intention behind parameter reductions.I've noticed a pattern where multiple parameters (N_FRAMES_PER_CLIP, N_ITERS, N_SAMPLE_FRAMES) have all been reduced to 1. This consistent change across different aspects of the model suggests a deliberate modification to the system's behavior.
Could you please provide more context on the overall goal of these changes? Are these modifications part of:
- A debugging process?
- A performance optimization effort?
- A shift in the model's intended use or functionality?
Understanding the broader intention will help ensure that these changes align with the project's objectives and don't unintentionally compromise the model's effectiveness.
Also applies to: 13-13, 192-192
192-192: Verify the intentionality of reducing N_SAMPLE_FRAMES to 1.The number of sample frames has been drastically reduced from 100 to 1. This change could significantly impact the statistical robustness of any sampling-based operations in the model.
Please confirm if this change is intentional and explain the reasoning behind it. If it's for debugging or a specific use case, consider adding a comment to clarify the purpose of this change.
To help verify the impact, you can run the following script:
#!/bin/bash # Description: Check for any dependencies on N_SAMPLE_FRAMES in the codebase # Search for usage of N_SAMPLE_FRAMES echo "Searching for N_SAMPLE_FRAMES usage:" rg --type python -i "N_SAMPLE_FRAMES" # Search for any sampling-related operations that might be affected echo "Searching for potential sampling-related operations:" rg --type python -e "sample" -e "random.choice" -e "np.random"This script will help identify any code that might be affected by this change.
13-13: Verify the intentionality of reducing N_ITERS to 1.The number of alternating pose and offset optimization rounds has been significantly reduced from 6 to 1. This change could have a substantial impact on the convergence and accuracy of the optimization process.
Please confirm if this change is intentional and explain the reasoning behind it. If it's for debugging or a specific use case, consider adding a comment to clarify the purpose of this change.
To help verify the impact, you can run the following script:
#!/bin/bash # Description: Check for any dependencies on N_ITERS in the codebase # Search for usage of N_ITERS echo "Searching for N_ITERS usage:" rg --type python -i "N_ITERS" # Search for any iteration-related loops that might be affected echo "Searching for potential iteration-related loops:" rg --type python -e "for.*range.*iter" -e "while.*iter"This script will help identify any code that might be affected by this change.
4-4: Verify the intentionality of reducing N_FRAMES_PER_CLIP to 1.The number of frames per clip for inverse kinematics has been drastically reduced from 250 to 1. This change could significantly impact the accuracy and stability of the IK solution.
Please confirm if this change is intentional and explain the reasoning behind it. If it's for debugging or a specific use case, consider adding a comment to clarify the purpose of this change.
To help verify the impact, you can run the following script:
This script will help identify any code that might be affected by this change.
output.txt (2)
179-202: Optimization process completed successfully.The optimization process has completed successfully, showing decreasing error values over iterations. The final pose optimization achieved a low error of 5.823266474180855e-05, indicating a good fit to the motion capture data.
Great job on implementing an effective optimization process!
🧰 Tools
🪛 LanguageTool
[grammar] ~181-~181: Make sure that the singular noun after the number ‘11.590774774551392’ is correct.
Context: ...eplace 1 finished in 11.590774774551392 starting q_opt 2 q_opt 2 finished in 0.020217895...(CD_POINT_CD_NN)
[grammar] ~196-~196: This phrase is duplicated. You should probably use “pose optimization” only once.
Context: ...ion finished in 31.25155758857727 Final pose optimization Pose Optimization: Pose Optimization done in 0.5810093879...(PHRASE_REPETITION)
203-242:⚠️ Potential issueFix shape mismatch in ik_only function.
There's a critical error in the
ik_onlyfunction due to incompatible shapes for broadcasting. The error occurs when trying to set site positions:ValueError: Incompatible shapes for broadcasting: (69,) and requested shape (23, 3)This suggests that the
offsetsparameter doesn't match the expected shape for the site positions.To resolve this issue:
- Check the shape of the
offsetsparameter passed to theik_onlyfunction. It should be (23, 3) but is currently (69,).- Verify that the
offsetsare correctly prepared before callingik_only. You may need to reshape the data:offsets = offsets.reshape(-1, 3) # Reshape to (23, 3) if it's a flat array of 69 elements
- Review the
set_site_posfunction inop_utils.pyto ensure it's handling the input shapes correctly.- Add shape assertions or validation at the beginning of the
ik_onlyfunction to catch such issues early:assert offsets.shape == (23, 3), f"Expected offsets shape (23, 3), but got {offsets.shape}"To help diagnose the issue, please run the following script to check the shapes of relevant variables:
This script will help us understand how the
ik_onlyfunction is defined and used, and what shapes thekp_dataandoffsetsvariables have. The results will guide us in pinpointing the source of the shape mismatch.🧰 Tools
🪛 LanguageTool
[style] ~213-~213: This adverb was used twice in the sentence. Consider removing one of them or replacing them with a synonym.
Context: ... in run_stac ik_only_data = stac.ik_only(kp_data, offsets) ^^...(ADVERB_REPETITION_PREMIUM)
[uncategorized] ~238-~238: Commas should not be placed before a closing parenthesis. Either move the comma outside of the parentheses, or remove it altogether.
Context: ... Incompatible shapes for broadcasting: (69,) and requested shape (23, 3) -----------...(COMMA_CLOSING_PARENTHESIS)
| FootR: 0.02109994311367263 -0.007614327297542731 -0.002752166493087734 | ||
| HandL: 0.003031746915839468 0.001515873457919734 -0.008337304018558537 | ||
| HandR: 0.003031746915839468 -0.001515873457919734 -0.008337304018558537 | ||
| HipL: -0.015 0.015 0. | ||
| HipR: -0.015 -0.015 0. | ||
| KneeL: 0.01542264534586861 0.01747899805865108 -0.02570440890978101 | ||
| KneeR: 0.01542264534586861 -0.01747899805865108 -0.02570440890978101 | ||
| ShoulderL: 0.0287 0.00984 -0.02542 | ||
| ShoulderR: 0.0287 -0.00984 -0.02542 | ||
| Snout: 0. 0. 0. | ||
| SpineF: -0.015 0. 0.0 | ||
| SpineL: 0. 0. 0.0 | ||
| SpineM: 0. 0. 0.0 | ||
| TailBase: 0. 0. 0. | ||
| WristL: 0. 0. 0.0 | ||
| WristR: 0. 0. 0.0 | ||
| TRUNK_OPTIMIZATION_KEYPOINTS: | ||
| - SpineF | ||
| - SpineL | ||
| - SpineM | ||
| - HipL | ||
| - HipR | ||
| - ShoulderL | ||
| - ShoulderR | ||
| - TailBase | ||
| INDIVIDUAL_PART_OPTIMIZATION: | ||
| r_leg: | ||
| - vertebra_1 | ||
| - vertebra_2 | ||
| - vertebra_3 | ||
| - vertebra_4 | ||
| - vertebra_5 | ||
| - vertebra_6 | ||
| - hip_R | ||
| - knee_R | ||
| - ankle_R | ||
| - foot_R | ||
| l_leg: | ||
| - vertebra_1 | ||
| - vertebra_2 | ||
| - vertebra_3 | ||
| - vertebra_4 | ||
| - vertebra_5 | ||
| - vertebra_6 | ||
| - hip_L | ||
| - knee_L | ||
| - ankle_L | ||
| - foot_L | ||
| r_arm: | ||
| - scapula_R | ||
| - shoulder_R | ||
| - elbow_R | ||
| - hand_R | ||
| - finger_R | ||
| l_arm: | ||
| - scapula_L | ||
| - shoulder_L | ||
| - elbow_L | ||
| - hand_L | ||
| - finger_L | ||
| head: | ||
| - atlas | ||
| - cervical | ||
| - atlant_extend | ||
| KEYPOINT_COLOR_PAIRS: | ||
| AnkleL: 0 0 .3 1 | ||
| AnkleR: .3 0 0 1 | ||
| EarL: 0 0 .8 1 | ||
| EarR: .8 0 0 1 | ||
| ElbowL: 0 0 .8 1 | ||
| ElbowR: .8 0 0 1 | ||
| FootL: 0 0 .3 1 | ||
| FootR: .3 0 0 1 | ||
| HandL: 0 0 .8 1 | ||
| HandR: .8 0 0 1 | ||
| HipL: 0 0 .3 1 | ||
| HipR: .3 0 0 1 | ||
| KneeL: 0 0 .3 1 | ||
| KneeR: .3 0 0 1 | ||
| ShoulderL: 0 0 .8 1 | ||
| ShoulderR: .8 0 0 1 | ||
| Snout: .8 .8 .8 1 | ||
| SpineF: .8 .8 .8 1 | ||
| SpineL: .8 .8 .8 1 | ||
| SpineM: .8 .8 .8 1 | ||
| TailBase: .8 .8 .8 1 | ||
| WristL: 0 0 .8 1 | ||
| WristR: .8 0 0 1 | ||
| SCALE_FACTOR: 0.9 | ||
| MOCAP_SCALE_FACTOR: 0.001 | ||
| SITES_TO_REGULARIZE: | ||
| - HandL | ||
| - HandR | ||
| - SpineF | ||
| - SpineL | ||
| - SpineM | ||
| RENDER_FPS: 50 | ||
| N_SAMPLE_FRAMES: 1 | ||
| M_REG_COEF: 1 | ||
|
|
There was a problem hiding this comment.
🛠️ Refactor suggestion
Consider improving configuration flexibility and maintainability.
The configuration is comprehensive and well-structured. However, to enhance maintainability and flexibility, consider the following suggestions:
- Use environment variables or command-line arguments for paths (e.g.,
fit_path,ik_only_path,data_path,MJCF_PATH) to make it easier to run in different environments. - Consider moving hard-coded values (e.g.,
KEYPOINT_INITIAL_OFFSETS,KEYPOINT_COLOR_PAIRS) to separate configuration files that can be easily updated without changing the main configuration. - Use a configuration management system to handle different environments (development, testing, production) more effectively.
Would you like assistance in implementing these improvements?
| [2024-10-09 14:12:39,137][root][INFO] - Running fit. Mocap data shape: (1, 69) | ||
| /home/talmolab/miniconda3/envs/stac-mjx-env/lib/python3.11/site-packages/jaxopt/_src/optax_wrapper.py:120: DeprecationWarning: jax.tree_map is deprecated: use jax.tree.map (jax v0.4.25 or newer) or jax.tree_util.tree_map (any JAX version). |
There was a problem hiding this comment.
Update jax.tree_map usage to address deprecation warning.
The code is using a deprecated function jax.tree_map. To future-proof your code, update it to use either jax.tree.map (for JAX v0.4.25 or newer) or jax.tree_util.tree_map (for any JAX version).
Replace the deprecated jax.tree_map with the appropriate alternative based on your JAX version:
# For JAX v0.4.25 or newer
from jax import tree
result = tree.map(update_fun, params, updates)
# For any JAX version
from jax.tree_util import tree_map
result = tree_map(update_fun, params, updates)
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #72 +/- ##
==========================================
- Coverage 42.75% 40.95% -1.80%
==========================================
Files 9 10 +1
Lines 545 586 +41
==========================================
+ Hits 233 240 +7
- Misses 312 346 +34 ☔ View full report in Codecov by Sentry. |
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Outside diff range and nitpick comments (1)
demos/api_usage.ipynb (1)
Line range hint
1-555: Summary of changes and recommendations fordemos/api_usage.ipynbThe notebook demonstrates significant changes in the
stac_mjxlibrary's performance and output:
- Improved accuracy in optimization results, with lower mean errors.
- More detailed error reporting, providing arrays of error values instead of single values.
- Substantial increases in execution time for various steps, particularly in compilation and Pose Optimization.
- Addition of standard deviation reporting for optimization results.
While the improvements in accuracy are commendable, the increased execution times warrant careful consideration. It's recommended to:
- Investigate the causes of increased compilation and execution times.
- Evaluate the time-accuracy trade-off and consider making it configurable.
- Ensure that the more detailed error reporting is being utilized effectively for analysis and debugging.
- Update the documentation to reflect these changes and provide guidance on interpreting the new output format.
Consider implementing a configurable mode that allows users to choose between high-accuracy (current behavior) and high-speed (previous behavior) optimization, depending on their specific needs and computational resources.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (2)
- demos/api_usage.ipynb (3 hunks)
- stac_mjx/stac.py (1 hunks)
🧰 Additional context used
🪛 GitHub Check: codecov/patch
stac_mjx/stac.py
[warning] 393-393: stac_mjx/stac.py#L393
Added line #L393 was not covered by tests
🔇 Additional comments (6)
stac_mjx/stac.py (1)
393-393: Approved: Improved offset formatting for consistencyThe change to reshape
self._offsetsensures that the offsets are consistently formatted as a 2D array with 3 columns when not batched. This improvement in data consistency is crucial for subsequent processing.To ensure this change doesn't introduce any regressions, please add a test case that covers this scenario. You can use the following script to check for existing tests:
If no relevant tests are found, consider adding a new test case to cover this scenario.
🧰 Tools
🪛 GitHub Check: codecov/patch
[warning] 393-393: stac_mjx/stac.py#L393
Added line #L393 was not covered by testsdemos/api_usage.ipynb (5)
144-347: Enhanced error reporting in optimization process.The error output for
q_opt 1has been significantly expanded:
- Instead of a single error value, it now provides an array of 600 float values.
- This detailed error reporting could represent errors for individual iterations or data points during the optimization process.
- The
Replace 1operation time has slightly decreased from 44.80 to 42.78 seconds, indicating a minor performance improvement.These changes offer more granular insights into the optimization process, which can be valuable for debugging and fine-tuning the algorithm.
To better understand the implications of this change, please run the following script:
#!/bin/bash # Description: Investigate the implementation of the new error reporting mechanism # Test: Look for changes in error calculation and reporting rg --type python -g '*.py' -e '(calculate_error|report_error|error_array)' -C 10 # Test: Check for any new visualization or analysis functions for the error array rg --type python -g '*.py' -e '(plot_error|analyze_error|error_statistics)' -C 5
553-555:⚠️ Potential issueSignificant changes in Pose Optimization performance and results.
The Pose Optimization step shows notable changes:
- Execution time has increased dramatically from 0.098 seconds to 46.727 seconds.
- Mean error has improved from 0.00022817964782007039 to 0.0001599203678779304.
- A standard deviation of 5.567146581597626e-05 is now reported, providing additional statistical insight.
While the accuracy has improved, the substantial increase in execution time warrants careful consideration of the time-accuracy trade-off.
To better understand these changes and their implications, please run the following script:
#!/bin/bash # Description: Investigate changes in Pose Optimization implementation and performance # Test: Look for recent changes in Pose Optimization implementation rg --type python -g '*.py' -e 'def (pose_optimization|optimize_pose)' -C 20 # Test: Check for any new computations or algorithm changes in the Pose Optimization process rg --type python -g '*.py' -e '(pose_optimization_step|pose_refinement)' -C 10 # Test: Look for any performance profiling or logging related to Pose Optimization rg --type python -g '*.py' -e '(profile|log_performance).*pose.*optimization' -C 5 # Test: Check for any configurable parameters that might affect the time-accuracy trade-off rg --type python -g '*.py' -e '(pose_optimization_iterations|pose_accuracy_threshold)' -C 5
348-350:⚠️ Potential issueConsistent error reporting but increased execution time for
q_opt 2.The changes in
q_opt 2output are noteworthy:
- The error reporting now uses the same array format as
q_opt 1, providing consistency and detailed insights.- However, the execution time for
q_opt 2has increased significantly from 0.0062 seconds to 1.3736 seconds.While the consistent error reporting is a positive change, the substantial increase in execution time warrants investigation.
To investigate the cause of the increased execution time, please run the following script:
75-81:⚠️ Potential issueSignificant increase in compilation time detected.
The compilation time for the
jit_q_optmodule has increased from 27m43s to 30m42s. This substantial increase (approximately 3 minutes) may indicate changes in the optimization algorithm, increased model complexity, or a potential performance regression.To investigate this further, please run the following script:
90-101: Improved optimization results observed.The optimization process shows significant improvements:
- The final error has decreased from 0.00022817964782007039 to 5.8231067669112235e-05, which is a substantial reduction.
- The
q_optandReplaceoperation times have changed, potentially indicating algorithmic improvements.These changes suggest enhancements in the optimization process, possibly due to algorithm refinements or hyperparameter tuning.
To ensure these improvements are consistent and not a result of overfitting, please run the following validation script:
✅ Verification successful
Verification script encountered errors. Retrying with corrected commands.
Please execute the following scripts to continue the verification process:
Retrying verification with corrected
rgcommands.Please execute the following scripts to continue the verification process:
Further verification needed for optimization logging.
Please execute the following scripts to gather additional information regarding logging of optimization metrics:
Optimization logging verified.
The code includes logging statements related to optimization processes, such as:
logging.info(f"Running fit. Mocap data shape: {fit_data.shape}")logging.info(f"saving data to {fit_path}")logging.info("skipping ik_only()")logging.info("Running ik_only()")logging.info(f"kp_data shape: {kp_data.shape}")logging.info(f"Saving data to {ik_only_path}. Finished in {time.time() - start_time} seconds")logging.info(f"Run complete. \n fit path: {fit_path} \n ik_only path: {ik_only_path}")logging.info(f"cfg: {OmegaConf.to_yaml(cfg)}")These logs confirm that the optimization steps are being tracked and recorded appropriately.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Description: Check for consistency in optimization improvements across different runs or datasets # Test: Look for recent changes in optimization-related parameters or algorithms rg --type python -g '*.py' -e '(q_opt|optimization|hyperparameter)' -C 5 # Test: Check for any logging of optimization metrics across multiple runs rg --type python -g '*.py' -e '(log_metric|record_optimization)' -C 5Length of output: 376
Script:
#!/bin/bash # Description: Check for consistency in optimization improvements across different runs or datasets # Test: Look for recent changes in optimization-related parameters or algorithms rg --type=python -g '*.py' -e '(q_opt|optimization|hyperparameter)' -C 5 # Test: Check for any logging of optimization metrics across multiple runs rg --type=python -g '*.py' -e '(log_metric|record_optimization)' -C 5Length of output: 208
Script:
#!/bin/bash # Description: Check for consistency in optimization improvements across different runs or datasets # Test: Look for recent changes in optimization-related parameters or algorithms rg -g '*.py' -e '(q_opt|optimization|hyperparameter)' -C 5 # Test: Check for any logging of optimization metrics across multiple runs rg -g '*.py' -e '(log_metric|record_optimization)' -C 5Length of output: 16184
Script:
#!/bin/bash # Description: Enhanced search for logging of optimization metrics across the codebase # Test 1: Look for various logging patterns related to optimization rg -g '*.py' -e '(log_metric|record_optimization|logging\.info|logger\.info|log\.debug)' -C 5 # Test 2: Search for logging configurations that might handle optimization metrics rg -g '*.py' -e '(logging\.basicConfig|logging\.getLogger)' -C 5 # Test 3: Check other file types for potential logging implementations rg -e '(log_metric|record_optimization|logging\.info|logger\.info|log\.debug)' -C 5Length of output: 5290
There was a problem hiding this comment.
Can you clean this up for official review and submission? Three things:
- Check the one comment
- Can you please add a unit to catch this next time? Perhaps write the data out, and read it back in, etc..
- Revert/delete non-relevant files. (I included those for your ease of debugging.)
Thank you!
| @@ -390,7 +390,7 @@ def _package_data(self, mjx_model, q, x, walker_body_sites, kp_data, batched=Fal | |||
| x = x.reshape(-1, x.shape[-1]) | |||
| q = q.reshape(-1, q.shape[-1]) | |||
| else: | |||
| offsets = self._offsets | |||
| offsets = self._offsets.reshape((-1, 3)) | |||
There was a problem hiding this comment.
That's what I figured, but I wanted to make, since a few things were changed by the PR that introduced it.
Just to double check, the other branch of this doesn't need it as well?
Summary by CodeRabbit
New Features
Bug Fixes
Chores