New get_state_set method

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "rlviser-py"
-version = "0.6.1"
+version = "0.6.2"
 edition = "2021"
 description = "Python implementation that manages a UDP connection to RLViser"
 license = "MIT"

gym_renderer.py

@@ -2,7 +2,6 @@
 
 import rlviser_py as rlviser
 import RocketSim as rsim
-
 from rlgym.api import Renderer
 from rlgym.rocket_league.api import Car, GameState
 from rlgym.rocket_league.common_values import BOOST_LOCATIONS
@@ -22,6 +21,7 @@ def render(self, state: GameState, shared_info: Dict[str, Any]) -> Any:
  ball.pos = rsim.Vec(*state.ball.position)
  ball.vel = rsim.Vec(*state.ball.linear_velocity)
  ball.ang_vel = rsim.Vec(*state.ball.angular_velocity)
+ ball.rot_mat = rsim.RotMat(*state.ball.rotation_mtx.transpose().flatten())
 
  car_data = []
  for idx, car in enumerate(state.cars.values()):

pygymtest.py

@@ -12,49 +12,50 @@
 
 from gym_renderer import RLViserRenderer
 
-env = RLGym(
- state_mutator=MutatorSequence(
- FixedTeamSizeMutator(blue_size=2, orange_size=2),
- KickoffMutator()
- ),
- obs_builder=DefaultObs(zero_padding=None),
- action_parser=RepeatAction(LookupTableAction(), repeats=8),
- reward_fn=CombinedReward(
- (GoalReward(), 10.),
- (TouchReward(), 0.1)
- ),
- termination_cond=GoalCondition(),
- truncation_cond=AnyCondition(
- TimeoutCondition(timeout=300.),
- NoTouchTimeoutCondition(timeout=30.)
- ),
- transition_engine=RocketSimEngine(),
- renderer=RLViserRenderer()
-)
-
-# simulate 2 episodes
-for _ in range(2):
- obs_dict = env.reset()
- steps = 0
- ep_reward = {agent_id: 0. for agent_id in env.agents}
- t0 = time.time()
- while True:
- env.render()
-
- actions = {}
- for agent_id, action_space in env.action_spaces.items():
- # agent.act(obs) | Your agent should go here
- actions[agent_id] = np.random.randint(action_space, size=(1,))
-
- obs_dict, reward_dict, terminated_dict, truncated_dict = env.step(actions)
- for agent_id, reward in reward_dict.items():
- ep_reward[agent_id] += reward
-
- if any(chain(terminated_dict.values(), truncated_dict.values())):
- break
-
- time.sleep(max(0, t0 + steps / 15 - time.time()))
- steps += 1
-
- ep_time = time.time() - t0
- print(f"Steps per second: {steps / ep_time:.0f} | Episode time: {ep_time:.2f} | Episode Reward: {max(ep_reward.values()):.2f}")
+if __name__ == "__main__":
+ env = RLGym(
+ state_mutator=MutatorSequence(
+ FixedTeamSizeMutator(blue_size=2, orange_size=2),
+ KickoffMutator()
+ ),
+ obs_builder=DefaultObs(zero_padding=None),
+ action_parser=RepeatAction(LookupTableAction(), repeats=1),
+ reward_fn=CombinedReward(
+ (GoalReward(), 10.),
+ (TouchReward(), 0.1)
+ ),
+ termination_cond=GoalCondition(),
+ truncation_cond=AnyCondition(
+ TimeoutCondition(timeout=300.),
+ NoTouchTimeoutCondition(timeout=30.)
+ ),
+ transition_engine=RocketSimEngine(),
+ renderer=RLViserRenderer(120)
+ )
+
+ # simulate 2 episodes
+ for _ in range(2):
+ obs_dict = env.reset()
+ steps = 0
+ ep_reward = {agent_id: 0. for agent_id in env.agents}
+ t0 = time.time()
+ while True:
+ actions = {}
+ for agent_id, action_space in env.action_spaces.items():
+ # agent.act(obs) | Your agent should go here
+ actions[agent_id] = np.random.randint(action_space, size=(1,))
+
+ for _ in range(8):
+ obs_dict, reward_dict, terminated_dict, truncated_dict = env.step(actions)
+ env.render()
+ time.sleep(max(0, t0 + steps / 1200 - time.time()))
+ steps += 1
+
+ for agent_id, reward in reward_dict.items():
+ ep_reward[agent_id] += reward
+
+ if any(chain(terminated_dict.values(), truncated_dict.values())):
+ break
+
+ ep_time = time.time() - t0
+ print(f"Steps per second: {steps / ep_time:.0f} | Episode time: {ep_time:.2f} | Episode Reward: {max(ep_reward.values()):.2f}")

pytest.py

@@ -3,42 +3,43 @@
 import rlviser_py as vis
 import RocketSim as rs
 
-game_mode = rs.GameMode.HOOPS
-
-# Create example arena
-arena = rs.Arena(game_mode)
-
-# Set boost pad locations
-vis.set_boost_pad_locations([pad.get_pos().as_tuple() for pad in arena.get_boost_pads()])
-
-# Setup example arena
-car = arena.add_car(rs.Team.BLUE)
-car.set_state(rs.CarState(pos=rs.Vec(z=17), vel=rs.Vec(x=50), boost=100))
-arena.ball.set_state(rs.BallState(pos=rs.Vec(y=400, z=100), ang_vel=rs.Vec(x=5)))
-car.set_controls(rs.CarControls(throttle=1, steer=1, boost=True))
-
-# Run for 3 seconds
-TIME = 3
-
-steps = 0
-start_time = time.time()
-for i in range(round(TIME * arena.tick_rate)):
- arena.step(1)
-
- # Render the current game state
- pad_states = [pad.get_state().is_active for pad in arena.get_boost_pads()]
- ball = arena.ball.get_state()
- car_data = [
- (car.id, car.team, car.get_config(), car.get_state())
- for car in arena.get_cars()
- ]
-
- vis.render(steps, arena.tick_rate, game_mode, pad_states, ball, car_data)
-
- # sleep to simulate running real time (it will run a LOT after otherwise)
- time.sleep(max(0, start_time + steps / arena.tick_rate - time.time()))
- steps += 1
-
-# Tell RLViser to exit
-print("Exiting...")
-vis.quit()
+if __name__ == "__main__":
+ game_mode = rs.GameMode.HOOPS
+
+ # Create example arena
+ arena = rs.Arena(game_mode)
+
+ # Set boost pad locations
+ vis.set_boost_pad_locations([pad.get_pos().as_tuple() for pad in arena.get_boost_pads()])
+
+ # Setup example arena
+ car = arena.add_car(rs.Team.BLUE)
+ car.set_state(rs.CarState(pos=rs.Vec(z=17), vel=rs.Vec(x=50), boost=100))
+ arena.ball.set_state(rs.BallState(pos=rs.Vec(y=400, z=100), ang_vel=rs.Vec(x=5)))
+ car.set_controls(rs.CarControls(throttle=1, steer=1, boost=True))
+
+ # Run for 3 seconds
+ TIME = 3
+
+ steps = 0
+ start_time = time.time()
+ for i in range(round(TIME * arena.tick_rate)):
+ arena.step(1)
+
+ # Render the current game state
+ pad_states = [pad.get_state().is_active for pad in arena.get_boost_pads()]
+ ball = arena.ball.get_state()
+ car_data = [
+ (car.id, car.team, car.get_config(), car.get_state())
+ for car in arena.get_cars()
+ ]
+
+ vis.render(steps, arena.tick_rate, game_mode, pad_states, ball, car_data)
+
+ # sleep to simulate running real time (it will run a LOT after otherwise)
+ time.sleep(max(0, start_time + steps / arena.tick_rate - time.time()))
+ steps += 1
+
+ # Tell RLViser to exit
+ print("Exiting...")
+ vis.quit()

rlviser_py.pyi

@@ -1,11 +1,19 @@
-from typing import List, Tuple
+from typing import List, Optional, Tuple
 
 from RocketSim import BallState, CarConfig, CarState, GameMode, Team
 
-def set_boost_pad_locations(locations: List[Tuple[float, float, float]]) -> ...:
+type TVec3 = Tuple[float, float, float]
+type TRotmat = Tuple[TVec3, TVec3, TVec3]
+type TBall = Tuple[TVec3, TRotmat, TVec3, TVec3]
+type TCar = Tuple[int, TVec3, TRotmat, TVec3, TVec3, float, bool, bool, bool]
+
+def set_boost_pad_locations(locations: List[TVec3]) -> ...:
+ pass
+
+def get_state_set() -> Optional[Tuple[List[float], TBall, List[TCar]]]:
  pass
 
-def render(tick_count: int, tick_rate: float, game_mode: GameMode, boost_pad_states: List[bool], ball: BallState, cars: List[Tuple[int, Team, CarConfig, CarState]]) -> ...:
+def render(tick_count: int, tick_rate: float, game_mode: GameMode, boost_pad_states: List[bool], ball: BallState, cars: List[Tuple[int, CarState]]) -> ...:
  pass
 
 def quit() -> ...: