add parser in train_model_base file

train_model_base.py

@@ -1,54 +1,112 @@
+import argparse
+
 import numpy as np
 
 import rnl as vault
 
 if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Train or setup environment.")
+    parser.add_argument(
+        "mode", choices=["train", "run"], help="Mode to run: 'train' or 'run'"
+    )
+
+    # Parse arguments
+    args = parser.parse_args()
+
     # 1.step -> config robot
+    # base_radius: float, vel_linear: Tuple, val_angular: Tuple, wheel_distance: float, weight: float, threshold: float
     param_robot = vault.robot(
-        radius=40,  # (centimeters)
-        vel_linear=[0.0, 2.0],  # [min, max]
-        val_angular=[1.0, 2.0],  # [min, max]
+        base_radius=0.033,  # (centimeters) # TODO: RANDOMIZE
+        vel_linear=[0.0, 2.0],  # [min, max] # TODO: RANDOMIZE
+        val_angular=[1.0, 2.0],  # [min, max] # TODO: RANDOMIZE
+        wheel_distance=0.16,  # (centimeters) # TODO: RANDOMIZE
+        weight=1.0,  # (kilograms) # TODO: RANDOMIZE
+        threshold=0.01,  # (centimeters) # TODO: RANDOMIZE
     )
 
     # 2.step -> config sensors [for now only lidar sensor!!]
     param_sensor = vault.sensor(
-        fov=4 * np.pi, num_rays=20, min_range=0.0, max_range=6.0  # int
+        fov=4 * np.pi,  # TODO: RANDOMIZE
+        num_rays=20,  # TODO: RANDOMIZE
+        min_range=0.0,  # TODO: RANDOMIZE
+        max_range=6.0,  # TODO: RANDOMIZE
     )
+
     # 3.step -> config env
     param_env = vault.make(
-        map="None",
-        mode="normal",  # hard (muda tudo), normal (mapa fixo)
-        timestep=1000,
-        fps=100,
-        threshold=0.05,
-        grid_lenght=5,
-        physical="normal",
+        map_file="None",  # TODO: RANDOMIZE
+        random_mode="normal",  # hard (muda tudo), normal (mapa fixo)
+        timestep=1000,  # TODO: RANDOMIZE
+        grid_dimension=5,  # TODO: RANDOMIZE
+        friction=0.4,  # TODO: RANDOMIZE
+        porcentage_obstacles=0.1,  # TODO: RANDOMIZE
+        max_step=1000,  # TODO: RANDOMIZE
     )
 
-    # 4.step -> config train robot
-    model = vault.Trainer(param_robot, param_sensor, param_env, pretrained_model=False)
-    # 5.step -> train robot
-    model.learn(
-        max_timestep=800000,
-        use_mutation=False,
-        freq_evolution=10000,
-        log=False,
-        batch_size=64,
-        lr=0.0001,
-        pop_size=6,
-        hidden_size=[800, 600],
-        no_mut=0.4,
-        arch_mut=0.2,
-        new_layer=0.2,
-        param_mut=0.2,
-        act_mut=0,
-        hp_mut=0.2,
-        mut_strength=0.1,
-        seed=1,
-        num_envs=1,
-        device="mps",
-        learn_step=10,
-        n_step=3,
-        memory_size=1000000,
-        target_score=200.0,
-    )
+    if args.mode == "train":
+        # 4.step -> config train robot
+        model = vault.Trainer(
+            param_robot, param_sensor, param_env, pretrained_model=False
+        )
+        # 5.step -> train robot
+        model.learn(
+            max_timestep=800000,
+            memory_size=1000000,
+            gamma=0.99,
+            n_step=3,
+            alpha=0.6,
+            beta=0.4,
+            tau=0.001,
+            prior_eps=0.000001,
+            num_atoms=51,
+            v_min=-200,
+            v_max=200,
+            epsilon_start=1.0,
+            epsilon_end=0.1,
+            epsilon_decay=0.995,
+            batch_size=64,
+            lr=0.0001,
+            seed=1,
+            num_envs=2,
+            device="mps",
+            learn_step=10,
+            target_score=200,
+            max_steps=1000000,
+            evaluation_steps=10000,
+            evaluation_loop=1,
+            learning_delay=0,
+            n_step_memory=1,
+            checkpoint=100,
+            checkpoint_path="checkpoints",
+            overwrite_checkpoints=False,
+            use_wandb=False,
+            wandb_api_key="",
+            accelerator=False,
+            use_mutation=True,
+            freq_evolution=10000,
+            population_size=6,
+            no_mutation=0.4,
+            arch_mutation=0.2,
+            new_layer=0.2,
+            param_mutation=0.2,
+            active_mutation=0.0,
+            hp_mutation=0.2,
+            hp_mutation_selection=["lr", "batch_size"],
+            mutation_strength=0.1,
+            evolution_steps=10000,
+            save_elite=False,
+            elite_path="elite",
+            tourn_size=2,
+            elitism=True,
+            hidden_size=[800, 600],
+        )
+    else:
+        # 4.step -> config render
+        param_render = vault.render(fps=100, controller=True, rgb_array=True)
+
+        # 5.step -> config train robot
+        model = vault.Trainer(
+            param_robot, param_sensor, param_env, param_render, pretrained_model=False
+        )
+        # 6.step -> run robot
+        model.run()