feat: rl session

crates/beet_ml/src/environments/frozen_lake/frozen_lake_plugin.rs

@@ -42,13 +42,14 @@ impl Plugin for FrozenLakePlugin {
 				TranslateGrid,
 				StepEnvironment<FrozenLakeQTableSession>,
 			)>::default(),
-			EpisodeRunnerPlugin::<FrozenLakeEpParams>::default(),
+			RlSessionPlugin::<FrozenLakeEpParams>::default(),
 		))
 		.add_systems(Startup, init_frozen_lake_assets)
 		.add_systems(Update, reward_grid.in_set(PostTickSet))
 		.add_systems(
 			Update,
-			(spawn_frozen_lake_static, spawn_frozen_lake).in_set(PostTickSet),
+			(spawn_frozen_lake_session, spawn_frozen_lake_episode)
+				.in_set(PostTickSet),
 		)
 		.init_resource::<RlRng>();
 

crates/beet_ml/src/environments/frozen_lake/spawn_frozen_lake.rs

@@ -33,7 +33,7 @@ pub fn init_frozen_lake_assets(
 	});
 }
 
-pub fn spawn_frozen_lake_static(
+pub fn spawn_frozen_lake_session(
 	mut events: EventReader<StartSession<FrozenLakeEpParams>>,
 	mut commands: Commands,
 	assets: Res<FrozenLakeAssets>,
@@ -56,7 +56,8 @@ pub fn spawn_frozen_lake_static(
 							.with_scale(tile_scale),
 						..default()
 					},
-					EpisodeOwner(event.trainer),
+					SessionEntity(event.session),
+					DespawnOnSessionEnd,
 				));
 			}
 		}
@@ -76,7 +77,8 @@ pub fn spawn_frozen_lake_static(
 								.with_scale(object_scale),
 							..default()
 						},
-						EpisodeOwner(event.trainer),
+						SessionEntity(event.session),
+						DespawnOnSessionEnd,
 					));
 				}
 				FrozenLakeCell::Goal => {
@@ -87,7 +89,8 @@ pub fn spawn_frozen_lake_static(
 								.with_scale(object_scale),
 							..default()
 						},
-						EpisodeOwner(event.trainer),
+						SessionEntity(event.session),
+						DespawnOnSessionEnd,
 					));
 				}
 				FrozenLakeCell::Ice => {}
@@ -99,7 +102,7 @@ pub fn spawn_frozen_lake_static(
 }
 
 
-pub fn spawn_frozen_lake(
+pub fn spawn_frozen_lake_episode(
 	mut events: EventReader<StartEpisode<FrozenLakeEpParams>>,
 	mut commands: Commands,
 	assets: Res<FrozenLakeAssets>,
@@ -128,10 +131,10 @@ pub fn spawn_frozen_lake(
 							RlAgentBundle {
 								state: map.agent_position(),
 								action: GridDirection::sample(),
-								table: QTable::default(),
 								env: FrozenLakeEnv::new(map.clone(), false),
 								params: event.params.learn_params.clone(),
-								trainer: EpisodeOwner(event.trainer),
+								session: SessionEntity(event.session),
+								despawn: DespawnOnEpisodeEnd,
 							},
 						))
 						.with_children(|parent| {
@@ -146,7 +149,7 @@ pub fn spawn_frozen_lake(
 								.with_children(|parent| {
 									parent.spawn((
 										TranslateGrid::new(
-											Duration::from_secs(1),
+											Duration::from_millis(1),
 										),
 										TargetAgent(agent),
 										RunTimer::default(),

crates/beet_ml/src/rl_realtime/mod.rs

@@ -1,9 +1,9 @@
 pub mod rl_agent;
 #[allow(unused_imports)]
 pub use self::rl_agent::*;
-pub mod episode_runner;
+pub mod rl_session;
 #[allow(unused_imports)]
-pub use self::episode_runner::*;
+pub use self::rl_session::*;
 pub mod rl_components;
 #[allow(unused_imports)]
 pub use self::rl_components::*;
@@ -13,6 +13,6 @@ pub use self::step_environment::*;
 pub mod rl_plugin;
 #[allow(unused_imports)]
 pub use self::rl_plugin::*;
-pub mod rl_session;
+pub mod rl_session_types;
 #[allow(unused_imports)]
-pub use self::rl_session::*;
+pub use self::rl_session_types::*;

crates/beet_ml/src/rl_realtime/rl_agent.rs

@@ -5,22 +5,19 @@ use bevy::prelude::*;
 
 
 #[derive(Bundle)]
-pub struct RlAgentBundle<
-	Env: Component + Environment<State = Table::State, Action = Table::Action>,
-	Table: Component + QSource,
-> {
-	pub state: Table::State,
-	pub action: Table::Action,
-	pub table: Table,
+pub struct RlAgentBundle<Env: Component + Environment> {
+	pub state: Env::State,
+	pub action: Env::Action,
 	pub env: Env,
 	pub params: QLearnParams,
-	pub trainer: EpisodeOwner,
+	pub session: SessionEntity,
+	pub despawn:DespawnOnEpisodeEnd
 }
 
-
-impl<
-		Env: Component + Environment<State = Table::State, Action = Table::Action>,
-		Table: Component + QSource,
-	> RlAgentBundle<Env, Table>
-{
-}
+// #[derive(Bundle)]
+// pub struct RlSessionBundle<S: RlSessionTypes>
+// where
+// 	S::QSource: Component,
+// {
+// 	pub source: S::QSource,
+// }

crates/beet_ml/src/rl_realtime/rl_plugin.rs

@@ -8,11 +8,11 @@ pub struct RlPlugin;
 impl Plugin for RlPlugin {
 	fn build(&self, app: &mut App) {
 		let world = app.world_mut();
-		world.init_component::<EpisodeOwner>();
+		world.init_component::<SessionEntity>();
 
 		let mut registry =
 			world.get_resource::<AppTypeRegistry>().unwrap().write();
 
-		registry.register::<EpisodeOwner>();
+		registry.register::<SessionEntity>();
 	}
 }