feat(phyai): WALL OSS Support#25
Conversation
📝 WalkthroughWalkthroughIntroduces the complete WALL-OSS-FLOW plugin for PhyAI: a frozen ChangesWALL-OSS-FLOW Plugin
Sequence Diagram(s)sequenceDiagram
participant CLI as CLI / Example Script
participant Engine
participant WallOSSEntry
participant WallOSSFlowWS1Scheduler
participant WallOSSFlowRunner
participant Qwen2_5_VLMoEForAction
CLI->>Engine: Engine(plugin="walloss", args=WallOSSArgs)
Engine->>WallOSSEntry: setup(args)
WallOSSEntry->>WallOSSFlowRunner: setup() — load weights, apply precision policy
WallOSSEntry->>WallOSSFlowWS1Scheduler: setup()
CLI->>Engine: engine.step(WallOSSFlowRequest | WallOSSFlowPredictRequest)
Engine->>WallOSSEntry: step(request)
WallOSSEntry->>WallOSSFlowWS1Scheduler: step(request)
WallOSSFlowWS1Scheduler->>WallOSSFlowWS1Scheduler: _validate() / _validate_predict()
WallOSSFlowWS1Scheduler->>WallOSSFlowWS1Scheduler: _to_wallx_batch() / _to_wallx_predict_batch()
WallOSSFlowWS1Scheduler->>WallOSSFlowRunner: forward(batch) or predict(batch, ...)
WallOSSFlowRunner->>Qwen2_5_VLMoEForAction: mode="validate" or mode="predict"
Qwen2_5_VLMoEForAction-->>WallOSSFlowRunner: logits / action tensor
WallOSSFlowRunner-->>CLI: output
CLI->>Engine: engine.close()
Engine->>WallOSSEntry: close()
WallOSSEntry->>WallOSSFlowWS1Scheduler: close() → runner.close()
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces the first-version WALL-OSS-FLOW plugin for the PhyAI engine, including end-to-end benchmark examples, configuration helpers, a model runner wrapper, and a single-GPU scheduler. The reviewer feedback focuses on improving robustness and performance: supporting sharded checkpoints, validating CLI arguments in the dry-run script, verifying the loaded JSON config structure, moving normalizer setup out of the inference hot path to reduce overhead, and explicitly casting pixel values to the target dtype to prevent runtime mismatches.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| if not (self.checkpoint_dir / "model.safetensors").is_file(): | ||
| raise FileNotFoundError( | ||
| f"missing model.safetensors under checkpoint_dir: {self.checkpoint_dir}" | ||
| ) |
There was a problem hiding this comment.
Strictly checking for a single model.safetensors file will fail for sharded checkpoints (which use model.safetensors.index.json and multiple shard files) or checkpoints saved in .bin format. Since MoE models are typically very large, they are almost always sharded. Checking for alternative weight file patterns makes the runner much more robust.
| if not (self.checkpoint_dir / "model.safetensors").is_file(): | |
| raise FileNotFoundError( | |
| f"missing model.safetensors under checkpoint_dir: {self.checkpoint_dir}" | |
| ) | |
| has_weights = ( | |
| (self.checkpoint_dir / "model.safetensors").is_file() | |
| or (self.checkpoint_dir / "model.safetensors.index.json").is_file() | |
| or (self.checkpoint_dir / "pytorch_model.bin").is_file() | |
| or (self.checkpoint_dir / "pytorch_model.bin.index.json").is_file() | |
| ) | |
| if not has_weights: | |
| raise FileNotFoundError( | |
| f"missing model weights (model.safetensors or pytorch_model.bin) under checkpoint_dir: {self.checkpoint_dir}" | |
| ) |
| args = parser.parse_args() | ||
|
|
There was a problem hiding this comment.
The dry-run script is missing validation for several CLI arguments. If --checkpoint or --wall-x-root are not directories, or if --n-timed is non-positive, the script will fail later with unhelpful errors (such as StatisticsError or ValueError from min()). Adding explicit validation here ensures robust error handling.
| args = parser.parse_args() | |
| args = parser.parse_args() | |
| if not args.checkpoint.is_dir(): | |
| raise NotADirectoryError(f"--checkpoint must be a directory: {args.checkpoint}") | |
| if not args.wall_x_root.is_dir(): | |
| raise NotADirectoryError(f"--wall-x-root must be a directory: {args.wall_x_root}") | |
| if args.n_warmup < 0: | |
| raise ValueError(f"--n-warmup must be non-negative, got {args.n_warmup}.") | |
| if args.n_timed <= 0: | |
| raise ValueError(f"--n-timed must be positive, got {args.n_timed}.") |
| cfg = json.loads(cfg_path.read_text()) | ||
|
|
There was a problem hiding this comment.
It is safer to verify that the loaded JSON config is a dictionary before calling .get() on it. If config.json is malformed or contains a top-level list, this prevents an unhelpful AttributeError.
| cfg = json.loads(cfg_path.read_text()) | |
| cfg = json.loads(cfg_path.read_text()) | |
| if not isinstance(cfg, dict): | |
| raise ValueError(f"config.json must be a JSON object, got {type(cfg).__name__}") |
| model = model.to(self.device) | ||
| model = self._apply_precision_policy(model) | ||
| model.eval() | ||
| self.model = model |
There was a problem hiding this comment.
Calling setup_default_normalizers() lazily inside the predict() hot path introduces unnecessary Python overhead (multiple getattr calls) on every single inference step. Furthermore, calling it after _apply_precision_policy means the normalizers are not automatically cast to the model's target precision (e.g., bfloat16), potentially leading to dtype mismatches or implicit casting overhead. Moving it to setup() resolves both issues.
| model = model.to(self.device) | |
| model = self._apply_precision_policy(model) | |
| model.eval() | |
| self.model = model | |
| model = model.to(self.device) | |
| self.model = model | |
| self.setup_default_normalizers() | |
| self.model = self._apply_precision_policy(self.model) | |
| self.model.eval() |
| model = self._require_model() | ||
| self.setup_default_normalizers() | ||
| return model( |
There was a problem hiding this comment.
Remove the lazy call to self.setup_default_normalizers() from the predict() hot path now that it is handled during setup().
model = self._require_model()
return model(| "input_ids": req.input_ids.to(device=self.device, dtype=torch.long), | ||
| "attention_mask": req.attention_mask.to(device=self.device, dtype=torch.long), | ||
| "moe_token_types": req.moe_token_types.to(device=self.device, dtype=torch.bool), | ||
| "pixel_values": req.pixel_values.to(device=self.device), |
There was a problem hiding this comment.
In _to_wallx_predict_batch, pixel_values is moved to self.device but its dtype is not explicitly cast to self.dtype. If pixel_values is in float32 but the vision tower expects self.dtype (e.g., bfloat16), this can cause a runtime dtype mismatch error or implicit casting overhead. Casting it to self.dtype ensures consistency.
| "pixel_values": req.pixel_values.to(device=self.device), | |
| "pixel_values": req.pixel_values.to(device=self.device, dtype=self.dtype), |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 8
🧹 Nitpick comments (1)
examples/run_walloss_flow_predict_dryrun.py (1)
73-74: ⚡ Quick winRefactor: Move
sys.pathmodification out ofmake_predict_request.Lines 73-74 modify global
sys.pathstate inside a "make request" function, which is a surprising side effect. The same modification is also done inmain()at lines 177-178, creating redundancy.♻️ Proposed refactor
Remove lines 73-74 from
make_predict_requestand rely on themain()function to set upsys.pathonce before callingmake_predict_request:def make_predict_request( *, checkpoint: Path, wall_x_root: Path, device: torch.device, ) -> WallOSSFlowPredictRequest: - if str(wall_x_root) not in sys.path: - sys.path.insert(0, str(wall_x_root)) - from wall_x.model.action_head import NormalizerThis makes
make_predict_requesta pure function without global side effects and eliminates the duplication.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/run_walloss_flow_predict_dryrun.py` around lines 73 - 74, Remove the sys.path modification (the lines that check if wall_x_root is in sys.path and insert it) from the make_predict_request function since the same modification is already being performed in the main function before make_predict_request is called. This eliminates the redundant global state modification inside make_predict_request and makes it a pure function without side effects. The sys.path setup should be handled once by main as part of its initialization responsibility.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@examples/run_walloss_flow_predict_dryrun.py`:
- Line 84: The AutoProcessor.from_pretrained() call uses trust_remote_code=True
which creates a remote code execution vulnerability if the checkpoint source is
untrusted. Since this is an example/test script, either add a warning comment
above the line documenting the trust assumption about the checkpoint source and
explaining why the flag is necessary, or attempt to remove
trust_remote_code=True entirely and test if the processor can be loaded without
it. If you keep the flag, include a comment explaining that this example assumes
the checkpoint comes from a trusted source and documenting why arbitrary code
execution is acceptable in this context.
- Around line 170-175: The script accepts checkpoint and wall_x_root as
command-line arguments but does not validate that these paths exist or have the
correct type before using them downstream, which leads to unclear errors. After
the parser.parse_args() call, add validation logic to ensure that
args.checkpoint exists as a directory and args.wall_x_root exists as a
directory, raising a clear error with helpful context if either validation
fails.
- Around line 180-184: Add a check for CUDA availability before using
CUDA-specific operations in the device initialization block. Before calling
torch.cuda.manual_seed_all(0), verify that torch.cuda.is_available() returns
True, and only execute the CUDA seeding if available. Additionally, validate
that the device specified in args.device (which defaults to "cuda") is actually
accessible on the system before creating the torch.device object, either by
checking availability upfront or falling back to CPU if the requested device is
unavailable.
In `@examples/run_walloss_flow.py`:
- Around line 133-137: Add a check for CUDA availability before using
CUDA-specific operations. Before the torch.manual_seed and
torch.cuda.manual_seed_all calls, verify torch.cuda.is_available() and
conditionally execute torch.cuda.manual_seed_all(0) only when CUDA is available.
Additionally, use the result of this check to set the device variable to either
"cuda" or "cpu" appropriately, so that the device assignment on the line with
torch.device works correctly on both GPU and CPU-only systems.
- Around line 115-116: The `--n-warmup` argument parser definition lacks
validation to ensure it is a positive integer, while other similar arguments
like `batch_size`, `seq_len`, and `n_timed` are validated. Add validation logic
for the `n_warmup` argument in the argument parsing section to check that its
value is positive (greater than 0), consistent with the validation pattern used
for the other numeric arguments in the script.
- Line 93: The assignment `logits = outputs.logits` on line 93 assumes the
outputs object has a `.logits` attribute without validation, which will fail
with AttributeError if the model returns a dict, tuple, or object without this
attribute. Add defensive extraction logic before this line that validates the
output structure and safely retrieves the logits, handling different possible
return formats (such as checking if outputs is a dict with a logits key, or if
it's an object with a logits attribute, or extracting from a tuple/list by index
if applicable based on the model's expected behavior in validate mode).
In `@phyai/src/phyai/models/walloss/configuration_walloss.py`:
- Around line 15-21: The _sum_dim function and the from_checkpoint method
currently accept zero and negative values for dimensions and action_horizon
respectively, which causes failures later with less actionable error messages.
Add validation in _sum_dim after computing the sum to check that the result is
positive (greater than 0), raising a ValueError with a clear message if not.
Similarly, add validation in the from_checkpoint method to ensure the
action_horizon parameter is positive before it is used, raising a ValueError
with appropriate context if validation fails.
In `@phyai/src/phyai/models/walloss/scheduler_ws1_walloss.py`:
- Around line 191-218: The validation block is missing two checks: first, add
validation to ensure that pixel_values.shape[0] matches the total number of
patches implied by image_grid_thw (computed as the product of the three values
in each image_grid_thw row). Second, replace the use of the `or` operator when
assigning action_dim and action_horizon, as this treats an explicit value of 0
as falsy and falls back to the config value instead of validating it as an
explicit override. Change these assignments to use explicit None checks (e.g.,
if req.action_dim is not None else cfg.action_dim) to properly distinguish
between an unset parameter and an explicitly provided value, including 0.
---
Nitpick comments:
In `@examples/run_walloss_flow_predict_dryrun.py`:
- Around line 73-74: Remove the sys.path modification (the lines that check if
wall_x_root is in sys.path and insert it) from the make_predict_request function
since the same modification is already being performed in the main function
before make_predict_request is called. This eliminates the redundant global
state modification inside make_predict_request and makes it a pure function
without side effects. The sys.path setup should be handled once by main as part
of its initialization responsibility.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: a0e3ae93-92aa-4143-ac0f-00cdfa736a4b
📒 Files selected for processing (8)
examples/run_walloss_flow.pyexamples/run_walloss_flow_predict_dryrun.pyphyai/src/phyai/engine.pyphyai/src/phyai/models/walloss/__init__.pyphyai/src/phyai/models/walloss/configuration_walloss.pyphyai/src/phyai/models/walloss/main_walloss.pyphyai/src/phyai/models/walloss/model_runner_walloss.pyphyai/src/phyai/models/walloss/scheduler_ws1_walloss.py
| cfg = WallOSSFlowConfig.from_checkpoint(checkpoint) | ||
| raw_cfg = json.loads((checkpoint / "config.json").read_text()) | ||
|
|
||
| processor = AutoProcessor.from_pretrained(str(checkpoint), trust_remote_code=True) |
There was a problem hiding this comment.
Security risk: trust_remote_code=True allows arbitrary code execution.
Line 84 loads the processor with trust_remote_code=True, which permits executing arbitrary Python code from the checkpoint. If the checkpoint comes from an untrusted source, this creates a remote code execution (RCE) vulnerability.
Since this is an example/test script (not production), consider adding a warning comment and documenting the trust assumption:
📋 Suggested mitigation
+ # SECURITY NOTE: trust_remote_code=True allows checkpoint code execution.
+ # Only use with trusted checkpoints from verified sources.
processor = AutoProcessor.from_pretrained(str(checkpoint), trust_remote_code=True)Alternatively, if the processor can be loaded without trust_remote_code=True, prefer that approach.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/run_walloss_flow_predict_dryrun.py` at line 84, The
AutoProcessor.from_pretrained() call uses trust_remote_code=True which creates a
remote code execution vulnerability if the checkpoint source is untrusted. Since
this is an example/test script, either add a warning comment above the line
documenting the trust assumption about the checkpoint source and explaining why
the flag is necessary, or attempt to remove trust_remote_code=True entirely and
test if the processor can be loaded without it. If you keep the flag, include a
comment explaining that this example assumes the checkpoint comes from a trusted
source and documenting why arbitrary code execution is acceptable in this
context.
| parser.add_argument("--checkpoint", type=Path, required=True) | ||
| parser.add_argument("--wall-x-root", type=Path, default=Path("/phyai_workspace/src/wall-x")) | ||
| parser.add_argument("--device", type=str, default="cuda") | ||
| parser.add_argument("--n-warmup", type=int, default=0) | ||
| parser.add_argument("--n-timed", type=int, default=1) | ||
| args = parser.parse_args() |
There was a problem hiding this comment.
Missing path validation: checkpoint and wall-x-root not validated.
The script does not validate that args.checkpoint is a directory or that args.wall_x_root exists before attempting to use them. This will lead to unclear errors downstream (e.g., file not found when loading config.json or modules).
✅ Proposed fix
After line 175, add:
args = parser.parse_args()
+
+ if not args.checkpoint.is_dir():
+ raise NotADirectoryError(f"--checkpoint must be a directory: {args.checkpoint}")
+ if not args.wall_x_root.is_dir():
+ raise NotADirectoryError(f"--wall-x-root must be a directory: {args.wall_x_root}")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/run_walloss_flow_predict_dryrun.py` around lines 170 - 175, The
script accepts checkpoint and wall_x_root as command-line arguments but does not
validate that these paths exist or have the correct type before using them
downstream, which leads to unclear errors. After the parser.parse_args() call,
add validation logic to ensure that args.checkpoint exists as a directory and
args.wall_x_root exists as a directory, raising a clear error with helpful
context if either validation fails.
| torch.manual_seed(0) | ||
| torch.cuda.manual_seed_all(0) | ||
|
|
||
| device = torch.device(args.device) | ||
| dtype = torch.bfloat16 |
There was a problem hiding this comment.
Missing CUDA availability check.
Lines 181 and 183 use torch.cuda.manual_seed_all and torch.device(args.device) (defaulting to "cuda") without checking torch.cuda.is_available(). On CPU-only systems or when --device cuda is specified without GPU support, this will fail with unclear errors.
🔧 Proposed fix
torch.manual_seed(0)
- torch.cuda.manual_seed_all(0)
-
device = torch.device(args.device)
+ if device.type == "cuda":
+ if not torch.cuda.is_available():
+ raise RuntimeError("CUDA device requested but CUDA is not available.")
+ torch.cuda.manual_seed_all(0)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| torch.manual_seed(0) | |
| torch.cuda.manual_seed_all(0) | |
| device = torch.device(args.device) | |
| dtype = torch.bfloat16 | |
| torch.manual_seed(0) | |
| device = torch.device(args.device) | |
| if device.type == "cuda": | |
| if not torch.cuda.is_available(): | |
| raise RuntimeError("CUDA device requested but CUDA is not available.") | |
| torch.cuda.manual_seed_all(0) | |
| dtype = torch.bfloat16 |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/run_walloss_flow_predict_dryrun.py` around lines 180 - 184, Add a
check for CUDA availability before using CUDA-specific operations in the device
initialization block. Before calling torch.cuda.manual_seed_all(0), verify that
torch.cuda.is_available() returns True, and only execute the CUDA seeding if
available. Additionally, validate that the device specified in args.device
(which defaults to "cuda") is actually accessible on the system before creating
the torch.device object, either by checking availability upfront or falling back
to CPU if the requested device is unavailable.
| if outputs is None: | ||
| raise RuntimeError("benchmark did not run any engine.step call.") | ||
|
|
||
| logits = outputs.logits |
There was a problem hiding this comment.
Missing validation: outputs.logits access may fail.
Line 93 accesses outputs.logits without validating the output structure. If the model's mode="validate" returns a dict, tuple, or object without a .logits attribute, this will raise AttributeError.
🛡️ Proposed fix with defensive extraction
- logits = outputs.logits
+ if isinstance(outputs, dict):
+ logits = outputs["logits"]
+ elif hasattr(outputs, "logits"):
+ logits = outputs.logits
+ else:
+ raise RuntimeError(f"cannot extract logits from output type={type(outputs)}")
return logits, {📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| logits = outputs.logits | |
| if isinstance(outputs, dict): | |
| logits = outputs["logits"] | |
| elif hasattr(outputs, "logits"): | |
| logits = outputs.logits | |
| else: | |
| raise RuntimeError(f"cannot extract logits from output type={type(outputs)}") | |
| return logits, { |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/run_walloss_flow.py` at line 93, The assignment `logits =
outputs.logits` on line 93 assumes the outputs object has a `.logits` attribute
without validation, which will fail with AttributeError if the model returns a
dict, tuple, or object without this attribute. Add defensive extraction logic
before this line that validates the output structure and safely retrieves the
logits, handling different possible return formats (such as checking if outputs
is a dict with a logits key, or if it's an object with a logits attribute, or
extracting from a tuple/list by index if applicable based on the model's
expected behavior in validate mode).
| parser.add_argument("--n-warmup", type=int, default=1) | ||
| parser.add_argument("--n-timed", type=int, default=3) |
There was a problem hiding this comment.
Missing validation: --n-warmup can be negative.
The script validates that batch_size, seq_len, and n_timed are positive but does not validate n_warmup. Negative warmup iterations are semantically incorrect.
✅ Proposed fix
parser.add_argument("--n-warmup", type=int, default=1)
parser.add_argument("--n-timed", type=int, default=3)After line 131, add:
if args.n_timed <= 0:
raise ValueError(f"--n-timed must be positive, got {args.n_timed}.")
+ if args.n_warmup < 0:
+ raise ValueError(f"--n-warmup must be non-negative, got {args.n_warmup}.")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/run_walloss_flow.py` around lines 115 - 116, The `--n-warmup`
argument parser definition lacks validation to ensure it is a positive integer,
while other similar arguments like `batch_size`, `seq_len`, and `n_timed` are
validated. Add validation logic for the `n_warmup` argument in the argument
parsing section to check that its value is positive (greater than 0), consistent
with the validation pattern used for the other numeric arguments in the script.
| torch.manual_seed(0) | ||
| torch.cuda.manual_seed_all(0) | ||
|
|
||
| device = torch.device("cuda") | ||
| dtype = torch.bfloat16 |
There was a problem hiding this comment.
Missing CUDA availability check.
Lines 134 and 136 use torch.cuda.manual_seed_all and torch.device("cuda") without first checking torch.cuda.is_available(). On CPU-only systems, this will fail with unclear errors.
🔧 Proposed fix
torch.manual_seed(0)
+ if not torch.cuda.is_available():
+ raise RuntimeError("CUDA is required but not available on this system.")
torch.cuda.manual_seed_all(0)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| torch.manual_seed(0) | |
| torch.cuda.manual_seed_all(0) | |
| device = torch.device("cuda") | |
| dtype = torch.bfloat16 | |
| torch.manual_seed(0) | |
| if not torch.cuda.is_available(): | |
| raise RuntimeError("CUDA is required but not available on this system.") | |
| torch.cuda.manual_seed_all(0) | |
| device = torch.device("cuda") | |
| dtype = torch.bfloat16 |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/run_walloss_flow.py` around lines 133 - 137, Add a check for CUDA
availability before using CUDA-specific operations. Before the torch.manual_seed
and torch.cuda.manual_seed_all calls, verify torch.cuda.is_available() and
conditionally execute torch.cuda.manual_seed_all(0) only when CUDA is available.
Additionally, use the result of this check to set the device variable to either
"cuda" or "cpu" appropriately, so that the device assignment on the line with
torch.device works correctly on both GPU and CPU-only systems.
| def _sum_dim(config: dict[str, int] | None, *, name: str) -> int: | ||
| if not isinstance(config, dict) or not config: | ||
| raise ValueError(f"{name} must be a non-empty dict, got {type(config).__name__}.") | ||
| try: | ||
| return sum(int(v) for v in config.values()) | ||
| except Exception as exc: | ||
| raise ValueError(f"failed to sum {name}: {config!r}") from exc |
There was a problem hiding this comment.
Validate positive dimensions and horizon at config-load time.
_sum_dim currently accepts zero/negative values, and from_checkpoint accepts non-positive action_horizon. Those invalid values flow into scheduler contracts and fail later with less actionable errors.
Proposed fix
def _sum_dim(config: dict[str, int] | None, *, name: str) -> int:
if not isinstance(config, dict) or not config:
raise ValueError(f"{name} must be a non-empty dict, got {type(config).__name__}.")
try:
- return sum(int(v) for v in config.values())
+ total = 0
+ for key, value in config.items():
+ dim = int(value)
+ if dim <= 0:
+ raise ValueError(f"{name}.{key} must be > 0, got {value!r}.")
+ total += dim
+ return total
except Exception as exc:
raise ValueError(f"failed to sum {name}: {config!r}") from exc
@@
- return cls(
+ horizon = int(action_horizon)
+ if horizon <= 0:
+ raise ValueError(f"action_horizon must be > 0, got {action_horizon!r}.")
+
+ return cls(
checkpoint_dir=checkpoint_dir,
model_type=model_type,
vocab_size=int(cfg["vocab_size"]),
action_dim=action_dim,
proprio_dim=proprio_dim,
- action_horizon=int(action_horizon),
+ action_horizon=horizon,
dataset_name=str(dataset_name),
)Also applies to: 67-75
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@phyai/src/phyai/models/walloss/configuration_walloss.py` around lines 15 -
21, The _sum_dim function and the from_checkpoint method currently accept zero
and negative values for dimensions and action_horizon respectively, which causes
failures later with less actionable error messages. Add validation in _sum_dim
after computing the sum to check that the result is positive (greater than 0),
raising a ValueError with a clear message if not. Similarly, add validation in
the from_checkpoint method to ensure the action_horizon parameter is positive
before it is used, raising a ValueError with appropriate context if validation
fails.
| if req.pixel_values.ndim != 2: | ||
| raise ValueError(f"pixel_values must be 2-D, got shape {tuple(req.pixel_values.shape)}.") | ||
|
|
||
| if req.image_grid_thw.ndim != 2 or req.image_grid_thw.shape[-1] != 3: | ||
| raise ValueError( | ||
| "image_grid_thw must have shape (num_images, 3), " | ||
| f"got {tuple(req.image_grid_thw.shape)}." | ||
| ) | ||
|
|
||
| if tuple(req.proprioception.shape) != (B, 1, cfg.proprio_dim): | ||
| raise ValueError( | ||
| f"proprioception shape {tuple(req.proprioception.shape)} != " | ||
| f"({B}, 1, {cfg.proprio_dim})." | ||
| ) | ||
|
|
||
| if tuple(req.agent_pos_mask.shape) != (B, 1, cfg.proprio_dim): | ||
| raise ValueError( | ||
| f"agent_pos_mask shape {tuple(req.agent_pos_mask.shape)} != " | ||
| f"({B}, 1, {cfg.proprio_dim})." | ||
| ) | ||
|
|
||
| action_dim = req.action_dim or cfg.action_dim | ||
| action_horizon = req.action_horizon or cfg.action_horizon | ||
| if tuple(req.dof_mask.shape) != (B, action_horizon, action_dim): | ||
| raise ValueError( | ||
| f"dof_mask shape {tuple(req.dof_mask.shape)} != " | ||
| f"({B}, {action_horizon}, {action_dim})." | ||
| ) |
There was a problem hiding this comment.
Tighten predict-request validation for grid/token consistency and explicit overrides.
The predict path validates rank, but it does not verify that pixel_values.shape[0] matches the total patches implied by image_grid_thw. Also, using or for action_dim/action_horizon treats explicit 0 as “unset” instead of invalid input.
Proposed fix
if req.image_grid_thw.ndim != 2 or req.image_grid_thw.shape[-1] != 3:
raise ValueError(
"image_grid_thw must have shape (num_images, 3), "
f"got {tuple(req.image_grid_thw.shape)}."
)
+ expected_tokens = int(req.image_grid_thw.to(dtype=torch.long).prod(dim=1).sum().item())
+ if req.pixel_values.shape[0] != expected_tokens:
+ raise ValueError(
+ f"pixel_values first dim {req.pixel_values.shape[0]} != "
+ f"sum(prod(image_grid_thw)) {expected_tokens}."
+ )
@@
- action_dim = req.action_dim or cfg.action_dim
- action_horizon = req.action_horizon or cfg.action_horizon
+ action_dim = cfg.action_dim if req.action_dim is None else int(req.action_dim)
+ action_horizon = (
+ cfg.action_horizon if req.action_horizon is None else int(req.action_horizon)
+ )
+ if action_dim <= 0 or action_horizon <= 0:
+ raise ValueError(
+ f"action_dim/action_horizon must be > 0, got {action_dim}/{action_horizon}."
+ )
if tuple(req.dof_mask.shape) != (B, action_horizon, action_dim):
raise ValueError(
f"dof_mask shape {tuple(req.dof_mask.shape)} != "
f"({B}, {action_horizon}, {action_dim})."
)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@phyai/src/phyai/models/walloss/scheduler_ws1_walloss.py` around lines 191 -
218, The validation block is missing two checks: first, add validation to ensure
that pixel_values.shape[0] matches the total number of patches implied by
image_grid_thw (computed as the product of the three values in each
image_grid_thw row). Second, replace the use of the `or` operator when assigning
action_dim and action_horizon, as this treats an explicit value of 0 as falsy
and falls back to the config value instead of validating it as an explicit
override. Change these assignments to use explicit None checks (e.g., if
req.action_dim is not None else cfg.action_dim) to properly distinguish between
an unset parameter and an explicitly provided value, including 0.
chenghuaWang
changed the title
1 participant
Summary
This PR adds the first PhyAI Engine plugin integration for WALL-OSS-FLOW.
The implementation follows the existing PhyAI Engine/plugin pattern and registers a new
wallossplugin. It includes a dedicated configuration wrapper, model runner, single-card scheduler, request object, and reproducible dry-run examples.The first version keeps the official WALL-OSS-FLOW model path isolated inside the new plugin and does not modify the PI0.5 plugin or other existing model plugins.
Changes
Added
phyai.models.wallossWallOSSArgsWallOSSEntryWallOSSConfigWallOSSRunnerWallOSSWS1SchedulerAdded examples:
examples/run_walloss_flow.pyexamples/run_walloss_flow_predict_dryrun.pyRegistered the
wallossplugin inphyai.engineVerification
The plugin was verified locally with the WALL-OSS-FLOW checkpoint.
Validation / smoke test:
(1, 50, 153715)torch.bfloat16cuda:0Policy prediction dry run:
(1, 32, 20)torch.float32cuda:0Official WALL-X vs PhyAI Engine fixed-seed comparison:
max_abs_diff:0.0torch.equal:TrueBenchmark
Local benchmark results:
57.22 ms711.79 msScope
This PR only adds the WALL-OSS-FLOW Engine plugin. WALL-OSS-0.5 is kept in a separate follow-up branch / PR so that the two integrations can be reviewed independently.
Summary by CodeRabbit
Release Notes
New Features