Supports vectorized environments for pick and place demo

[kellyguo11]

Description

Adds support for vectorizing the pick and place demo to test performance for multi-environments.

Type of change

• New feature (non-breaking change which adds functionality)

Checklist

• I have read and understood the contribution guidelines
• I have run the pre-commit checks with ./isaaclab.sh --format
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• I have updated the changelog and the corresponding version in the extension's config/extension.toml file
• I have added my name to the CONTRIBUTORS.md or my name already exists there

[greptile-apps]

Greptile Overview

Greptile Summary

This PR successfully vectorizes the pick-and-place demo to support multiple environments simultaneously. The implementation converts scalar control flags to per-environment boolean tensors, updates keyboard event handlers to broadcast commands to all environments, and properly indexes tensors for selective environment updates.

Key changes:

• Adds --num_envs CLI argument (default: 32) for configuring environment count
• Converts go_to_cube and go_to_target from single booleans to per-environment boolean tensors
• Updates _apply_action() to use if statements (not elif) for independent processing of cube/target tracking modes
• Implements tensor masking (self.go_to_cube, self.go_to_target) to selectively update only relevant environments
• Fixes gripper command shape by removing unnecessary unsqueeze(dim=1)
• Adds simulation config parameters (use_fabric=True, enable_scene_query_support=True) required for multi-environment surface grippers

The previous comment correctly identified the elif → if issue which has been resolved.

Confidence Score: 4/5

• Safe to merge with minor validation recommended - vectorization logic is sound
• The implementation correctly handles tensor operations for vectorized environments. The elif → if fix ensures independent processing of control modes. Tensor indexing with boolean masks is properly implemented. Minor concern: no runtime testing mentioned in checklist, though code structure appears correct.
• No files require special attention - the single file changed has straightforward vectorization updates

Important Files Changed

File Analysis

Filename	Score	Overview
scripts/demos/pick_and_place.py	4/5	Adds vectorized environment support with proper tensor indexing and if statements (was elif) for independent environment control

Sequence Diagram

sequenceDiagram
    participant User
    participant KeyboardHandler
    participant PickAndPlaceEnv
    participant Robot[N]
    participant Cube[N]
    participant Gripper[N]
    
    User->>KeyboardHandler: Press key (A/D/W/S/Q/E)
    KeyboardHandler->>PickAndPlaceEnv: _on_keyboard_event()
    PickAndPlaceEnv->>PickAndPlaceEnv: Update control flags for ALL envs<br/>(go_to_cube[:] or go_to_target[:])
    
    loop Simulation Step
        PickAndPlaceEnv->>PickAndPlaceEnv: _apply_action()
        
        alt go_to_cube.any()
            PickAndPlaceEnv->>Robot[N]: Get head positions [go_to_cube]
            PickAndPlaceEnv->>Cube[N]: Get cube positions [go_to_cube]
            PickAndPlaceEnv->>PickAndPlaceEnv: Calculate proportional control<br/>for selected envs
            PickAndPlaceEnv->>PickAndPlaceEnv: instant_controls[go_to_cube] = ...
        end
        
        alt go_to_target.any()
            PickAndPlaceEnv->>Robot[N]: Get head positions [go_to_target]
            PickAndPlaceEnv->>PickAndPlaceEnv: Calculate proportional control<br/>for selected envs
            PickAndPlaceEnv->>PickAndPlaceEnv: instant_controls[go_to_target] = ...
        end
        
        PickAndPlaceEnv->>Robot[N]: set_joint_effort_target(x, y, z)<br/>for ALL envs
        PickAndPlaceEnv->>Gripper[N]: set_grippers_command()<br/>for ALL envs
    end

Loading

[greptile-apps]

_{1 file reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{1 file reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[ooctipus]

LGTM!

[ooctipus]

maybe false?