Repository: Genesis-Embodied-AI/Genesis
File: genesis/engine/entities/rigid_entity/rigid_entity.py (Python)

Code before:

        scene_contact_info = self._solver.collider.contact_data.to_torch(gs.device)
        n_contacts = self._solver.collider.n_contacts.to_torch(gs.device)

        if exclude_self_contact:
            valid_mask = torch.logical_xor(
                torch.logical_and(
                    scene_contact_info["geom_a"] >= self.geom_start,
                    scene_contact_info["geom_a"] < self.geom_end,
                ),
                torch.logical_and(
                    scene_contact_info["geom_b"] >= self.geom_start,
                    scene_contact_info["geom_b"] < self.geom_end,
                ),
            )
        else:
            valid_mask = torch.logical_or(
                torch.logical_and(
                    scene_contact_info["geom_a"] >= self.geom_start,
                    scene_contact_info["geom_a"] < self.geom_end,
                ),
                torch.logical_and(
                    scene_contact_info["geom_b"] >= self.geom_start,
                    scene_contact_info["geom_b"] < self.geom_end,
                ),
            )
        if with_entity is not None:
            if self.idx == with_entity.idx:
                gs.raise_exception("`with_entity` cannot be the same as the caller entity.")
            valid_mask = torch.logical_and(
                valid_mask,
                torch.logical_or(
                    torch.logical_and(
                        scene_contact_info["geom_a"] >= with_entity.geom_start,
                        scene_contact_info["geom_a"] < with_entity.geom_end,
                    ),
                    torch.logical_and(
                        scene_contact_info["geom_b"] >= with_entity.geom_start,
                        scene_contact_info["geom_b"] < with_entity.geom_end,
                    ),
                ),
            )

        # Create an index tensor for contacts and compare against per-env n_contacts.
        # Assumes valid_mask.shape[0] corresponds to the maximum number of contacts.
        contact_indices = torch.arange(valid_mask.shape[0], device=valid_mask.device).unsqueeze(1)
        valid_mask = torch.logical_and(valid_mask, contact_indices < n_contacts)

        if self._solver.n_envs == 0:
            valid_idx = torch.nonzero(valid_mask.squeeze(), as_tuple=True)[0]
            contact_info = {

Review comment (suggestion):
to reduce repeated code, can you do?

```
logical_operation = torch.logical_xor if exclude_self_contact else torch.logical_or
valid_mask = logical_operation(...)
```

Code after:
        exclude_self_contact: bool
            Exclude the self collision from the returning contacts. Defaults to False.

        Returns
        -------
        contact_info : dict
            The contact information.
        """
        scene_contact_info = self._solver.collider.contact_data.to_torch(gs.device)
        n_contacts = self._solver.collider.n_contacts.to_torch(gs.device)

        logical_operation = torch.logical_xor if exclude_self_contact else torch.logical_or
        if with_entity is not None and self.idx == with_entity.idx:
            if exclude_self_contact:
                gs.raise_exception("`with_entity` is self but `exclude_self_contact` is True.")
            logical_operation = torch.logical_and

        valid_mask = logical_operation(
            torch.logical_and(
                scene_contact_info["geom_a"] >= self.geom_start,
                scene_contact_info["geom_a"] < self.geom_end,
            ),
            torch.logical_and(
                scene_contact_info["geom_b"] >= self.geom_start,
                scene_contact_info["geom_b"] < self.geom_end,
            ),
        )
        if with_entity is not None and self.idx != with_entity.idx:
            valid_mask = torch.logical_and(
                valid_mask,
                torch.logical_or(
                    torch.logical_and(
                        scene_contact_info["geom_a"] >= with_entity.geom_start,
                        scene_contact_info["geom_a"] < with_entity.geom_end,
                    ),
                    torch.logical_and(
                        scene_contact_info["geom_b"] >= with_entity.geom_start,
                        scene_contact_info["geom_b"] < with_entity.geom_end,
                    ),
                ),
            )

        # Create an index tensor for contacts and compare against per-env n_contacts.
        # Assumes valid_mask.shape[0] corresponds to the maximum number of contacts.
        contact_indices = torch.arange(valid_mask.shape[0], device=valid_mask.device).unsqueeze(1)
        valid_mask = torch.logical_and(valid_mask, contact_indices < n_contacts)

        if self._solver.n_envs == 0:
            valid_idx = torch.nonzero(valid_mask.squeeze(), as_tuple=True)[0]
            contact_info = {
                "geom_a": scene_contact_info["geom_a"][valid_idx, 0],