def compute_reward()

in gym_genesis/tasks/cube_stack.py [0:0]

• 7 lines of code
• 1 McCabe index (conditional complexity)

    def compute_reward(self):
        pos_1 = self.cube_1.get_pos()  # (B, 3)
        pos_2 = self.cube_2.get_pos()  # (B, 3)

        xy_dist = torch.norm(pos_1[:, :2] - pos_2[:, :2], dim=1)  # (B,)
        z_diff = pos_1[:, 2] - pos_2[:, 2]  # (B,)

        reward = ((xy_dist < 0.05) & (z_diff > 0.03)).float()  # (B,)
        return reward.cpu().numpy()