Revert reward size validation to unblock OC cogwheel test

Summary:
D76681825 added a new validation logic for reward tensor size hence broke the OC [cogwheel test](https://www.internalfb.com/conveyor/aps/ads_rl/releases/461.1/nodes/cogwheel_aps_pearl_ads_bidding_oc_on_mast/runs/1623320241869408?pipeline_names[0]=auto&node_details_tab=run_details_summary). [Error](https://www.internalfb.com/mlhub/pipelines/runs/mast/aps-cogwheel_pearl_ads_bidding_oc_896952826-1bc277ccc6?job_attempt=0&tab=summary)
To unblock this, revert this check to allow muti-dimensions for reward. We will re-evaluate this check once the diff owner D76681825 rodrigodesalvobraz is back.

Reviewed By: zhongjiema

Differential Revision: D77566142

fbshipit-source-id: 0bea1ac142f1e7e951c1c36bb14e5d602ab52c98

Author: Weixuan Kan

pearl/replay_buffers/transition.py

@@ -153,14 +153,10 @@ def __post_init__(self) -> None:
             self.action.ndim >= 1
         ), f"action has shape {self.action.shape}, but must be (batch_size,) or (batch_size, ...)"
 
-        # Allow reward to have shape (batch_size,) or (batch_size, 1)
-        is_1d = self.reward.ndim == 1
-        is_2d_with_1_col = self.reward.ndim == 2 and self.reward.shape[1] == 1
-        valid_reward_shape = is_1d or is_2d_with_1_col
-        assert valid_reward_shape, (
-            f"reward has shape {self.reward.shape}, "
-            f"but must be (batch_size,) or (batch_size, 1)"
-        )
+        # Allow reward to have shape (batch_size,) or (batch_size, ...)
+        assert (
+            self.reward.ndim >= 1
+        ), f"reward has shape {self.reward.shape}, but must be (batch_size,) or (batch_size, ...)"
 
         assert self.state.shape[0] == self.reward.shape[0], (
             f"state has shape {self.state.shape}, "